Explicitly convert arrays to scalars

google-deepmind · Dec 16, 2023 · c3d87e6 · c3d87e6
1 parent dc0d0c5
commit c3d87e6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/tutorial.ipynb b/python/tutorial.ipynb
@@ -1080,7 +1080,7 @@
         "  for i in range(n_steps):\n",
         "    mujoco.mj_step(model, data)\n",
         "    sim_time[i] = data.time\n",
-        "    angle[i] = data.joint('root').qpos\n",
+        "    angle[i] = data.joint('root').qpos[0]\n",
         "    energy[i] = data.energy[0] + data.energy[1]\n",
         "\n",
         "  # plot\n",
@@ -1198,7 +1198,7 @@
         "  for i in range(n_steps):\n",
         "    mujoco.mj_step(model, data)\n",
         "    if data.warning.number.any():\n",
-        "      warning_index = np.nonzero(data.warning.number)[0]\n",
+        "      warning_index = np.nonzero(data.warning.number)[0][0]\n",
         "      warning = mujoco.mjtWarning(warning_index).name\n",
         "      print(f'stopped due to divergence ({warning}) at timestep {i}.\\n')\n",
         "      break\n",