[wip] Add test notebook for contact models

examples/Parallel_computing.ipynb

@@ -22,6 +22,10 @@
    "source": [
     "# @title Imports and setup\n",
     "import sys\n",
+    "import os\n",
+    "\n",
+    "# Deactivate GPU to avoid out of memory errors\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
     "\n",
     "from IPython.display import HTML, clear_output, display\n",
     "\n",
@@ -37,11 +41,13 @@
     "import jax.numpy as jnp\n",
     "import jax_dataclasses\n",
     "import rod\n",
-    "from rod.builder.primitives import SphereBuilder\n",
+    "from rod.builder.primitives import SphereBuilder, BoxBuilder\n",
     "\n",
     "import jaxsim.typing as jtp\n",
     "from jaxsim import logging\n",
     "\n",
+    "from jaxsim.mujoco import MujocoVideoRecorder, MujocoModelHelper, RodModelToMjcf\n",
+    "\n",
     "logging.set_logging_level(logging.LoggingLevel.INFO)\n",
     "logging.info(f\"Running on {jax.devices()}\")"
    ]
@@ -64,7 +70,7 @@
     "# @title Create a sphere model\n",
     "model_sdf_string = rod.Sdf(\n",
     "    version=\"1.7\",\n",
-    "    model=SphereBuilder(radius=0.10, mass=1.0, name=\"sphere\")\n",
+    "    model=BoxBuilder(x=0.30, y=0.30, z=0.30, mass=1.0, name=\"box\")\n",
     "    .build_model()\n",
     "    .add_link()\n",
     "    .add_inertial()\n",
@@ -112,6 +118,21 @@
     "integrator_state = integrator.init(x0=data.state, t0=0.0, dt=dt)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mcjf_string, assets = RodModelToMjcf.convert(rod_model=model_sdf_string.model)\n",
+    "mj_helper = MujocoModelHelper.build_from_xml(\n",
+    "    mjcf_description=mcjf_string, assets=assets\n",
+    ")\n",
+    "recorder = MujocoVideoRecorder(\n",
+    "    model=mj_helper.model, assets=mj_helper.data, fps=int(1 / dt), width=640, height=480\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -129,11 +150,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data = data.replace(\n",
-    "    soft_contacts_params=js.contact.estimate_good_soft_contacts_parameters(\n",
-    "        model, number_of_active_collidable_points_steady_state=3\n",
-    "    )\n",
-    ")"
+    "# data = data.replace(\n",
+    "#     soft_contacts_params=js.contact.estimate_good_soft_contacts_parameters(\n",
+    "#         model, number_of_active_collidable_points_steady_state=3\n",
+    "#     )\n",
+    "# )"
    ]
   },
   {
@@ -150,8 +171,8 @@
    "outputs": [],
    "source": [
     "# Primary Calculations\n",
-    "envs_per_row = 4  # @slider(2, 10, 1)\n",
-    "\n",
+    "envs_per_row = 1  # @slider(2, 10, 1)\n",
+    "initial_height = 0.7\n",
     "env_spacing = 0.5\n",
     "edge_len = env_spacing * (2 * envs_per_row - 1)\n",
     "\n",
@@ -162,7 +183,7 @@
     "    xx, yy = jnp.meshgrid(edge, edge)\n",
     "\n",
     "    poses = [\n",
-    "        [[xx[i, j], yy[i, j], 0.2 + 0.1 * (i * envs_per_row + j)], [0, 0, 0]]\n",
+    "        [[xx[i, j], yy[i, j], initial_height], [0, 0, 0]]\n",
     "        for i in range(xx.shape[0])\n",
     "        for j in range(yy.shape[0])\n",
     "    ]\n",
@@ -229,12 +250,13 @@
    "outputs": [],
    "source": [
     "# Define a function to simulate multiple model instances\n",
-    "simulate_vectorized = jax.vmap(simulate, in_axes=(None, None, 0))\n",
+    "# simulate_vectorized = jax.vmap(simulate, in_axes=(None, None, 0))\n",
     "\n",
     "# Run and time the simulation\n",
     "now = time.perf_counter()\n",
     "\n",
-    "x_t = simulate_vectorized(data, integrator_state, poses[:, 0])\n",
+    "# x_t = simulate_vectorized(data, integrator_state, poses[:, 0]).\n",
+    "x_t = simulate(data, integrator_state, poses)\n",
     "\n",
     "comp_time = time.perf_counter() - now\n",
     "\n",
@@ -246,6 +268,22 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "for pose in x_t:\n",
+    "    mj_helper.set_base_position(pose)\n",
+    "    recorder.record_frame()\n",
+    "\n",
+    "\n",
+    "recorder.write_video(path=Path.cwd() / Path(\"sphere.mp4\"), exists_ok=True)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -300,7 +338,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,