From fd19ecdc11b9dc2a6c4a2da49f13d2d66faf05ed Mon Sep 17 00:00:00 2001
From: Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
Date: Tue, 28 Jan 2025 17:48:14 +0100
Subject: [PATCH] [DROP] Add examples for parallel visualization

---
 .../jaxsim_as_physics_engine_advanced.ipynb   | 78 ++++++++++++++++++-
 1 file changed, 76 insertions(+), 2 deletions(-)

diff --git a/examples/jaxsim_as_physics_engine_advanced.ipynb b/examples/jaxsim_as_physics_engine_advanced.ipynb
index 1d76c9817..de85498c3 100644
--- a/examples/jaxsim_as_physics_engine_advanced.ipynb
+++ b/examples/jaxsim_as_physics_engine_advanced.ipynb
@@ -299,6 +299,17 @@
     "    )\n",
     ")(jnp.vstack(subkeys))\n",
     "\n",
+    "# Reset the x and y position to a grid.\n",
+    "data_batch_t0 = data_batch_t0.reset_base_position(\n",
+    "    jnp.array(\n",
+    "        [\n",
+    "            jnp.linspace(-1, 1, batch_size),\n",
+    "            jnp.linspace(-1, 1, batch_size),\n",
+    "            data_batch_t0.base_position()[:, 2],\n",
+    "        ]\n",
+    "    ).T\n",
+    ")\n",
+    "\n",
     "print(\"W_p_B(t0)=\\n\", data_batch_t0.base_position()[0:10])"
    ]
   },
@@ -420,6 +431,69 @@
     "plt.show()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jaxsim.mujoco\n",
+    "\n",
+    "mjcf_string, assets = jaxsim.mujoco.ModelToMjcf.convert(model.built_from)\n",
+    "\n",
+    "# Create a helper for each parallel instance.\n",
+    "mj_model_helpers = [\n",
+    "    jaxsim.mujoco.MujocoModelHelper.build_from_xml(\n",
+    "        mjcf_description=mjcf_string, assets=assets\n",
+    "    )\n",
+    "    for _ in range(batch_size)\n",
+    "]\n",
+    "\n",
+    "# Create the video recorder.\n",
+    "recorder = jaxsim.mujoco.MujocoVideoRecorder(\n",
+    "    model=mj_model_helpers[0].model,\n",
+    "    data=[helper.data for helper in mj_model_helpers],\n",
+    "    fps=int(1 / model.time_step),\n",
+    "    width=320 * 2,\n",
+    "    height=240 * 2,\n",
+    ")\n",
+    "\n",
+    "for data_t in data_trajectory_list:\n",
+    "\n",
+    "    for helper, base_position, base_quaternion, joint_position in zip(\n",
+    "        mj_model_helpers,\n",
+    "        data_t.base_position(),\n",
+    "        data.base_orientation(True),\n",
+    "        data.joint_positions(),\n",
+    "        strict=True,\n",
+    "    ):\n",
+    "        helper.set_base_position(position=base_position)\n",
+    "        helper.set_base_orientation(orientation=base_quaternion, dcm=True)\n",
+    "\n",
+    "        if model.dofs() > 0:\n",
+    "            helper.set_joint_positions(\n",
+    "                positions=joint_position, joint_names=model.joint_names()\n",
+    "            )\n",
+    "\n",
+    "    # Record a new video frame.\n",
+    "    recorder.record_frame()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import mediapy as media\n",
+    "\n",
+    "# from pathlib import Path\n",
+    "\n",
+    "# recorder.write_video(Path(\"sphere.mp4\"))\n",
+    "\n",
+    "media.show_video(recorder.frames, fps=recorder.fps)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -454,7 +528,7 @@
    "toc_visible": true
   },
   "kernelspec": {
-   "display_name": "rsl",
+   "display_name": "jaxsimsprint",
    "language": "python",
    "name": "python3"
   },
@@ -468,7 +542,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,