From 8fa7fe7f5b2463b7a0978049be40710e3a511e4c Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Tue, 10 Sep 2024 04:20:21 -0700
Subject: [PATCH] Add a tutorial section on multiple frames in the same scene.
 Fixes #1757, #1356.

PiperOrigin-RevId: 672902542
Change-Id: Ie9fae2151354206b1418c129e2bd2ab016695ec3
---
 python/tutorial.ipynb | 76 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)

diff --git a/python/tutorial.ipynb b/python/tutorial.ipynb
index 36f76de5..15b5ad39 100644
--- a/python/tutorial.ipynb
+++ b/python/tutorial.ipynb
@@ -1925,6 +1925,82 @@
     "media.show_video(frames, fps=framerate)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "p6wHrjRg1EGF"
+   },
+   "source": [
+    "## Multiple frames in the same scene\n",
+    "\n",
+    "Sometimes one would like to draw the same geometry multiple times, for example when a model is tracking states from motion-capture, it's nice to have the data\n",
+    "visualized next to the model. Unlike `mjv_updateScene` (called by the `Renderer`'s `update_scene` method) which clears the scene at every call, `mjv_addGeoms` will add visual geoms to an existing scene:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "T4b_8n6t1ASk"
+   },
+   "outputs": [],
+   "source": [
+    "# Get MuJoCo's standard humanoid model.\n",
+    "print('Getting MuJoCo humanoid XML description from GitHub:')\n",
+    "!git clone https://github.com/google-deepmind/mujoco\n",
+    "with open('mujoco/model/humanoid/humanoid.xml', 'r') as f:\n",
+    "  xml = f.read()\n",
+    "\n",
+    "# Load the model, make two MjData's.\n",
+    "model = mujoco.MjModel.from_xml_string(xml)\n",
+    "data = mujoco.MjData(model)\n",
+    "data2 = mujoco.MjData(model)\n",
+    "\n",
+    "# Episode parameters.\n",
+    "duration = 3       # (seconds)\n",
+    "framerate = 60     # (Hz)\n",
+    "data.qpos[0:2] = [-.5, -.5]  # Initial x-y position (m)\n",
+    "data.qvel[2] = 4   # Initial vertical velocity (m/s)\n",
+    "ctrl_phase = 2 * np.pi * np.random.rand(model.nu)  # Control phase\n",
+    "ctrl_freq = 1     # Control frequency\n",
+    "\n",
+    "# Visual options for the \"ghost\" model.\n",
+    "vopt2 = mujoco.MjvOption()\n",
+    "vopt2.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True  # Transparent.\n",
+    "pert = mujoco.MjvPerturb()  # Empty MjvPerturb object\n",
+    "# We only want dynamic objects (the humanoid). Static objects (the floor)\n",
+    "# should not be re-drawn. The mjtCatBit flag lets us do that, though we could\n",
+    "# equivalently use mjtVisFlag.mjVIS_STATIC\n",
+    "catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC\n",
+    "\n",
+    "# Simulate and render.\n",
+    "frames = []\n",
+    "with mujoco.Renderer(model, 480, 640) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    # Sinusoidal control signal.\n",
+    "    data.ctrl = np.sin(ctrl_phase + 2 * np.pi * data.time * ctrl_freq)\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      # This draws the regular humanoid from `data`.\n",
+    "      renderer.update_scene(data)\n",
+    "\n",
+    "      # Copy qpos to data2, move the humanoid sideways, call mj_forward.\n",
+    "      data2.qpos = data.qpos\n",
+    "      data2.qpos[0] += 1.5\n",
+    "      data2.qpos[1] += 1\n",
+    "      mujoco.mj_forward(model, data2)\n",
+    "\n",
+    "      # Call mjv_addGeoms to add the ghost humanoid to the scene.\n",
+    "      mujoco.mjv_addGeoms(model, data2, vopt2, pert, catmask, renderer.scene)\n",
+    "\n",
+    "      # Render and add the frame.\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "# Render video at half real-time.\n",
+    "media.show_video(frames, fps=framerate/2)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {