Add cell to tutorial colab demonstrating procedural camera control.

https://youtu.be/cE3s_IfO4g4 PiperOrigin-RevId: 579466532 Change-Id: I433db285bcfa2b9fab0c11e1fef1e25af8187901
google-deepmind · Nov 4, 2023 · c45996b · c45996b
1 parent 7b9a0cd
commit c45996b
Show file tree

Hide file tree

Showing 3 changed files with 228 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -55,6 +55,10 @@ running on Google Colab:
  - For a more advanced example, see the LQR tutorial which creates an LQR
    controller to balance a humanoid on one leg using MuJoCo's dynamics
    derivatives: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb)
+ - The MJX tutorial provides usage examples of
+   [MuJoCo XLA](https://mujoco.readthedocs.io/en/stable/mjx.html), a branch of MuJoCo written in
+   JAX:
+   [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb)
 
 ## Installation
 

diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -29,7 +29,7 @@ General
 MJX
 ^^^
 - Added support for Newton solver (``mjSOL_NEWTON`` in :ref:`mjtSolver`).  The Newton solver significantly speeds up
-  simulation on GPU.  See updated benchmarks on an Nvidia A100 (reported numbers are steps-per-second):
+  simulation on GPU:
 
   .. list-table:: Steps-per-second, Conjugate Gradient vs. Newton on A100
      :header-rows: 1
@@ -73,6 +73,13 @@ Simulate
 
 Documentation
 ^^^^^^^^^^^^^
+
+.. youtube:: cE3s_IfO4g4
+   :align: right
+   :width: 240px
+
+- Added cell to the `tutorial colab <https://github.com/google-deepmind/mujoco#getting-started>`__ providing an example
+  of procedural camera control:
 - Added documentation for the :ref:`UI` framework.
 - Fixed typos and supported fields in docs (fixes :github:issue:`1105` and :github:issue:`1106`).
 

diff --git a/python/tutorial.ipynb b/python/tutorial.ipynb
@@ -72,7 +72,7 @@
       },
       "outputs": [],
       "source": [
-        "#@title Check if installation was successful\n",
+        "#@title Set up rendering, check installation\n",
         "\n",
         "from google.colab import files\n",
         "\n",
@@ -921,7 +921,7 @@
         "id": "g1MKUEL_eSCM"
       },
       "source": [
-        "Below is a model of a chaotic pendulum, similar to [this one](https://www.exploratorium.edu/exhibits/chaotic-pendulum) in the San Francisco Exploratorium. "
+        "Below is a model of a chaotic pendulum, similar to [this one](https://www.exploratorium.edu/exhibits/chaotic-pendulum) in the San Francisco Exploratorium."
       ]
     },
     {
@@ -1242,7 +1242,7 @@
         "free_body_MJCF = \"\"\"\n",
         "\u003cmujoco\u003e\n",
         "  \u003casset\u003e\n",
-        "    \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\" \n",
+        "    \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
         "    rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"edge\" markrgb=\".2 .3 .4\"/\u003e\n",
         "    \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\"\n",
         "    reflectance=\".2\"/\u003e\n",
@@ -1447,7 +1447,7 @@
         "  \u003cworldbody\u003e\n",
         "    \u003clight name=\"light\" pos=\"-.2 0 1\"/\u003e\n",
         "    \u003cgeom name=\"ground\" type=\"plane\" size=\".5 .5 10\" material=\"grid\"\n",
-        "     zaxis=\"-.3 0 1\" friction=\".1\"/\u003e \n",
+        "     zaxis=\"-.3 0 1\" friction=\".1\"/\u003e\n",
         "    \u003ccamera name=\"y\" pos=\"-.1 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/\u003e\n",
         "    \u003cbody pos=\"0 0 .1\"\u003e\n",
         "      \u003cjoint/\u003e\n",
@@ -1941,6 +1941,218 @@
         "    frames.append(pixels)\n",
         "media.show_video(frames, fps=framerate)"
       ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Zzzugf-qPExb"
+      },
+      "source": [
+        "## Camera control\n",
+        "\n",
+        "Cameras can be controlled dynamically in order to achieve cinematic effects. Run the three cells below to see the difference between rendering from a static and moving camera.\n",
+        "\n",
+        "The camera-control code smoothly transitions between two trajectories, one orbiting a fixed point, the other tracking a moving object. Parameter values in the code were obtained by iterating quickly on low-res videos."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-SW-K9WuPGrp"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Load the \"dominos\" model\n",
+        "dominos_xml = \"\"\"\n",
+        "\u003cmujoco\u003e\n",
+        "  \u003casset\u003e\n",
+        "    \u003ctexture type=\"skybox\" builtin=\"gradient\" rgb1=\".3 .5 .7\" rgb2=\"0 0 0\" width=\"32\" height=\"512\"/\u003e\n",
+        "    \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" width=\"512\" height=\"512\" rgb1=\".1 .2 .3\" rgb2=\".2 .3 .4\"/\u003e\n",
+        "    \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\" reflectance=\".2\"/\u003e\n",
+        "  \u003c/asset\u003e\n",
+        "\n",
+        "  \u003cstatistic meansize=\".01\"/\u003e\n",
+        "\n",
+        "  \u003cvisual\u003e\n",
+        "    \u003cglobal offheight=\"2160\" offwidth=\"3840\"/\u003e\n",
+        "    \u003cquality offsamples=\"8\"/\u003e\n",
+        "  \u003c/visual\u003e\n",
+        "\n",
+        "  \u003cdefault\u003e\n",
+        "    \u003cgeom type=\"box\" solref=\".005 1\"/\u003e\n",
+        "    \u003cdefault class=\"static\"\u003e\n",
+        "      \u003cgeom rgba=\".3 .5 .7 1\"/\u003e\n",
+        "    \u003c/default\u003e\n",
+        "  \u003c/default\u003e\n",
+        "\n",
+        "  \u003coption timestep=\"5e-4\"/\u003e\n",
+        "\n",
+        "  \u003cworldbody\u003e\n",
+        "    \u003clight pos=\".3 -.3 .8\" mode=\"trackcom\" diffuse=\"1 1 1\" specular=\".3 .3 .3\"/\u003e\n",
+        "    \u003clight pos=\"0 -.3 .4\" mode=\"targetbodycom\" target=\"box\" diffuse=\".8 .8 .8\" specular=\".3 .3 .3\"/\u003e\n",
+        "    \u003cgeom name=\"floor\" type=\"plane\" size=\"3 3 .01\" pos=\"-0.025 -0.295  0\" material=\"grid\"/\u003e\n",
+        "    \u003cgeom name=\"ramp\" pos=\".25 -.45 -.03\" size=\".04 .1 .07\" euler=\"-30 0 0\" class=\"static\"/\u003e\n",
+        "    \u003ccamera name=\"top\" pos=\"-0.37 -0.78 0.49\" xyaxes=\"0.78 -0.63 0 0.27 0.33 0.9\"/\u003e\n",
+        "\n",
+        "    \u003cbody name=\"ball\" pos=\".25 -.45 .1\"\u003e\n",
+        "      \u003cfreejoint name=\"ball\"/\u003e\n",
+        "      \u003cgeom name=\"ball\" type=\"sphere\" size=\".02\" rgba=\".65 .81 .55 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody pos=\".26 -.3 .03\" euler=\"0 0 -90.0\"\u003e\n",
+        "      \u003cfreejoint/\u003e\n",
+        "      \u003cgeom size=\".0015 .015 .03\" rgba=\"1 .5 .5 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody pos=\".26 -.27 .04\" euler=\"0 0 -81.0\"\u003e\n",
+        "      \u003cfreejoint/\u003e\n",
+        "      \u003cgeom size=\".002 .02 .04\" rgba=\"1 1 .5 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody pos=\".24 -.21 .06\" euler=\"0 0 -63.0\"\u003e\n",
+        "      \u003cfreejoint/\u003e\n",
+        "      \u003cgeom size=\".003 .03 .06\" rgba=\".5 1 .5 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody pos=\".2 -.16 .08\" euler=\"0 0 -45.0\"\u003e\n",
+        "      \u003cfreejoint/\u003e\n",
+        "      \u003cgeom size=\".004 .04 .08\" rgba=\".5 1 1 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody pos=\".15 -.12 .1\" euler=\"0 0 -27.0\"\u003e\n",
+        "      \u003cfreejoint/\u003e\n",
+        "      \u003cgeom size=\".005 .05 .1\" rgba=\".5 .5 1 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody pos=\".09 -.1 .12\" euler=\"0 0 -9.0\"\u003e\n",
+        "      \u003cfreejoint/\u003e\n",
+        "      \u003cgeom size=\".006 .06 .12\" rgba=\"1 .5 1 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody name=\"seasaw_wrapper\" pos=\"-.23 -.1 0\" euler=\"0 0 30\"\u003e\n",
+        "      \u003cgeom size=\".01 .01 .015\" pos=\"0 .05 .015\" class=\"static\"/\u003e\n",
+        "      \u003cgeom size=\".01 .01 .015\" pos=\"0 -.05 .015\" class=\"static\"/\u003e\n",
+        "      \u003cgeom type=\"cylinder\" size=\".01 .0175\" pos=\"-.09 0 .0175\" class=\"static\"/\u003e\n",
+        "      \u003cbody name=\"seasaw\" pos=\"0 0 .03\"\u003e\n",
+        "        \u003cjoint axis=\"0 1 0\"/\u003e\n",
+        "        \u003cgeom type=\"cylinder\" size=\".005 .039\" zaxis=\"0 1 0\" rgba=\".84 .15 .33 1\"/\u003e\n",
+        "        \u003cgeom size=\".1 .02 .005\" pos=\"0 0 .01\" rgba=\".84 .15 .33 1\"/\u003e\n",
+        "      \u003c/body\u003e\n",
+        "    \u003c/body\u003e\n",
+        "\n",
+        "    \u003cbody name=\"box\" pos=\"-.3 -.14 .05501\" euler=\"0 0 -30\"\u003e\n",
+        "      \u003cfreejoint name=\"box\"/\u003e\n",
+        "      \u003cgeom name=\"box\" size=\".01 .01 .01\" rgba=\".0 .7 .79 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "  \u003c/worldbody\u003e\n",
+        "\u003c/mujoco\u003e\n",
+        "\"\"\"\n",
+        "model = mujoco.MjModel.from_xml_string(dominos_xml)\n",
+        "data = mujoco.MjData(model)\n",
+        "renderer = mujoco.Renderer(model, height=1024, width=1440)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "a2WruafiPhPk"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Render from fixed camera\n",
+        "duration = 2.5  # (seconds)\n",
+        "framerate = 60  # (Hz)\n",
+        "\n",
+        "# Simulate and display video.\n",
+        "frames = []\n",
+        "mujoco.mj_resetData(model, data)  # Reset state and time.\n",
+        "while data.time \u003c duration:\n",
+        "  mujoco.mj_step(model, data)\n",
+        "  if len(frames) \u003c data.time * framerate:\n",
+        "    renderer.update_scene(data, camera='top')\n",
+        "    pixels = renderer.render()\n",
+        "    frames.append(pixels)\n",
+        "media.show_video(frames, fps=framerate)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Kie3y-27bQ3J"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Render from moving camera\n",
+        "duration = 3  # (seconds)\n",
+        "\n",
+        "# find time when box is thrown (speed \u003e 2cm/s)\n",
+        "throw_time = 0.0\n",
+        "mujoco.mj_resetData(model, data)\n",
+        "while data.time \u003c duration and not throw_time:\n",
+        "  mujoco.mj_step(model, data)\n",
+        "  box_speed = np.linalg.norm(data.joint('box').qvel[:3])\n",
+        "  if box_speed \u003e 0.02:\n",
+        "    throw_time = data.time\n",
+        "assert throw_time \u003e 0\n",
+        "\n",
+        "def mix(time, t0=0.0, width=1.0):\n",
+        "  \"\"\"Sigmoidal mixing function.\"\"\"\n",
+        "  t = (time - t0) / width\n",
+        "  s = 1 / (1 + np.exp(-t))\n",
+        "  return 1 - s, s\n",
+        "\n",
+        "def unit_cos(t):\n",
+        "  \"\"\"Unit cosine sigmoid from (0,0) to (1,1).\"\"\"\n",
+        "  return 0.5 - np.cos(np.pi*np.clip(t, 0, 1))/2\n",
+        "\n",
+        "def orbit_motion(t):\n",
+        "  \"\"\"Return orbit trajectory.\"\"\"\n",
+        "  distance = 0.9\n",
+        "  azimuth = 140 + 100 * unit_cos(t)\n",
+        "  elevation = -30\n",
+        "  lookat = data.geom('floor').xpos.copy()\n",
+        "  return distance, azimuth, elevation, lookat\n",
+        "\n",
+        "def track_motion():\n",
+        "  \"\"\"Return box-track trajectory.\"\"\"\n",
+        "  distance = 0.08\n",
+        "  azimuth = 280\n",
+        "  elevation = -10\n",
+        "  lookat = data.geom('box').xpos.copy()\n",
+        "  return distance, azimuth, elevation, lookat\n",
+        "\n",
+        "def cam_motion():\n",
+        "  \"\"\"Return sigmoidally-mixed {orbit, box-track} trajectory.\"\"\"\n",
+        "  d0, a0, e0, l0 = orbit_motion(data.time / throw_time)\n",
+        "  d1, a1, e1, l1 = track_motion()\n",
+        "  mix_time = 0.3\n",
+        "  w0, w1 = mix(data.time, throw_time, mix_time)\n",
+        "  return w0*d0+w1*d1, w0*a0+w1*a1, w0*e0+w1*e1, w0*l0+w1*l1\n",
+        "\n",
+        "# Make a camera.\n",
+        "cam = mujoco.MjvCamera()\n",
+        "mujoco.mjv_defaultCamera(cam)\n",
+        "\n",
+        "# Simulate and display video.\n",
+        "framerate = 60  # (Hz)\n",
+        "slowdown = 4    # 4x slow-down\n",
+        "mujoco.mj_resetData(model, data)\n",
+        "frames = []\n",
+        "while data.time \u003c duration:\n",
+        "  mujoco.mj_step(model, data)\n",
+        "  if len(frames) \u003c data.time * framerate * slowdown:\n",
+        "    cam.distance, cam.azimuth, cam.elevation, cam.lookat = cam_motion()\n",
+        "    renderer.update_scene(data, cam)\n",
+        "    pixels = renderer.render()\n",
+        "    frames.append(pixels)\n",
+        "media.show_video(frames, fps=framerate)"
+      ]
     }
   ],
   "metadata": {