jax_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['state_0', 'state_1', 'state_2', 'state_3', 'state_4', 'state_5', 'state_6', 'state_7', 'state_8', 'state_9', 'state_10', 'state_11', 'state_12', 'state_13', 'state_14', 'state_15', 'state_16', 'state_17', 'state_18', 'state_19', 'state_20', 'state_21', 'state_22', 'state_23', 'state_24', 'state_25', 'state_26', 'action_0', 'action_1', 'action_2', 'action_3', 'action_4', 'action_5', 'action_6', 'reward']\n",
      "ep 0\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mActorCritic\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mMujocoSim\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;66;03m# Instantiate your environment\u001b[39;00m\n",
      "File \u001b[0;32m~/rl-manipulator/ActorCritic.py:301\u001b[0m\n\u001b[1;32m    299\u001b[0m batch_train\u001b[38;5;241m=\u001b[39mjax\u001b[38;5;241m.\u001b[39mjit(ac\u001b[38;5;241m.\u001b[39mbatch_train)\n\u001b[1;32m    300\u001b[0m \u001b[38;5;66;03m# batch_train()\u001b[39;00m\n\u001b[0;32m--> 301\u001b[0m \u001b[43mac\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_train\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    302\u001b[0m \u001b[38;5;66;03m# ac.train()\u001b[39;00m\n",
      "File \u001b[0;32m~/rl-manipulator/ActorCritic.py:155\u001b[0m, in \u001b[0;36mActorCritic.batch_train\u001b[0;34m(self, episodes, batch_size)\u001b[0m\n\u001b[1;32m    151\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m200\u001b[39m):\n\u001b[1;32m    153\u001b[0m     action \u001b[38;5;241m=\u001b[39m select_action(state, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mactor_params)\n\u001b[0;32m--> 155\u001b[0m     next_state, reward, batch \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmjx_model\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maction\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    157\u001b[0m     \u001b[38;5;66;03m# jax.debug.print(f\"next_state.shape {next_state.shape} mean {jnp.mean(next_state, axis=0)}\")\u001b[39;00m\n\u001b[1;32m    158\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(batch_size):    \n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/api.py:1240\u001b[0m, in \u001b[0;36mvmap.<locals>.vmap_f\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1237\u001b[0m in_axes_flat \u001b[38;5;241m=\u001b[39m flatten_axes(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvmap in_axes\u001b[39m\u001b[38;5;124m\"\u001b[39m, in_tree, (in_axes, \u001b[38;5;241m0\u001b[39m), kws\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m   1238\u001b[0m axis_size_ \u001b[38;5;241m=\u001b[39m (axis_size \u001b[38;5;28;01mif\u001b[39;00m axis_size \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m\n\u001b[1;32m   1239\u001b[0m               _mapped_axis_size(fun, in_tree, args_flat, in_axes_flat, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvmap\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m-> 1240\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mbatching\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1241\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis_size_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axes_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1242\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mflatten_axes\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mvmap out_axes\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_tree\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axes\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1243\u001b[0m \u001b[43m    \u001b[49m\u001b[43mspmd_axis_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mspmd_axis_name\u001b[49m\n\u001b[1;32m   1244\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1245\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m tree_unflatten(out_tree(), out_flat)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/linear_util.py:192\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    189\u001b[0m gen \u001b[38;5;241m=\u001b[39m gen_static_args \u001b[38;5;241m=\u001b[39m out_store \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    191\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 192\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    193\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    194\u001b[0m   \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m    195\u001b[0m   \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m    196\u001b[0m   \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks\u001b[39;00m\n\u001b[1;32m    197\u001b[0m   \u001b[38;5;66;03m# only after this exception is handled, which can corrupt the global\u001b[39;00m\n\u001b[1;32m    198\u001b[0m   \u001b[38;5;66;03m# state.\u001b[39;00m\n\u001b[1;32m    199\u001b[0m   \u001b[38;5;28;01mwhile\u001b[39;00m stack:\n",
      "File \u001b[0;32m~/rl-manipulator/MujocoSim.py:92\u001b[0m, in \u001b[0;36mMujocoSim.step\u001b[0;34m(self, model, data, action)\u001b[0m\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep\u001b[39m(\u001b[38;5;28mself\u001b[39m, model, data, action):\n\u001b[1;32m     82\u001b[0m     \u001b[38;5;66;03m# Execute the action and return the new state and reward\u001b[39;00m\n\u001b[1;32m     83\u001b[0m     \u001b[38;5;66;03m# rng = jax.random.PRNGKey(0)\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     88\u001b[0m     \u001b[38;5;66;03m# fun_vmapped = jax.vmap(fun)\u001b[39;00m\n\u001b[1;32m     89\u001b[0m     \u001b[38;5;66;03m# batch=fun_vmapped(rng)\u001b[39;00m\n\u001b[1;32m     90\u001b[0m     data\u001b[38;5;241m.\u001b[39mreplace(ctrl\u001b[38;5;241m=\u001b[39maction)\n\u001b[0;32m---> 92\u001b[0m     data\u001b[38;5;241m=\u001b[39m\u001b[43mmjx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     94\u001b[0m     state\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_state(data)\n\u001b[1;32m     95\u001b[0m     reward\u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_reward(state,action)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/forward.py:56\u001b[0m, in \u001b[0;36mnamed_scope.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(fn)\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapper\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m     55\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(name \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(fn, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__name__\u001b[39m\u001b[38;5;124m'\u001b[39m)):\n\u001b[0;32m---> 56\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     57\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m res\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/forward.py:365\u001b[0m, in \u001b[0;36mstep\u001b[0;34m(m, d)\u001b[0m\n\u001b[1;32m    362\u001b[0m \u001b[38;5;129m@named_scope\u001b[39m\n\u001b[1;32m    363\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep\u001b[39m(m: Model, d: Data) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Data:\n\u001b[1;32m    364\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Advance simulation.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 365\u001b[0m   d \u001b[38;5;241m=\u001b[39m \u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mm\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43md\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    367\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m m\u001b[38;5;241m.\u001b[39mopt\u001b[38;5;241m.\u001b[39mintegrator \u001b[38;5;241m==\u001b[39m IntegratorType\u001b[38;5;241m.\u001b[39mEULER:\n\u001b[1;32m    368\u001b[0m     d \u001b[38;5;241m=\u001b[39m euler(m, d)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/forward.py:56\u001b[0m, in \u001b[0;36mnamed_scope.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(fn)\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapper\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m     55\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(name \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(fn, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__name__\u001b[39m\u001b[38;5;124m'\u001b[39m)):\n\u001b[0;32m---> 56\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     57\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m res\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/forward.py:348\u001b[0m, in \u001b[0;36mforward\u001b[0;34m(m, d)\u001b[0m\n\u001b[1;32m    345\u001b[0m \u001b[38;5;129m@named_scope\u001b[39m\n\u001b[1;32m    346\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(m: Model, d: Data) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Data:\n\u001b[1;32m    347\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Forward dynamics.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 348\u001b[0m   d \u001b[38;5;241m=\u001b[39m \u001b[43mfwd_position\u001b[49m\u001b[43m(\u001b[49m\u001b[43mm\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43md\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    349\u001b[0m   d \u001b[38;5;241m=\u001b[39m fwd_velocity(m, d)\n\u001b[1;32m    350\u001b[0m   d \u001b[38;5;241m=\u001b[39m fwd_actuation(m, d)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/forward.py:56\u001b[0m, in \u001b[0;36mnamed_scope.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(fn)\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapper\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m     55\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(name \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(fn, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__name__\u001b[39m\u001b[38;5;124m'\u001b[39m)):\n\u001b[0;32m---> 56\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     57\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m res\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/forward.py:66\u001b[0m, in \u001b[0;36mfwd_position\u001b[0;34m(m, d)\u001b[0m\n\u001b[1;32m     64\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Position-dependent computations.\"\"\"\u001b[39;00m\n\u001b[1;32m     65\u001b[0m \u001b[38;5;66;03m# TODO(robotics-simulation): tendon\u001b[39;00m\n\u001b[0;32m---> 66\u001b[0m d \u001b[38;5;241m=\u001b[39m \u001b[43msmooth\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkinematics\u001b[49m\u001b[43m(\u001b[49m\u001b[43mm\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43md\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     67\u001b[0m d \u001b[38;5;241m=\u001b[39m smooth\u001b[38;5;241m.\u001b[39mcom_pos(m, d)\n\u001b[1;32m     68\u001b[0m d \u001b[38;5;241m=\u001b[39m smooth\u001b[38;5;241m.\u001b[39mcamlight(m, d)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/smooth.py:86\u001b[0m, in \u001b[0;36mkinematics\u001b[0;34m(m, d)\u001b[0m\n\u001b[1;32m     82\u001b[0m   mat \u001b[38;5;241m=\u001b[39m math\u001b[38;5;241m.\u001b[39mquat_to_mat(quat)\n\u001b[1;32m     84\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m qpos, anchor, axis, pos, quat, mat\n\u001b[0;32m---> 86\u001b[0m qpos, xanchor, xaxis, xpos, xquat, xmat \u001b[38;5;241m=\u001b[39m \u001b[43mscan\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody_tree\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     87\u001b[0m \u001b[43m    \u001b[49m\u001b[43mm\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     88\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     89\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mjjjqqbb\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     90\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mqjjbbb\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     91\u001b[0m \u001b[43m    \u001b[49m\u001b[43mm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjnt_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     92\u001b[0m \u001b[43m    \u001b[49m\u001b[43mm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjnt_pos\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     93\u001b[0m \u001b[43m    \u001b[49m\u001b[43mm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjnt_axis\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     94\u001b[0m \u001b[43m    \u001b[49m\u001b[43md\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mqpos\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     95\u001b[0m \u001b[43m    \u001b[49m\u001b[43mm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mqpos0\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     96\u001b[0m \u001b[43m    \u001b[49m\u001b[43mm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody_pos\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     97\u001b[0m \u001b[43m    \u001b[49m\u001b[43mm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody_quat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     98\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    100\u001b[0m v_local_to_global \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mvmap(support\u001b[38;5;241m.\u001b[39mlocal_to_global)\n\u001b[1;32m    102\u001b[0m \u001b[38;5;66;03m# TODO(erikfrey): confirm that quats are more performant for mjx than mats\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/scan.py:479\u001b[0m, in \u001b[0;36mbody_tree\u001b[0;34m(m, f, in_types, out_types, reverse, *args)\u001b[0m\n\u001b[1;32m    476\u001b[0m     carry \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mtree_map(take_fn, y)\n\u001b[1;32m    478\u001b[0m   f_args \u001b[38;5;241m=\u001b[39m [_take(arg, ids) \u001b[38;5;28;01mfor\u001b[39;00m arg, ids \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(args, key_in_take[key])]\n\u001b[0;32m--> 479\u001b[0m   key_y[key] \u001b[38;5;241m=\u001b[39m \u001b[43m_nvmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcarry\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mf_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    481\u001b[0m \u001b[38;5;66;03m# slice None results from the final output\u001b[39;00m\n\u001b[1;32m    482\u001b[0m keys \u001b[38;5;241m=\u001b[39m [k \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m keys \u001b[38;5;28;01mif\u001b[39;00m key_y[k] \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m]\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/scan.py:130\u001b[0m, in \u001b[0;36m_nvmap\u001b[0;34m(f, *args)\u001b[0m\n\u001b[1;32m    127\u001b[0m   args \u001b[38;5;241m=\u001b[39m [a \u001b[38;5;28;01mif\u001b[39;00m n \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m n \u001b[38;5;28;01mfor\u001b[39;00m n, a \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(args, np_args)]\n\u001b[1;32m    128\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m f(\u001b[38;5;241m*\u001b[39margs)\n\u001b[0;32m--> 130\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mouter_f\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43min_axes\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/api.py:1240\u001b[0m, in \u001b[0;36mvmap.<locals>.vmap_f\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1237\u001b[0m in_axes_flat \u001b[38;5;241m=\u001b[39m flatten_axes(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvmap in_axes\u001b[39m\u001b[38;5;124m\"\u001b[39m, in_tree, (in_axes, \u001b[38;5;241m0\u001b[39m), kws\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m   1238\u001b[0m axis_size_ \u001b[38;5;241m=\u001b[39m (axis_size \u001b[38;5;28;01mif\u001b[39;00m axis_size \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m\n\u001b[1;32m   1239\u001b[0m               _mapped_axis_size(fun, in_tree, args_flat, in_axes_flat, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvmap\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m-> 1240\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mbatching\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1241\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis_size_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axes_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1242\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mflatten_axes\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mvmap out_axes\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_tree\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axes\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1243\u001b[0m \u001b[43m    \u001b[49m\u001b[43mspmd_axis_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mspmd_axis_name\u001b[49m\n\u001b[1;32m   1244\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1245\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m tree_unflatten(out_tree(), out_flat)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/linear_util.py:192\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    189\u001b[0m gen \u001b[38;5;241m=\u001b[39m gen_static_args \u001b[38;5;241m=\u001b[39m out_store \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    191\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 192\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    193\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    194\u001b[0m   \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m    195\u001b[0m   \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m    196\u001b[0m   \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks\u001b[39;00m\n\u001b[1;32m    197\u001b[0m   \u001b[38;5;66;03m# only after this exception is handled, which can corrupt the global\u001b[39;00m\n\u001b[1;32m    198\u001b[0m   \u001b[38;5;66;03m# state.\u001b[39;00m\n\u001b[1;32m    199\u001b[0m   \u001b[38;5;28;01mwhile\u001b[39;00m stack:\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/scan.py:128\u001b[0m, in \u001b[0;36m_nvmap.<locals>.outer_f\u001b[0;34m(np_args, *args)\u001b[0m\n\u001b[1;32m    126\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mouter_f\u001b[39m(\u001b[38;5;241m*\u001b[39margs, np_args\u001b[38;5;241m=\u001b[39mnp_args):\n\u001b[1;32m    127\u001b[0m   args \u001b[38;5;241m=\u001b[39m [a \u001b[38;5;28;01mif\u001b[39;00m n \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m n \u001b[38;5;28;01mfor\u001b[39;00m n, a \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(args, np_args)]\n\u001b[0;32m--> 128\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/smooth.py:82\u001b[0m, in \u001b[0;36mkinematics.<locals>.fn\u001b[0;34m(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat)\u001b[0m\n\u001b[1;32m     80\u001b[0m anchor \u001b[38;5;241m=\u001b[39m jp\u001b[38;5;241m.\u001b[39mstack(anchors) \u001b[38;5;28;01mif\u001b[39;00m anchors \u001b[38;5;28;01melse\u001b[39;00m jp\u001b[38;5;241m.\u001b[39mempty((\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m3\u001b[39m))\n\u001b[1;32m     81\u001b[0m axis \u001b[38;5;241m=\u001b[39m jp\u001b[38;5;241m.\u001b[39mstack(axes) \u001b[38;5;28;01mif\u001b[39;00m axes \u001b[38;5;28;01melse\u001b[39;00m jp\u001b[38;5;241m.\u001b[39mempty((\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m3\u001b[39m))\n\u001b[0;32m---> 82\u001b[0m mat \u001b[38;5;241m=\u001b[39m \u001b[43mmath\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mquat_to_mat\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquat\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m qpos, anchor, axis, pos, quat, mat\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/math.py:194\u001b[0m, in \u001b[0;36mquat_to_mat\u001b[0;34m(q)\u001b[0m\n\u001b[1;32m    177\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Converts a quaternion into a 9-dimensional rotation matrix.\"\"\"\u001b[39;00m\n\u001b[1;32m    178\u001b[0m q \u001b[38;5;241m=\u001b[39m jp\u001b[38;5;241m.\u001b[39mouter(q, q)\n\u001b[1;32m    180\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m jp\u001b[38;5;241m.\u001b[39marray([\n\u001b[1;32m    181\u001b[0m     [\n\u001b[1;32m    182\u001b[0m         q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m+\u001b[39m q[\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m] \u001b[38;5;241m-\u001b[39m q[\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m2\u001b[39m] \u001b[38;5;241m-\u001b[39m q[\u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m3\u001b[39m],\n\u001b[1;32m    183\u001b[0m         \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m (q[\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m] \u001b[38;5;241m-\u001b[39m q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m3\u001b[39m]),\n\u001b[1;32m    184\u001b[0m         \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m (q[\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m3\u001b[39m] \u001b[38;5;241m+\u001b[39m q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m2\u001b[39m]),\n\u001b[1;32m    185\u001b[0m     ],\n\u001b[1;32m    186\u001b[0m     [\n\u001b[1;32m    187\u001b[0m         \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m (q[\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m] \u001b[38;5;241m+\u001b[39m q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m3\u001b[39m]),\n\u001b[1;32m    188\u001b[0m         q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m-\u001b[39m q[\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m] \u001b[38;5;241m+\u001b[39m q[\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m2\u001b[39m] \u001b[38;5;241m-\u001b[39m q[\u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m3\u001b[39m],\n\u001b[1;32m    189\u001b[0m         \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m (q[\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m3\u001b[39m] \u001b[38;5;241m-\u001b[39m q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m]),\n\u001b[1;32m    190\u001b[0m     ],\n\u001b[1;32m    191\u001b[0m     [\n\u001b[1;32m    192\u001b[0m         \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m (q[\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m3\u001b[39m] \u001b[38;5;241m-\u001b[39m q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m2\u001b[39m]),\n\u001b[1;32m    193\u001b[0m         \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m (q[\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m3\u001b[39m] \u001b[38;5;241m+\u001b[39m q[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m]),\n\u001b[0;32m--> 194\u001b[0m         \u001b[43mq\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mq\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mq\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mq\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m]\u001b[49m,\n\u001b[1;32m    195\u001b[0m     ],\n\u001b[1;32m    196\u001b[0m ])\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:736\u001b[0m, in \u001b[0;36m_forward_operator_to_aval.<locals>.op\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m    735\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mop\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs):\n\u001b[0;32m--> 736\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mgetattr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maval\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m_\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mname\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:264\u001b[0m, in \u001b[0;36m_defer_to_unrecognized_arg.<locals>.deferring_binary_op\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m    262\u001b[0m args \u001b[38;5;241m=\u001b[39m (other, \u001b[38;5;28mself\u001b[39m) \u001b[38;5;28;01mif\u001b[39;00m swap \u001b[38;5;28;01melse\u001b[39;00m (\u001b[38;5;28mself\u001b[39m, other)\n\u001b[1;32m    263\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(other, _accepted_binop_types):\n\u001b[0;32m--> 264\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mbinary_op\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;66;03m# Note: don't use isinstance here, because we don't want to raise for\u001b[39;00m\n\u001b[1;32m    266\u001b[0m \u001b[38;5;66;03m# subclasses, e.g. NamedTuple objects that may override operators.\u001b[39;00m\n\u001b[1;32m    267\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(other) \u001b[38;5;129;01min\u001b[39;00m _rejected_binop_types:\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/pjit.py:298\u001b[0m, in \u001b[0;36m_cpp_pjit.<locals>.cache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    296\u001b[0m \u001b[38;5;129m@api_boundary\u001b[39m\n\u001b[1;32m    297\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcache_miss\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m--> 298\u001b[0m   outs, out_flat, out_tree, args_flat, jaxpr, attrs_tracked \u001b[38;5;241m=\u001b[39m \u001b[43m_python_pjit_helper\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    299\u001b[0m \u001b[43m      \u001b[49m\u001b[43mjit_info\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    300\u001b[0m   executable \u001b[38;5;241m=\u001b[39m _read_most_recent_pjit_call_executable(jaxpr)\n\u001b[1;32m    301\u001b[0m   maybe_fastpath_data \u001b[38;5;241m=\u001b[39m _get_fastpath_data(\n\u001b[1;32m    302\u001b[0m       executable, out_tree, args_flat, out_flat, attrs_tracked, jaxpr\u001b[38;5;241m.\u001b[39meffects,\n\u001b[1;32m    303\u001b[0m       jit_info\u001b[38;5;241m.\u001b[39mabstracted_axes)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/pjit.py:176\u001b[0m, in \u001b[0;36m_python_pjit_helper\u001b[0;34m(jit_info, *args, **kwargs)\u001b[0m\n\u001b[1;32m    173\u001b[0m   args_flat \u001b[38;5;241m=\u001b[39m [\u001b[38;5;241m*\u001b[39minit_states, \u001b[38;5;241m*\u001b[39margs_flat]\n\u001b[1;32m    175\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 176\u001b[0m   out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mpjit_p\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    177\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m pxla\u001b[38;5;241m.\u001b[39mDeviceAssignmentMismatchError \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    178\u001b[0m   fails, \u001b[38;5;241m=\u001b[39m e\u001b[38;5;241m.\u001b[39margs\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/core.py:2784\u001b[0m, in \u001b[0;36mAxisPrimitive.bind\u001b[0;34m(self, *args, **params)\u001b[0m\n\u001b[1;32m   2782\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[1;32m   2783\u001b[0m   top_trace \u001b[38;5;241m=\u001b[39m find_top_trace(args)\n\u001b[0;32m-> 2784\u001b[0m   axis_main \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mmax\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43maxis_frame\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmain_trace\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43ma\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mused_axis_names\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2785\u001b[0m \u001b[43m                  \u001b[49m\u001b[43mdefault\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mgetattr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mlevel\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2786\u001b[0m   top_trace \u001b[38;5;241m=\u001b[39m (top_trace \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m axis_main \u001b[38;5;129;01mor\u001b[39;00m axis_main\u001b[38;5;241m.\u001b[39mlevel \u001b[38;5;241m<\u001b[39m top_trace\u001b[38;5;241m.\u001b[39mlevel\n\u001b[1;32m   2787\u001b[0m                \u001b[38;5;28;01melse\u001b[39;00m axis_main\u001b[38;5;241m.\u001b[39mwith_cur_sublevel())\n\u001b[1;32m   2788\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbind_with_trace(top_trace, args, params)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/jax/_src/core.py:2784\u001b[0m, in \u001b[0;36m<genexpr>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m   2782\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[1;32m   2783\u001b[0m   top_trace \u001b[38;5;241m=\u001b[39m find_top_trace(args)\n\u001b[0;32m-> 2784\u001b[0m   axis_main \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmax\u001b[39m((axis_frame(a)\u001b[38;5;241m.\u001b[39mmain_trace \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m used_axis_names(\u001b[38;5;28mself\u001b[39m, params)),\n\u001b[1;32m   2785\u001b[0m                   default\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, key\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mlambda\u001b[39;00m t: \u001b[38;5;28mgetattr\u001b[39m(t, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlevel\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m))\n\u001b[1;32m   2786\u001b[0m   top_trace \u001b[38;5;241m=\u001b[39m (top_trace \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m axis_main \u001b[38;5;129;01mor\u001b[39;00m axis_main\u001b[38;5;241m.\u001b[39mlevel \u001b[38;5;241m<\u001b[39m top_trace\u001b[38;5;241m.\u001b[39mlevel\n\u001b[1;32m   2787\u001b[0m                \u001b[38;5;28;01melse\u001b[39;00m axis_main\u001b[38;5;241m.\u001b[39mwith_cur_sublevel())\n\u001b[1;32m   2788\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbind_with_trace(top_trace, args, params)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "\n",
    "from ActorCritic import *\n",
    "from MujocoSim import *\n",
    "\n",
    "# Instantiate your environment\n",
    "env = MujocoSim()\n",
    "\n",
    "# Create ActorCritic instance\n",
    "ac = ActorCritic(env, actor_params, critic_params)\n",
    "\n",
    "# # Train the model\n",
    "# ac.train(episodes=1000)\n",
    "\n",
    "\n",
    "\n",
    "# Add this method to the ActorCritic class\n",
    "# ActorCritic.train_batch = train_batch\n",
    "# ActorCritic.batch_update = batch_update\n",
    "# ActorCritic.batch_select_action = batch_select_action\n",
    "\n",
    "# Train the model with batches\n",
    "# ac.train_batch(batch_size=32, episodes=1000)\n",
    "# tr=jax.jit(ac.train)\n",
    "# tr()\n",
    "\n",
    "batch_train=jax.jit(ac.batch_train)\n",
    "batch_train()\n",
    "# ac.batch_train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'W1': Array([[-1.7639800e-03,  1.3475739e-03, -1.0666488e-02, ...,\n",
      "        -1.7109284e-02,  1.5562554e-03,  2.1836185e-03],\n",
      "       [-1.1511814e-02,  2.6543327e-03,  5.7022562e-03, ...,\n",
      "         6.2514259e-03,  6.1439765e-03,  2.5876977e-03],\n",
      "       [-6.9878413e-03,  7.8353770e-03,  7.1274065e-03, ...,\n",
      "         7.8045707e-03, -1.1834394e-02,  2.3500472e-03],\n",
      "       ...,\n",
      "       [ 8.2423979e-01,  9.3815941e-03, -1.0896824e-02, ...,\n",
      "         8.7352598e-01,  2.5724322e-01, -4.6460144e-04],\n",
      "       [ 1.6423662e+01, -2.1316223e-04,  6.5006944e-03, ...,\n",
      "         1.7041855e+01,  4.9237900e+00, -1.1844464e-02],\n",
      "       [ 6.1574252e-03,  2.5239927e-03,  1.3379987e-02, ...,\n",
      "         2.8177933e-03,  1.4925360e-02,  1.6257811e-02]], dtype=float32), 'b1': Array([ 1.6409216e+01,  0.0000000e+00,  0.0000000e+00,  1.9004103e+01,\n",
      "        1.4679784e+01,  0.0000000e+00,  0.0000000e+00, -1.0753442e-02,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
      "        9.6775770e+00,  2.4143900e+01,  0.0000000e+00,  1.5218526e+01,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
      "        4.2558050e+00,  0.0000000e+00, -7.1592117e-03,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -3.5333224e-03,\n",
      "        0.0000000e+00,  0.0000000e+00, -1.0402502e-02,  2.3977291e+01,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
      "        1.6203636e+01,  1.7557173e+01,  1.1929279e+01,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -4.2796843e-03,\n",
      "        0.0000000e+00,  0.0000000e+00, -6.4210212e-03, -1.0339653e-02,\n",
      "        0.0000000e+00,  7.6032248e+00,  1.1068655e+01,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  1.5226414e+01,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
      "        3.4720006e+00, -9.5652090e-03,  1.0660647e+00,  0.0000000e+00,\n",
      "        2.0217128e+00,  1.3127703e+01,  9.5274630e+00,  5.4256201e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  1.2648825e+01,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00,  2.3998026e+01, -1.4391742e-04,\n",
      "        8.5533476e+00,  2.0615557e+01,  1.3752448e+01,  2.8659348e+01,\n",
      "        0.0000000e+00,  0.0000000e+00,  4.5480667e+01,  0.0000000e+00,\n",
      "        2.1122217e+00,  2.3985600e+00,  0.0000000e+00,  0.0000000e+00,\n",
      "        1.7389385e+01,  0.0000000e+00,  0.0000000e+00, -7.6160124e-03,\n",
      "       -3.0325255e-03,  8.5609598e+00,  1.8118477e+01,  0.0000000e+00,\n",
      "        2.6536155e+00, -1.0561910e-02,  2.1372997e+01,  0.0000000e+00,\n",
      "        0.0000000e+00,  0.0000000e+00, -1.0487474e-02,  1.1494564e+01,\n",
      "        0.0000000e+00,  0.0000000e+00,  7.3569813e+00,  3.4747404e-01,\n",
      "        0.0000000e+00,  0.0000000e+00, -9.3712797e-03,  2.0728542e+01,\n",
      "        0.0000000e+00,  2.1886337e+01,  2.2987185e+01,  5.0200925e+00,\n",
      "        1.1607883e+01,  1.7031353e+01,  4.9234986e+00,  0.0000000e+00],      dtype=float32), 'W2': Array([[-8.52892971e+00, -8.67022705e+00,  1.16384401e+01,\n",
      "        -8.19744682e+00, -8.10352898e+00,  3.05660038e+01,\n",
      "        -8.66883183e+00],\n",
      "       [-2.23555346e-03,  1.57635454e-02,  6.65206369e-03,\n",
      "        -6.02840958e-03, -2.80043506e-03,  7.27308542e-03,\n",
      "        -1.27434898e-02],\n",
      "       [-5.96316904e-03, -5.35314763e-03, -1.36993024e-02,\n",
      "        -8.56473669e-03,  8.01949063e-04,  1.15627619e-02,\n",
      "        -2.06519803e-03],\n",
      "       [-9.57728672e+00, -9.73606491e+00,  1.30720768e+01,\n",
      "        -9.21071434e+00, -9.12463665e+00,  3.43249359e+01,\n",
      "        -9.73449326e+00],\n",
      "       [-7.50478697e+00, -7.62670612e+00,  1.02235250e+01,\n",
      "        -7.22447968e+00, -7.13950920e+00,  2.68929996e+01,\n",
      "        -7.61544752e+00],\n",
      "       [ 2.67935963e-03, -1.49912238e-02, -6.31031906e-03,\n",
      "         1.60926767e-02,  9.31915734e-03,  9.21037514e-03,\n",
      "        -4.50251671e-03],\n",
      "       [-9.94189829e-03, -1.23365857e-02,  1.10089099e-02,\n",
      "         9.39102657e-03,  2.29294710e-02,  3.95066664e-03,\n",
      "        -1.31530436e-02],\n",
      "       [ 2.30194535e-03, -8.85050092e-03, -8.13414715e-03,\n",
      "         1.17765814e-02, -7.92195834e-03, -9.26212221e-03,\n",
      "         8.79941508e-04],\n",
      "       [ 9.55727603e-03, -2.79511616e-04,  5.30077703e-03,\n",
      "        -1.44581718e-03,  1.33888703e-02, -1.32427579e-02,\n",
      "        -4.52246610e-03],\n",
      "       [-1.03035159e-02, -2.19071712e-02, -1.26467701e-02,\n",
      "        -1.01261605e-02,  1.02824140e-02, -6.77769491e-03,\n",
      "        -1.32416524e-02],\n",
      "       [ 9.76030994e-03,  4.61006211e-03,  1.94343571e-02,\n",
      "        -1.01098977e-02,  1.50212953e-02,  3.70889017e-03,\n",
      "        -3.90958181e-03],\n",
      "       [ 6.37731748e-03, -2.17269477e-03, -1.24115888e-02,\n",
      "        -2.17712373e-02,  2.40439340e-03, -1.91675255e-03,\n",
      "        -1.51250679e-02],\n",
      "       [-5.01468182e+00, -5.09047604e+00,  6.84179640e+00,\n",
      "        -4.84065008e+00, -4.77694798e+00,  1.79727669e+01,\n",
      "        -5.08364534e+00],\n",
      "       [-1.25608196e+01, -1.27704926e+01,  1.71267681e+01,\n",
      "        -1.20998535e+01, -1.19321556e+01,  4.49690781e+01,\n",
      "        -1.27581072e+01],\n",
      "       [-1.66859329e-02,  1.39956586e-02,  9.04677063e-03,\n",
      "        -7.08316686e-03,  2.15829071e-03, -6.01198385e-03,\n",
      "        -9.25715454e-03],\n",
      "       [-7.33381557e+00, -7.44178009e+00,  9.98631096e+00,\n",
      "        -7.04955721e+00, -6.97107840e+00,  2.62563858e+01,\n",
      "        -7.43322659e+00],\n",
      "       [ 9.02867131e-03,  3.22606112e-03, -8.75808019e-03,\n",
      "         7.52681773e-03,  4.86147776e-03,  4.81323060e-03,\n",
      "         2.40292586e-02],\n",
      "       [-6.69383584e-03,  9.62360203e-03,  1.09238373e-02,\n",
      "         4.54173796e-03,  5.36167854e-03, -8.98423139e-03,\n",
      "        -1.41656419e-04],\n",
      "       [ 1.88929159e-02,  1.62733917e-03, -1.02146305e-02,\n",
      "         1.74517687e-02,  3.08617600e-03, -1.20777069e-02,\n",
      "         9.86434519e-03],\n",
      "       [ 8.02630931e-03,  4.26331256e-03,  1.34296110e-02,\n",
      "        -1.16779264e-02, -6.55866647e-03, -5.55921905e-03,\n",
      "        -3.53876414e-04],\n",
      "       [-2.25518203e+00, -2.29606080e+00,  3.08542776e+00,\n",
      "        -2.20461392e+00, -2.15247393e+00,  8.12841702e+00,\n",
      "        -2.33190489e+00],\n",
      "       [ 1.62632344e-03, -1.13341017e-02,  5.46714012e-03,\n",
      "        -4.91014170e-03,  1.35129839e-02,  1.08377524e-02,\n",
      "        -8.02616775e-03],\n",
      "       [-1.72461700e-02, -4.72398754e-03,  2.21708938e-02,\n",
      "        -1.06455525e-02, -1.53461332e-02,  4.95541142e-04,\n",
      "         1.68536901e-02],\n",
      "       [ 2.16011405e-02,  2.90551782e-03, -1.45375747e-02,\n",
      "         1.10551110e-02,  3.09774373e-03,  1.72720160e-02,\n",
      "        -1.14590013e-02],\n",
      "       [-1.22162304e-03,  6.73977006e-03, -1.20173497e-02,\n",
      "         2.62198341e-03,  1.38052600e-02,  5.56129124e-03,\n",
      "         1.03985537e-02],\n",
      "       [-3.78480833e-03,  4.06509172e-03, -2.72586173e-03,\n",
      "         1.22060115e-02, -3.82714462e-03, -2.45974283e-03,\n",
      "         7.34328467e-04],\n",
      "       [ 1.00428006e-02, -6.18942082e-04,  2.02048961e-02,\n",
      "         6.88806549e-03,  1.76206100e-04,  1.11210449e-02,\n",
      "        -3.06158024e-03],\n",
      "       [ 2.94981245e-03, -7.43757281e-03, -3.47154564e-03,\n",
      "         1.81804560e-02,  2.78334296e-03, -2.64144270e-04,\n",
      "         5.46288025e-03],\n",
      "       [-9.96791106e-03,  8.61121714e-03, -1.11579821e-02,\n",
      "         6.99002878e-04,  2.93778977e-03, -2.11555790e-02,\n",
      "         6.70472067e-03],\n",
      "       [ 7.58774625e-03, -9.46320221e-03, -1.51735276e-03,\n",
      "        -9.69663821e-03, -2.06177775e-03,  1.77018158e-03,\n",
      "         1.50688486e-02],\n",
      "       [-1.12750614e-02,  1.14820562e-02, -1.00116907e-02,\n",
      "        -1.03512164e-02, -2.65574967e-03,  1.51991844e-03,\n",
      "        -9.18539008e-04],\n",
      "       [-1.13466215e+01, -1.15551090e+01,  1.55056782e+01,\n",
      "        -1.09286118e+01, -1.07936153e+01,  4.06890831e+01,\n",
      "        -1.15504103e+01],\n",
      "       [-8.56029976e-04, -1.00283127e-03,  9.33236722e-03,\n",
      "         8.17381777e-03,  5.18953311e-04, -8.95066652e-03,\n",
      "         1.10021490e-03],\n",
      "       [-5.43868123e-03,  1.36277610e-02, -5.60163520e-03,\n",
      "        -9.90571105e-04, -1.00633660e-02, -6.22560130e-03,\n",
      "         3.69148050e-03],\n",
      "       [ 1.42669084e-03, -9.48405825e-03, -1.53188063e-02,\n",
      "        -1.50805945e-02,  2.01605838e-02,  6.75952109e-03,\n",
      "        -4.95190872e-03],\n",
      "       [-1.91933708e-03,  1.34828687e-02, -1.04036592e-02,\n",
      "        -9.93888476e-04,  1.39908362e-02,  5.42675145e-03,\n",
      "         1.31656034e-02],\n",
      "       [-7.46883869e+00, -7.59583759e+00,  1.01791191e+01,\n",
      "        -7.16686964e+00, -7.10345650e+00,  2.67260113e+01,\n",
      "        -7.55773211e+00],\n",
      "       [-8.30317020e+00, -8.44512844e+00,  1.13142004e+01,\n",
      "        -7.98387432e+00, -7.86865616e+00,  2.96909962e+01,\n",
      "        -8.44413948e+00],\n",
      "       [-5.66216183e+00, -5.77589130e+00,  7.75240564e+00,\n",
      "        -5.45707369e+00, -5.38524818e+00,  2.03016872e+01,\n",
      "        -5.78242540e+00],\n",
      "       [-7.68417120e-03,  1.19979689e-02,  3.07967188e-03,\n",
      "        -4.26376471e-03,  1.25023921e-03,  6.22692984e-03,\n",
      "         2.98868632e-03],\n",
      "       [ 7.89726991e-03, -1.06854169e-02,  1.27869472e-02,\n",
      "         1.43679986e-02, -4.56493720e-03,  1.51716080e-03,\n",
      "         6.53003226e-04],\n",
      "       [ 6.77318824e-03,  1.14690829e-02,  5.22435317e-03,\n",
      "        -5.20466221e-03,  2.33171917e-02, -1.25126457e-02,\n",
      "         1.46009708e-02],\n",
      "       [-9.47053730e-03,  1.36938440e-02,  8.14578775e-03,\n",
      "         8.57250858e-03,  7.41413841e-03,  5.84573206e-03,\n",
      "         1.17821684e-02],\n",
      "       [-6.06535515e-03,  3.30756884e-03,  5.74090471e-03,\n",
      "        -1.11015234e-03, -3.99014866e-03,  8.15952290e-03,\n",
      "         1.20898858e-02],\n",
      "       [ 3.45717813e-03,  6.27425825e-03,  6.48623158e-04,\n",
      "         2.60872860e-02, -4.95703379e-03,  4.48500877e-03,\n",
      "         3.44471005e-03],\n",
      "       [ 7.46332156e-03, -2.15045847e-02,  6.84807729e-03,\n",
      "        -2.09729820e-02,  4.85602813e-03,  3.06886644e-03,\n",
      "        -1.20068761e-02],\n",
      "       [ 7.47603225e-03,  8.57086200e-03, -9.75384004e-03,\n",
      "        -6.93832291e-03,  9.94026382e-03, -3.26356781e-03,\n",
      "        -1.44198388e-02],\n",
      "       [ 6.44598482e-03, -1.46521535e-03, -8.02924111e-03,\n",
      "        -1.26858661e-02,  6.70234347e-03, -4.88430075e-03,\n",
      "         6.95254141e-03],\n",
      "       [-1.05395112e-02,  4.77318885e-03, -9.08821449e-03,\n",
      "        -2.15563783e-03,  4.60964721e-03, -1.72874948e-03,\n",
      "        -9.31204017e-03],\n",
      "       [-4.57184267e+00, -4.61311722e+00,  6.22148705e+00,\n",
      "        -4.39223099e+00, -4.35382748e+00,  1.63422260e+01,\n",
      "        -4.61900997e+00],\n",
      "       [-5.72575903e+00, -5.82132196e+00,  7.78577757e+00,\n",
      "        -5.49610090e+00, -5.42948723e+00,  2.04548550e+01,\n",
      "        -5.79529858e+00],\n",
      "       [-2.79516783e-02,  1.36684980e-02, -1.19638373e-03,\n",
      "         5.90700051e-03, -1.08773112e-02, -2.63076695e-03,\n",
      "        -4.96649044e-03],\n",
      "       [-7.47012300e-03,  8.75923876e-03,  2.25926540e-03,\n",
      "        -9.97500774e-03,  7.11163366e-03,  2.89358245e-03,\n",
      "         4.73866425e-03],\n",
      "       [-7.37383496e-04,  6.36800192e-04,  1.19799273e-02,\n",
      "         4.72359499e-03, -1.46303726e-02, -7.73388892e-03,\n",
      "        -2.80391541e-03],\n",
      "       [-8.12171936e+00, -8.23792362e+00,  1.10559330e+01,\n",
      "        -7.80190706e+00, -7.72176170e+00,  2.90637589e+01,\n",
      "        -8.23308277e+00],\n",
      "       [ 3.43902851e-03,  7.69946666e-04,  1.26938084e-02,\n",
      "         1.26634110e-02, -2.88947322e-03, -3.91164050e-03,\n",
      "         1.23805157e-03],\n",
      "       [-1.22534577e-02,  1.65379681e-02, -1.68769376e-03,\n",
      "         5.86400041e-03,  4.00504423e-03,  3.13585224e-05,\n",
      "         2.49811867e-03],\n",
      "       [ 1.27100218e-02,  1.15483496e-02, -6.74086949e-03,\n",
      "        -2.39019538e-03, -1.95149987e-04, -4.24581720e-03,\n",
      "        -1.60237178e-02],\n",
      "       [-5.80541207e-04,  3.68732377e-03, -4.58522991e-04,\n",
      "        -9.60281491e-03,  1.73103884e-02, -1.14519487e-03,\n",
      "         1.94285079e-04],\n",
      "       [ 4.66414541e-03, -2.45105615e-03, -3.56976944e-03,\n",
      "        -1.00249574e-02, -2.67415196e-02,  1.17552727e-02,\n",
      "         1.03959097e-02],\n",
      "       [ 2.38006935e-03,  8.78463779e-03, -3.65609158e-04,\n",
      "         5.10359602e-03,  9.72790364e-03,  1.90559274e-03,\n",
      "        -8.54096096e-03],\n",
      "       [-1.31240766e-03,  7.72109441e-03, -7.62961153e-03,\n",
      "         2.33820900e-02, -1.37993358e-02, -2.47004554e-02,\n",
      "        -3.14865308e-03],\n",
      "       [ 2.38493625e-02, -1.00639276e-02, -1.66255906e-02,\n",
      "        -1.04901846e-02,  1.71359219e-02,  3.46532767e-03,\n",
      "        -1.88651644e-02],\n",
      "       [-7.48514256e-04,  9.11393261e-04, -1.37545948e-03,\n",
      "         7.37898471e-03,  9.01585072e-03,  8.08385666e-03,\n",
      "         5.80555620e-03],\n",
      "       [-2.33246779e+00, -2.35702372e+00,  3.17377949e+00,\n",
      "        -2.25758100e+00, -2.21734619e+00,  8.35234070e+00,\n",
      "        -2.33519149e+00],\n",
      "       [ 7.45888008e-03,  1.66034549e-02,  5.69120841e-03,\n",
      "        -1.24395872e-02,  1.16653498e-02,  1.30123869e-02,\n",
      "         7.86039513e-03],\n",
      "       [-9.73387957e-01, -9.80107486e-01,  1.32415056e+00,\n",
      "        -9.31751966e-01, -9.27954614e-01,  3.48961329e+00,\n",
      "        -9.81512904e-01],\n",
      "       [-5.80337970e-03,  9.41497833e-03,  8.06426350e-03,\n",
      "         4.73152939e-03, -4.20396170e-03,  4.18418413e-03,\n",
      "        -1.25556421e-02],\n",
      "       [-1.57333314e+00, -1.62756586e+00,  2.16364264e+00,\n",
      "        -1.50808203e+00, -1.52072060e+00,  5.68011618e+00,\n",
      "        -1.61353087e+00],\n",
      "       [-7.25487185e+00, -7.37566185e+00,  9.90634727e+00,\n",
      "        -6.97828341e+00, -6.90453434e+00,  2.59853973e+01,\n",
      "        -7.36834288e+00],\n",
      "       [-4.51126099e+00, -4.58535767e+00,  6.12118626e+00,\n",
      "        -4.33064222e+00, -4.28926325e+00,  1.61323013e+01,\n",
      "        -4.58985043e+00],\n",
      "       [-3.03285933e+00, -3.07640862e+00,  4.12406063e+00,\n",
      "        -2.89964104e+00, -2.87729740e+00,  1.08204298e+01,\n",
      "        -3.06991744e+00],\n",
      "       [ 5.11945784e-03,  8.25598557e-03, -7.11416965e-03,\n",
      "        -6.28586253e-03,  9.88419447e-03, -6.53124508e-03,\n",
      "         1.37922885e-02],\n",
      "       [ 3.44852149e-03,  8.64699297e-03, -1.60147529e-02,\n",
      "        -8.88988958e-04, -2.53874925e-03, -3.93393682e-03,\n",
      "         3.42588359e-03],\n",
      "       [-3.51662375e-03,  8.02879035e-03,  7.96551630e-03,\n",
      "         1.05292285e-02, -2.42968481e-02, -1.98662397e-03,\n",
      "        -7.99137633e-03],\n",
      "       [ 1.53445490e-02, -8.80418066e-03, -1.69282500e-02,\n",
      "        -1.18849333e-02,  5.10878721e-03,  1.16355531e-02,\n",
      "         6.31901203e-03],\n",
      "       [ 1.69924311e-02,  4.54971241e-03, -3.80055024e-03,\n",
      "        -8.45329557e-03, -1.37875704e-02, -3.76761454e-04,\n",
      "         9.44503210e-03],\n",
      "       [ 1.25891734e-02,  1.55286733e-02,  7.72145286e-05,\n",
      "         2.00749207e-02, -9.80000384e-03, -8.64620600e-03,\n",
      "        -1.02823758e-02],\n",
      "       [-6.98176050e+00, -7.08410168e+00,  9.50850677e+00,\n",
      "        -6.71863222e+00, -6.63374853e+00,  2.50097427e+01,\n",
      "        -7.07887125e+00],\n",
      "       [-1.70109570e-02, -8.02952424e-03, -1.46700733e-03,\n",
      "        -3.39253014e-03, -5.20012761e-03,  1.90237761e-02,\n",
      "        -4.24596760e-03],\n",
      "       [ 4.06803656e-03,  1.78170279e-02,  6.24854211e-03,\n",
      "        -6.23275433e-03, -3.14142974e-03, -1.07603883e-02,\n",
      "         8.89648974e-04],\n",
      "       [ 1.67515799e-02,  5.26508316e-03, -7.19598494e-03,\n",
      "         7.01875146e-03,  1.94742065e-02,  2.22609024e-02,\n",
      "        -1.07546235e-02],\n",
      "       [-1.12352362e+01, -1.14151258e+01,  1.53085060e+01,\n",
      "        -1.08027391e+01, -1.06692801e+01,  4.02568207e+01,\n",
      "        -1.14208384e+01],\n",
      "       [-1.21008401e-04,  2.18944787e-03, -1.02690347e-02,\n",
      "         1.39459474e-02,  9.44015197e-03, -1.21001638e-02,\n",
      "        -4.52996930e-04],\n",
      "       [-4.56731319e+00, -4.64840174e+00,  6.25178576e+00,\n",
      "        -4.42586184e+00, -4.36534166e+00,  1.64251957e+01,\n",
      "        -4.64050198e+00],\n",
      "       [-1.11902323e+01, -1.13807869e+01,  1.52428112e+01,\n",
      "        -1.07576599e+01, -1.06272726e+01,  4.00722694e+01,\n",
      "        -1.13608208e+01],\n",
      "       [-7.82247353e+00, -7.93120813e+00,  1.06514711e+01,\n",
      "        -7.53604746e+00, -7.43801689e+00,  2.80364838e+01,\n",
      "        -7.94799948e+00],\n",
      "       [-1.62595520e+01, -1.64710789e+01,  2.21156731e+01,\n",
      "        -1.56293535e+01, -1.54605970e+01,  5.81977043e+01,\n",
      "        -1.64637337e+01],\n",
      "       [-5.75296907e-03,  8.93194973e-03,  1.23904981e-02,\n",
      "         1.41957160e-02,  4.42796992e-03,  4.45373077e-03,\n",
      "        -8.36576521e-03],\n",
      "       [-1.15610305e-02, -4.88896342e-03, -1.20765576e-02,\n",
      "        -2.56475937e-02, -1.95098426e-02, -6.37684204e-03,\n",
      "         6.63520582e-03],\n",
      "       [-2.20361290e+01, -2.24085121e+01,  3.00333061e+01,\n",
      "        -2.11620808e+01, -2.09377308e+01,  7.89235840e+01,\n",
      "        -2.23864613e+01],\n",
      "       [-1.55591837e-03,  3.41278361e-03, -7.47419195e-04,\n",
      "        -2.32900237e-03,  6.03694469e-03, -3.07815638e-03,\n",
      "        -2.39110906e-02],\n",
      "       [-1.61873221e+00, -1.62272847e+00,  2.17622280e+00,\n",
      "        -1.56446612e+00, -1.54764998e+00,  5.76852942e+00,\n",
      "        -1.62219417e+00],\n",
      "       [-1.70500576e+00, -1.73394084e+00,  2.30668187e+00,\n",
      "        -1.63285482e+00, -1.61055934e+00,  6.10220194e+00,\n",
      "        -1.71522903e+00],\n",
      "       [-2.95227543e-02, -7.95140397e-03,  1.56437531e-02,\n",
      "         5.29354112e-03,  1.79456780e-04, -1.21933473e-02,\n",
      "         3.55989463e-03],\n",
      "       [ 8.80683586e-03,  5.04600443e-03, -1.69482583e-03,\n",
      "         1.07700657e-02, -8.50117393e-03, -7.22834608e-03,\n",
      "         1.75444297e-02],\n",
      "       [-9.05940056e+00, -9.19135189e+00,  1.23293676e+01,\n",
      "        -8.70809174e+00, -8.60548782e+00,  3.24147072e+01,\n",
      "        -9.19037151e+00],\n",
      "       [ 2.87559000e-03, -8.25631432e-03,  7.04111159e-03,\n",
      "        -2.37314403e-02, -4.57078312e-03,  3.35643505e-04,\n",
      "        -1.29328873e-02],\n",
      "       [ 4.59383940e-03, -1.37127629e-02, -2.68566841e-03,\n",
      "         4.77457372e-03,  1.08231828e-02,  1.87916111e-03,\n",
      "         6.24818122e-03],\n",
      "       [-1.59825850e-02,  4.61590383e-03,  7.01797288e-03,\n",
      "        -6.20370591e-03, -8.34321894e-04, -5.75719122e-03,\n",
      "         3.97415133e-03],\n",
      "       [ 7.03130383e-03,  2.77546905e-02, -7.64628407e-03,\n",
      "         1.92800199e-03, -7.90680107e-03, -3.89903085e-03,\n",
      "        -3.40028433e-03],\n",
      "       [-4.50647163e+00, -4.59815311e+00,  6.14307356e+00,\n",
      "        -4.32129860e+00, -4.28642654e+00,  1.61727009e+01,\n",
      "        -4.58071852e+00],\n",
      "       [-9.28103256e+00, -9.42728615e+00,  1.26328163e+01,\n",
      "        -8.90517807e+00, -8.81997681e+00,  3.31865768e+01,\n",
      "        -9.39961815e+00],\n",
      "       [ 1.24841109e-02, -8.46529473e-03, -1.06279284e-03,\n",
      "        -4.50919895e-03,  2.01034453e-03,  6.98173931e-03,\n",
      "         1.59351539e-03],\n",
      "       [-1.55112374e+00, -1.57635021e+00,  2.14229250e+00,\n",
      "        -1.50016165e+00, -1.46955776e+00,  5.58215094e+00,\n",
      "        -1.56992352e+00],\n",
      "       [ 6.32024719e-04,  8.61488190e-03, -8.20358563e-03,\n",
      "        -2.47113593e-03, -2.11310200e-03, -1.84177887e-04,\n",
      "        -3.98610486e-03],\n",
      "       [-1.15173864e+01, -1.17123852e+01,  1.56801062e+01,\n",
      "        -1.10799742e+01, -1.09388657e+01,  4.12207832e+01,\n",
      "        -1.16941013e+01],\n",
      "       [ 3.65351490e-03, -3.17907799e-03, -4.25604125e-03,\n",
      "         5.54517237e-03, -1.80684787e-04,  7.58556882e-04,\n",
      "        -1.33797002e-03],\n",
      "       [ 7.18388893e-03, -1.91843770e-02, -4.01440251e-04,\n",
      "         2.12118011e-02, -1.29199680e-02,  1.11826137e-03,\n",
      "         2.12392583e-02],\n",
      "       [ 2.78871413e-03, -5.71626192e-03, -5.12762507e-03,\n",
      "         1.34568857e-02,  1.70601439e-03,  5.44311153e-03,\n",
      "        -4.85346513e-03],\n",
      "       [-6.79550646e-03, -1.20632658e-02,  3.49226338e-03,\n",
      "         1.31541979e-03,  3.34482943e-03, -2.21167970e-02,\n",
      "         1.41299982e-03],\n",
      "       [-5.25172424e+00, -5.34935999e+00,  7.16095018e+00,\n",
      "        -5.04209661e+00, -4.98744678e+00,  1.88151112e+01,\n",
      "        -5.34045029e+00],\n",
      "       [ 4.57903091e-03, -2.86875269e-03, -1.17423171e-02,\n",
      "         4.11790004e-03,  6.55229390e-03,  4.97741671e-03,\n",
      "        -1.57830082e-02],\n",
      "       [-8.01514369e-03,  5.47602121e-03,  9.54920240e-03,\n",
      "         1.22941984e-02,  1.69288891e-03, -5.21602528e-03,\n",
      "         6.49594842e-03],\n",
      "       [-3.46966815e+00, -3.54615784e+00,  4.74267626e+00,\n",
      "        -3.34131694e+00, -3.33749294e+00,  1.24798164e+01,\n",
      "        -3.54669785e+00],\n",
      "       [-9.83893573e-01, -9.83317554e-01,  1.35889280e+00,\n",
      "        -9.61072922e-01, -9.21885788e-01,  3.55606008e+00,\n",
      "        -1.01044858e+00],\n",
      "       [ 6.69422187e-03, -7.34642940e-03, -4.54162247e-03,\n",
      "        -9.71861184e-03, -9.32834204e-03,  1.58640277e-02,\n",
      "        -2.00623646e-03],\n",
      "       [-9.08224564e-03,  3.33716453e-05,  6.99682254e-03,\n",
      "         1.47401588e-02,  1.52977751e-04, -2.23430973e-02,\n",
      "        -1.06945103e-02],\n",
      "       [-2.40677339e-03, -6.54417090e-03, -8.17077234e-03,\n",
      "        -9.96397808e-03,  9.79982503e-03, -1.03539806e-02,\n",
      "         1.17314141e-02],\n",
      "       [-1.02055635e+01, -1.03819542e+01,  1.39008589e+01,\n",
      "        -9.81601238e+00, -9.71861935e+00,  3.65870018e+01,\n",
      "        -1.03764143e+01],\n",
      "       [-1.29399337e-02, -8.82760808e-03,  8.32744688e-03,\n",
      "         5.17643457e-05, -1.28714712e-02,  1.14479847e-02,\n",
      "         8.54625553e-03],\n",
      "       [-1.21517344e+01, -1.23277006e+01,  1.65511169e+01,\n",
      "        -1.17125425e+01, -1.15504141e+01,  4.35195236e+01,\n",
      "        -1.23118916e+01],\n",
      "       [-1.12352200e+01, -1.14260607e+01,  1.53182106e+01,\n",
      "        -1.08141546e+01, -1.06762533e+01,  4.02766685e+01,\n",
      "        -1.14229250e+01],\n",
      "       [-3.23826766e+00, -3.26992035e+00,  4.39728117e+00,\n",
      "        -3.08116603e+00, -3.07644820e+00,  1.15475206e+01,\n",
      "        -3.27242255e+00],\n",
      "       [-6.60257006e+00, -6.70513153e+00,  9.01069164e+00,\n",
      "        -6.34790659e+00, -6.29539156e+00,  2.36551914e+01,\n",
      "        -6.68676996e+00],\n",
      "       [-8.95123768e+00, -9.06781483e+00,  1.21790276e+01,\n",
      "        -8.58454132e+00, -8.49886703e+00,  3.20257874e+01,\n",
      "        -9.07110786e+00],\n",
      "       [-2.20177031e+00, -2.24435377e+00,  3.00059748e+00,\n",
      "        -2.10903358e+00, -2.10233545e+00,  7.90702534e+00,\n",
      "        -2.23762178e+00],\n",
      "       [ 1.75844296e-03, -7.84207880e-03, -1.17232846e-02,\n",
      "         1.17919187e-03, -1.61634833e-02, -1.72208284e-03,\n",
      "        -1.03150471e-03]], dtype=float32), 'b2': Array([-17.728003, -19.979641,  25.528618, -16.001337, -15.690265,\n",
      "        63.69301 , -19.82239 ], dtype=float32)}\n"
     ]
    }
   ],
   "source": [
    "print(ac.actor_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mujoco\n",
    "from mujoco import mjx\n",
    "import jax\n",
    "import mediapy as media\n",
    "\n",
    "\n",
    "\n",
    "# Make model, data, and renderer\n",
    "mj_model = mujoco.MjModel.from_xml_path('aloha/scene.xml')\n",
    "mj_data = mujoco.MjData(mj_model)\n",
    "renderer = mujoco.Renderer(mj_model)\n",
    "\n",
    "\n",
    "\n",
    "mjx_model = mjx.put_model(mj_model)\n",
    "mjx_data = mjx.put_data(mj_model, mj_data)\n",
    "\n",
    "# print(mj_data.qpos, type(mj_data.qpos))\n",
    "# print(mjx_data.qpos, type(mjx_data.qpos), mjx_data.qpos.devices())\n",
    "\n",
    "\n",
    "# enable joint visualization option:\n",
    "scene_option = mujoco.MjvOption()\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True\n",
    "\n",
    "duration = 3.8  # (seconds)\n",
    "framerate = 30  # (Hz)\n",
    "\n",
    "frames = []\n",
    "mujoco.mj_resetData(mj_model, mj_data)\n",
    "# while mj_data.time < duration:\n",
    "#   mujoco.mj_step(mj_model, mj_data)\n",
    "#   if len(frames) < mj_data.time * framerate:\n",
    "#     renderer.update_scene(mj_data, scene_option=scene_option)\n",
    "#     pixels = renderer.render()\n",
    "#     frames.append(pixels)\n",
    "\n",
    "# # Simulate and display video.\n",
    "# media.show_video(frames, fps=framerate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "rng = jax.random.PRNGKey(0)\n",
    "rng = jax.random.split(rng, 24)\n",
    "\n",
    "\n",
    "fun=lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (15,)))\n",
    "fun_vmapped = jax.vmap(fun)\n",
    "batch=fun_vmapped(rng)\n",
    "\n",
    "\n",
    "jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))\n",
    "batch = jit_step(mjx_model, batch)\n",
    "\n",
    "\n",
    "print(batch.qpos)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}