Add cell to tutorial colab showing how to add visual geoms to a scene.

PiperOrigin-RevId: 456044324
Change-Id: I2c9efe85caa75356ed6b6a41f6ad0d5e97ffee87

Author: yuvaltassa

dm_control/mujoco/tutorial.ipynb

@@ -477,6 +477,77 @@
         "ax.set_axis_off()"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "R_jkaXd_9Eiu"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Adding arbitrary 3D geometry {vertical-output: true}\n",
+        "\n",
+        "def get_geom_speed(physics, geom_name):\n",
+        "  \"\"\"Returns the speed of a geom.\"\"\"\n",
+        "  geom_vel = np.zeros(6)\n",
+        "  geom_type = mujoco.mjtObj.mjOBJ_GEOM\n",
+        "  geom_id = mujoco.mj_name2id(physics.model.ptr, geom_type, geom_name)\n",
+        "  mujoco.mj_objectVelocity(physics.model.ptr, physics.data.ptr,\n",
+        "                           geom_type, geom_id, geom_vel, 0)\n",
+        "  return np.linalg.norm(geom_vel)\n",
+        "\n",
+        "def add_visual_capsule(scene, point1, point2, radius, rgba):\n",
+        "  \"\"\"Adds one capsule to an mjvScene.\"\"\"\n",
+        "  if scene.ngeom \u003e= scene.maxgeom:\n",
+        "    return\n",
+        "  scene.ngeom += 1  # increment ngeom\n",
+        "  # initialise a new capsule, add it to the scene using mjv_makeConnector\n",
+        "  mujoco.mjv_initGeom(scene.geoms[scene.ngeom-1],\n",
+        "                      mujoco.mjtGeom.mjGEOM_CAPSULE, np.zeros(3),\n",
+        "                      np.zeros(3), np.zeros(9), rgba.astype(np.float32))\n",
+        "  mujoco.mjv_makeConnector(scene.geoms[scene.ngeom-1],\n",
+        "                           mujoco.mjtGeom.mjGEOM_CAPSULE, radius,\n",
+        "                           point1[0], point1[1], point1[2],\n",
+        "                           point2[0], point2[1], point2[2])\n",
+        "\n",
+        " # traces of time, position and speed\n",
+        "time = []\n",
+        "position = []\n",
+        "speed = []\n",
+        "offset = physics.model.jnt_axis[0]/8  # offset along the joint axis\n",
+        "\n",
+        "def scene_callback(physics, scn):\n",
+        "  \"\"\"Draw position trace, speed modifies width and colours.\"\"\"\n",
+        "  if len(position) \u003e 1:\n",
+        "    for i in range(len(position)-1):\n",
+        "      rgba=np.array((np.clip(speed[i]/10, 0, 1),\n",
+        "                     np.clip(1-speed[i]/10, 0, 1),\n",
+        "                     .5, 1.))\n",
+        "      radius=.003*(1+speed[i])\n",
+        "      point1 = position[i] + offset*time[i]\n",
+        "      point2 = position[i+1] + offset*time[i+1]\n",
+        "      add_visual_capsule(scn, point1, point2, radius, rgba)\n",
+        "\n",
+        "duration = 6    # (seconds)\n",
+        "framerate = 30  # (Hz)\n",
+        "\n",
+        "# Simulate and display video.\n",
+        "frames = []\n",
+        "physics.reset()  # Reset state and time\n",
+        "while physics.data.time \u003c duration:\n",
+        "  # append data to the traces\n",
+        "  position.append(physics.named.data.geom_xpos[\"green_sphere\"].copy())\n",
+        "  time.append(physics.data.time)\n",
+        "  speed.append(get_geom_speed(physics, \"green_sphere\"))\n",
+        "  physics.step()\n",
+        "  if len(frames) \u003c physics.data.time * framerate:\n",
+        "    camera = mujoco.Camera(physics, max_geom=10000,\n",
+        "                           scene_callback=scene_callback)\n",
+        "    pixels = camera.render()\n",
+        "    frames.append(pixels)\n",
+        "display_video(frames, framerate)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {