Merge pull request #159 from s1lent4gnt:lilkm/fix-humanoid_joystick-nb

PiperOrigin-RevId: 794777571
Change-Id: Ia069206a21e563111c33ae2198e6b41e0a942ce4

Author: copybara-github

mujoco_playground/experimental/learning/humanoid_joystick.ipynb

@@ -172,8 +172,6 @@
     "ep_length_mean, ep_length_std = [], []\n",
     "times = [datetime.now()]\n",
     "\n",
-    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
-    "\n",
     "\n",
     "def progress(num_steps, metrics):\n",
     "  # Log to wandb.\n",
@@ -186,20 +184,12 @@
     "  x_data.append(num_steps)\n",
     "  y_data.append(metrics[\"eval/episode_reward\"])\n",
     "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
-    "  ep_length_mean.append(metrics[\"eval/avg_episode_length\"])\n",
-    "  ep_length_std.append(metrics[\"eval/avg_episode_length_std\"])\n",
-    "\n",
-    "  axes[0].set_xlim([0, ppo_params.num_timesteps * 1.25])\n",
-    "  axes[0].set_xlabel(\"# environment steps\")\n",
-    "  axes[0].set_ylabel(\"reward per episode\")\n",
-    "  axes[0].set_title(f\"y={y_data[-1]:.3f}\")\n",
-    "  axes[0].errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
     "\n",
-    "  axes[1].set_xlim([0, ppo_params.num_timesteps * 1.25])\n",
-    "  axes[1].set_xlabel(\"# environment steps\")\n",
-    "  axes[1].set_ylabel(\"episode length\")\n",
-    "  axes[1].set_title(f\"y={ep_length_mean[-1]:.3f}\")\n",
-    "  axes[1].errorbar(x_data, ep_length_mean, yerr=ep_length_std, color=\"blue\")\n",
+    "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
+    "  plt.xlabel(\"# environment steps\")\n",
+    "  plt.ylabel(\"reward per episode\")\n",
+    "  plt.title(f\"y={y_data[-1]:.3f}\")\n",
+    "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
     "\n",
     "  display(plt.gcf())\n",
     "\n",
@@ -357,11 +347,11 @@
     "    rewards.append(\n",
     "        {k[7:]: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
     "    )\n",
-    "    linvel.append(eval_env.get_local_linvel(state.data))\n",
+    "    linvel.append(eval_env.get_local_linvel(state.data, \"pelvis\"))\n",
     "    angvel.append(eval_env.get_gyro(state.data))\n",
     "    track.append(\n",
     "        eval_env._reward_tracking_lin_vel(\n",
-    "            state.info[\"command\"], eval_env.get_local_linvel(state.data)\n",
+    "            state.info[\"command\"], eval_env.get_local_linvel(state.data, \"pelvis\")\n",
     "        )\n",
     "    )\n",
     "\n",
@@ -381,7 +371,7 @@
     "    qvels.append(qvel)\n",
     "    qpos_cost.append(jp.sum(jp.square(state.data.qpos[7:] - eval_env._default_pose)))\n",
     "\n",
-    "    xyz = np.array(state.data.xpos[eval_env.mj_model.body(\"torso\").id])\n",
+    "    xyz = np.array(state.data.xpos[eval_env.mj_model.body(\"torso_link\").id])\n",
     "    xyz += np.array([0, 0.0, 0])\n",
     "    x_axis = state.data.xmat[eval_env._torso_body_id, 0]\n",
     "    yaw = -np.arctan2(x_axis[1], x_axis[0])\n",