jaxsim 0.5.1.dev103__tar.gz → 0.5.1.dev126__tar.gz

{jaxsim-0.5.1.dev103 → jaxsim-0.5.1.dev126}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev103
+Version: 0.5.1.dev126
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>
@@ -81,6 +81,7 @@ Requires-Dist: pytest>=6.0; extra == "testing"
 Requires-Dist: pytest-benchmark; extra == "testing"
 Requires-Dist: pytest-icdiff; extra == "testing"
 Requires-Dist: robot-descriptions; extra == "testing"
+Requires-Dist: icub-models; extra == "testing"
 Provides-Extra: viz
 Requires-Dist: lxml; extra == "viz"
 Requires-Dist: mediapy; extra == "viz"
@@ -91,9 +92,7 @@ Requires-Dist: jaxsim[style,testing,viz]; extra == "all"
 
 # JaxSim
 
-JaxSim is a **differentiable physics engine** and **multibody dynamics library** designed for applications in control and robot learning, implemented with JAX.
-
-Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence.
+**JaxSim** is a **differentiable physics engine** and **multibody dynamics library** built with JAX, tailored for control and robotic learning applications.
 
 <div align="center">
 <br/>
@@ -108,41 +107,105 @@ Its design facilitates research and accelerates prototyping in the intersection
 </div>
 
 ## Features
+- Reduced-coordinate physics engine for **fixed-base** and **floating-base** robots.
+- Multibody dynamics library for model-based control algorithms.
+- Fully Python-based, leveraging [jax][jax] following a functional programming paradigm.
+- Seamless execution on CPUs, GPUs, and TPUs.
+- Supports JIT compilation and automatic vectorization for high performance.
+- Compatible with SDF models and URDF (via [sdformat][sdformat] conversion).
+
+## Usage
+
+### Using JaxSim as simulator
+
+
+```python
+import jax.numpy as jnp
+import jaxsim.api as js
+import icub_models
+import pathlib
+
+# Load the iCub model
+model_path = icub_models.get_model_file("iCubGazeboV2_5")
+joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
+          'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
+          'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
+          'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
+          'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
+          'r_ankle_roll')
+
+# Build and reduce the model
+model_description = pathlib.Path(model_path)
+full_model = js.model.JaxSimModel.build_from_model_description(
+    model_description=model_description, time_step=0.0001, is_urdf=True
+)
+model = js.model.reduce(model=full_model, considered_joints=joints)
+
+ndof = model.dofs()
+# Initialize data and simulation
+data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
+    base_position=jnp.array([0.0, 0.0, 1.0])
+)
+T = jnp.arange(start=0, stop=1.0, step=model.time_step)
+tau = jnp.zeros(ndof)
+
+# Simulate
+for t in T:
+    data, _ = js.model.step(model=model, data=data, link_forces=None, joint_force_references=tau)
 
-- Physics engine in reduced coordinates supporting fixed-base and floating-base robots.
-- Multibody dynamics library providing all the necessary components for developing model-based control algorithms.
-- Completely developed in Python with [`google/jax`][jax] following a functional programming paradigm.
-- Transparent support for running on CPUs, GPUs, and TPUs.
-- Full support for JIT compilation for increased performance.
-- Full support for automatic vectorization for massive parallelization of open-loop and closed-loop architectures.
-- Support for SDF models and, upon conversion with [sdformat][sdformat], URDF models.
-- Visualization based on the [passive viewer][passive_viewer_mujoco] of Mujoco.
-
-### JaxSim as a simulator
-
-- Wide range of fixed-step explicit Runge-Kutta integrators.
-- Support for variable-step integrators implemented as embedded Runge-Kutta schemes.
-- Improved stability by optionally integrating the base orientation on the $\text{SO}(3)$ manifold.
-- Soft contacts model supporting full friction cone and sticking-slipping transition.
-- Collision detection between points rigidly attached to bodies and uneven ground surfaces.
-
-### JaxSim as a multibody dynamics library
+```
 
-- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians.
-- Provides all the quantities included in the Euler-Poincarè formulation of the equations of motion.
-- Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].
-- Exposes all the necessary quantities to develop controllers in centroidal coordinates.
-- Supports running open-loop and full closed-loop control architectures on hardware accelerators.
+### Using JaxSim as a multibody dynamics library
+``` python
+import jax.numpy as jnp
+import jaxsim.api as js
+import icub_models
+import pathlib
+
+# Load the iCub model
+model_path = icub_models.get_model_file("iCubGazeboV2_5")
+joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
+          'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
+          'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
+          'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
+          'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
+          'r_ankle_roll')
+
+# Build and reduce the model
+model_description = pathlib.Path(model_path)
+full_model = js.model.JaxSimModel.build_from_model_description(
+    model_description=model_description, time_step=0.0001, is_urdf=True
+)
+model = js.model.reduce(model=full_model, considered_joints=joints)
+
+# Initialize model data
+data = js.data.JaxSimModelData.build(
+    model=model,
+    base_position=jnp.array([0.0, 0.0, 1.0],
+)
+
+# Frame and dynamics computations
+frame_index = js.frame.name_to_idx(model=model, frame_name="l_foot")
+W_H_F = js.frame.transform(model=model, data=data, frame_index=frame_index)  # Frame transformation
+W_J_F = js.frame.jacobian(model=model, data=data, frame_index=frame_index)  # Frame Jacobian
+
+# Dynamics properties
+M = js.model.free_floating_mass_matrix(model=model, data=data)  # Mass matrix
+h = js.model.free_floating_bias_forces(model=model, data=data)  # Bias forces
+g = js.model.free_floating_gravity_forces(model=model, data=data)  # Gravity forces
+C = js.model.free_floating_coriolis_matrix(model=model, data=data)  # Coriolis matrix
+
+# Print dynamics results
+print(f"M: shape={M.shape}, h: shape={h.shape}, g: shape={g.shape}, C: shape={C.shape}")
 
-### JaxSim for robot learning
+```
+### Additional features
 
-- Being developed with JAX, all the RBDAs support automatic differentiation both in forward and reverse modes.
+- Full support for automatic differentiation of RBDAs (forward and reverse modes) with JAX.
 - Support for automatically differentiating against kinematics and dynamics parameters.
 - All fixed-step integrators are forward and reverse differentiable.
 - All variable-step integrators are forward differentiable.
-- Ideal for sampling synthetic data for reinforcement learning (RL).
-- Ideal for designing physics-informed neural networks (PINNs) with loss functions requiring model-based quantities.
-- Ideal for combining model-based control with learning-based components.
+- Check the example folder for additional usecase !
 
 [jax]: https://github.com/google/jax/
 [sdformat]: https://github.com/gazebosim/sdformat
@@ -156,12 +219,6 @@ Its design facilitates research and accelerates prototyping in the intersection
 > JaxSim currently focuses on locomotion applications.
 > Only contacts between bodies and smooth ground surfaces are supported.
 
-## Documentation
-
-The JaxSim API documentation is available at [jaxsim.readthedocs.io][readthedocs].
-
-[readthedocs]: https://jaxsim.readthedocs.io/
-
 ## Installation
 
 <details>
@@ -233,6 +290,13 @@ pip install --no-deps -e .
 [venv]: https://docs.python.org/3/tutorial/venv.html
 [jax_gpu]: https://github.com/google/jax/#installation
 
+## Documentation
+
+The JaxSim API documentation is available at [jaxsim.readthedocs.io][readthedocs].
+
+[readthedocs]: https://jaxsim.readthedocs.io/
+
+
 ## Overview
 
 <details>

{jaxsim-0.5.1.dev103 → jaxsim-0.5.1.dev126}/README.md

@@ -1,8 +1,6 @@
 # JaxSim
 
-JaxSim is a **differentiable physics engine** and **multibody dynamics library** designed for applications in control and robot learning, implemented with JAX.
-
-Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence.
+**JaxSim** is a **differentiable physics engine** and **multibody dynamics library** built with JAX, tailored for control and robotic learning applications.
 
 <div align="center">
 <br/>
@@ -17,41 +15,105 @@ Its design facilitates research and accelerates prototyping in the intersection
 </div>
 
 ## Features
+- Reduced-coordinate physics engine for **fixed-base** and **floating-base** robots.
+- Multibody dynamics library for model-based control algorithms.
+- Fully Python-based, leveraging [jax][jax] following a functional programming paradigm.
+- Seamless execution on CPUs, GPUs, and TPUs.
+- Supports JIT compilation and automatic vectorization for high performance.
+- Compatible with SDF models and URDF (via [sdformat][sdformat] conversion).
+
+## Usage
+
+### Using JaxSim as simulator
+
+
+```python
+import jax.numpy as jnp
+import jaxsim.api as js
+import icub_models
+import pathlib
+
+# Load the iCub model
+model_path = icub_models.get_model_file("iCubGazeboV2_5")
+joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
+          'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
+          'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
+          'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
+          'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
+          'r_ankle_roll')
+
+# Build and reduce the model
+model_description = pathlib.Path(model_path)
+full_model = js.model.JaxSimModel.build_from_model_description(
+    model_description=model_description, time_step=0.0001, is_urdf=True
+)
+model = js.model.reduce(model=full_model, considered_joints=joints)
+
+ndof = model.dofs()
+# Initialize data and simulation
+data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
+    base_position=jnp.array([0.0, 0.0, 1.0])
+)
+T = jnp.arange(start=0, stop=1.0, step=model.time_step)
+tau = jnp.zeros(ndof)
+
+# Simulate
+for t in T:
+    data, _ = js.model.step(model=model, data=data, link_forces=None, joint_force_references=tau)
 
-- Physics engine in reduced coordinates supporting fixed-base and floating-base robots.
-- Multibody dynamics library providing all the necessary components for developing model-based control algorithms.
-- Completely developed in Python with [`google/jax`][jax] following a functional programming paradigm.
-- Transparent support for running on CPUs, GPUs, and TPUs.
-- Full support for JIT compilation for increased performance.
-- Full support for automatic vectorization for massive parallelization of open-loop and closed-loop architectures.
-- Support for SDF models and, upon conversion with [sdformat][sdformat], URDF models.
-- Visualization based on the [passive viewer][passive_viewer_mujoco] of Mujoco.
-
-### JaxSim as a simulator
-
-- Wide range of fixed-step explicit Runge-Kutta integrators.
-- Support for variable-step integrators implemented as embedded Runge-Kutta schemes.
-- Improved stability by optionally integrating the base orientation on the $\text{SO}(3)$ manifold.
-- Soft contacts model supporting full friction cone and sticking-slipping transition.
-- Collision detection between points rigidly attached to bodies and uneven ground surfaces.
-
-### JaxSim as a multibody dynamics library
+```
 
-- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians.
-- Provides all the quantities included in the Euler-Poincarè formulation of the equations of motion.
-- Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].
-- Exposes all the necessary quantities to develop controllers in centroidal coordinates.
-- Supports running open-loop and full closed-loop control architectures on hardware accelerators.
+### Using JaxSim as a multibody dynamics library
+``` python
+import jax.numpy as jnp
+import jaxsim.api as js
+import icub_models
+import pathlib
+
+# Load the iCub model
+model_path = icub_models.get_model_file("iCubGazeboV2_5")
+joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
+          'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
+          'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
+          'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
+          'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
+          'r_ankle_roll')
+
+# Build and reduce the model
+model_description = pathlib.Path(model_path)
+full_model = js.model.JaxSimModel.build_from_model_description(
+    model_description=model_description, time_step=0.0001, is_urdf=True
+)
+model = js.model.reduce(model=full_model, considered_joints=joints)
+
+# Initialize model data
+data = js.data.JaxSimModelData.build(
+    model=model,
+    base_position=jnp.array([0.0, 0.0, 1.0],
+)
+
+# Frame and dynamics computations
+frame_index = js.frame.name_to_idx(model=model, frame_name="l_foot")
+W_H_F = js.frame.transform(model=model, data=data, frame_index=frame_index)  # Frame transformation
+W_J_F = js.frame.jacobian(model=model, data=data, frame_index=frame_index)  # Frame Jacobian
+
+# Dynamics properties
+M = js.model.free_floating_mass_matrix(model=model, data=data)  # Mass matrix
+h = js.model.free_floating_bias_forces(model=model, data=data)  # Bias forces
+g = js.model.free_floating_gravity_forces(model=model, data=data)  # Gravity forces
+C = js.model.free_floating_coriolis_matrix(model=model, data=data)  # Coriolis matrix
+
+# Print dynamics results
+print(f"M: shape={M.shape}, h: shape={h.shape}, g: shape={g.shape}, C: shape={C.shape}")
 
-### JaxSim for robot learning
+```
+### Additional features
 
-- Being developed with JAX, all the RBDAs support automatic differentiation both in forward and reverse modes.
+- Full support for automatic differentiation of RBDAs (forward and reverse modes) with JAX.
 - Support for automatically differentiating against kinematics and dynamics parameters.
 - All fixed-step integrators are forward and reverse differentiable.
 - All variable-step integrators are forward differentiable.
-- Ideal for sampling synthetic data for reinforcement learning (RL).
-- Ideal for designing physics-informed neural networks (PINNs) with loss functions requiring model-based quantities.
-- Ideal for combining model-based control with learning-based components.
+- Check the example folder for additional usecase !
 
 [jax]: https://github.com/google/jax/
 [sdformat]: https://github.com/gazebosim/sdformat
@@ -65,12 +127,6 @@ Its design facilitates research and accelerates prototyping in the intersection
 > JaxSim currently focuses on locomotion applications.
 > Only contacts between bodies and smooth ground surfaces are supported.
 
-## Documentation
-
-The JaxSim API documentation is available at [jaxsim.readthedocs.io][readthedocs].
-
-[readthedocs]: https://jaxsim.readthedocs.io/
-
 ## Installation
 
 <details>
@@ -142,6 +198,13 @@ pip install --no-deps -e .
 [venv]: https://docs.python.org/3/tutorial/venv.html
 [jax_gpu]: https://github.com/google/jax/#installation
 
+## Documentation
+
+The JaxSim API documentation is available at [jaxsim.readthedocs.io][readthedocs].
+
+[readthedocs]: https://jaxsim.readthedocs.io/
+
+
 ## Overview
 
 <details>

{jaxsim-0.5.1.dev103 → jaxsim-0.5.1.dev126}/environment.yml

@@ -29,6 +29,7 @@ dependencies:
   - pytest
   - pytest-icdiff
   - robot_descriptions
+  - icub-models
   # [viz]
   - lxml
   - mediapy
@@ -55,6 +56,7 @@ dependencies:
   - sphinx-multiversion
   - sphinx_rtd_theme
   - sphinx-toolbox
+  - icub-models
   # ========================================
   # Other dependencies for GitHub Codespaces
   # ========================================

{jaxsim-0.5.1.dev103 → jaxsim-0.5.1.dev126}/examples/README.md

@@ -8,11 +8,13 @@ This folder contains Jupyter notebooks that demonstrate the practical usage of J
 | :--- | :---: | :--- |
 | [`jaxsim_as_multibody_dynamics_library`](./jaxsim_as_multibody_dynamics_library.ipynb) | [![Open In Colab][colab_badge]][ipynb_jaxsim_as_multibody_dynamics] | An example demonstrating how to use JaxSim as a multibody dynamics library. |
 | [`jaxsim_as_physics_engine.ipynb`](./jaxsim_as_physics_engine.ipynb) | [![Open In Colab][colab_badge]][ipynb_jaxsim_as_physics_engine] | An example demonstrating how to simulate vectorized models in parallel. |
+| [`jaxsim_as_physics_engine_advanced.ipynb`](./jaxsim_as_physics_engine_advanced.ipynb) | [![Open In Colab][colab_badge]][jaxsim_as_physics_engine_advanced] | An example showcasing advanced JaxSim usage, such as customizing the integrator, contact model, and more. |
 | [`jaxsim_for_robot_controllers.ipynb`](./jaxsim_for_robot_controllers.ipynb) | [![Open In Colab][colab_badge]][ipynb_jaxsim_closed_loop] | A basic example showing how to simulate a PD controller with gravity compensation for a 2-DOF cart-pole. |
 
 [colab_badge]: https://colab.research.google.com/assets/colab-badge.svg
 [ipynb_jaxsim_closed_loop]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_for_robot_controllers.ipynb
 [ipynb_jaxsim_as_physics_engine]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_physics_engine.ipynb
+[jaxsim_as_physics_engine_advanced]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_physics_engine_advanced.ipynb
 [ipynb_jaxsim_as_multibody_dynamics]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_multibody_dynamics_library.ipynb
 
 ## How to run the examples

jaxsim-0.5.1.dev126/examples/jaxsim_as_physics_engine.ipynb

@@ -0,0 +1,274 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "H-WgcgGQaTG7"
+   },
+   "source": [
+    "# JaxSim as a hardware-accelerated parallel physics engine\n",
+    "\n",
+    "This notebook shows how to use the key APIs to load a robot model and simulate multiple trajectories simultaneously.\n",
+    "\n",
+    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_physics_engine.ipynb\">\n",
+    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+    "</a>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "SgOSnrSscEkt"
+   },
+   "source": [
+    "## Prepare the environment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "fdqvAqMDaTG9"
+   },
+   "outputs": [],
+   "source": [
+    "# @title Imports and setup\n",
+    "import sys\n",
+    "from IPython.display import clear_output\n",
+    "\n",
+    "IS_COLAB = \"google.colab\" in sys.modules\n",
+    "\n",
+    "# Install JAX and Gazebo\n",
+    "if IS_COLAB:\n",
+    "    !{sys.executable} -m pip install --pre -qU jaxsim\n",
+    "    !apt install -qq lsb-release wget gnupg\n",
+    "    !wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg\n",
+    "    !echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null\n",
+    "    !apt -qq update\n",
+    "    !apt install -qq --no-install-recommends libsdformat13 gz-tools2\n",
+    "\n",
+    "    clear_output()\n",
+    "\n",
+    "# Set environment variable to avoid GPU out of memory errors\n",
+    "%env XLA_PYTHON_CLIENT_MEM_PREALLOCATE=false\n",
+    "\n",
+    "# ================\n",
+    "# Notebook imports\n",
+    "# ================\n",
+    "\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import jaxsim.api as js\n",
+    "from jaxsim import logging\n",
+    "import pathlib\n",
+    "import urllib.request\n",
+    "\n",
+    "logging.set_logging_level(logging.LoggingLevel.WARNING)\n",
+    "print(f\"Running on {jax.devices()}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "NqjuZKvOaTG_"
+   },
+   "source": [
+    "## Prepare the simulation\n",
+    "\n",
+    "JaxSim supports loading robot descriptions from both [SDF][sdformat] and [URDF][urdf] files. In this example, we will load the [ergoCub][ergocub] model urdf.\n",
+    "\n",
+    "[sdformat]: http://sdformat.org/\n",
+    "[urdf]: http://wiki.ros.org/urdf/\n",
+    "[ergocub]: https://ergocub.eu/\n",
+    "[rod]: https://github.com/ami-iit/rod\n",
+    "\n",
+    "### Create the model and its data\n",
+    " To define a simulation we need two main objects:\n",
+    "\n",
+    "- `model`: an object that defines the dynamics of the system.\n",
+    "- `data`: an object that contains the state of the system.\n",
+    "\n",
+    "\n",
+    "The `JaxSimModel` object contains the simulation time step, the integrator and the contact model.\n",
+    "To see the advanced usage, check the advanced example, where you will see how to pass explicitly an integrator class and state to the `model` object and how to change the contact model."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create the model "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "etQ577cFaTHA"
+   },
+   "outputs": [],
+   "source": [
+    "#  Create the JaxSim model.\n",
+    "url = \"https://raw.githubusercontent.com/icub-tech-iit/ergocub-software/refs/heads/master/urdf/ergoCub/robots/ergoCubSN001/model.urdf\"\n",
+    "\n",
+    "# Retrieve the file\n",
+    "model_path, _ = urllib.request.urlretrieve(url)\n",
+    "\n",
+    "model_description_path = pathlib.Path(model_path)\n",
+    "full_model = js.model.JaxSimModel.build_from_model_description(\n",
+    "    model_description=model_description_path,\n",
+    "    time_step=0.0001,\n",
+    "    is_urdf=True\n",
+    ")\n",
+    "\n",
+    "joints_list = tuple(('l_shoulder_pitch', 'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow',\n",
+    "               'r_shoulder_pitch', 'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow',\n",
+    "               'l_hip_pitch', 'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',\n",
+    "               'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll'))\n",
+    "\n",
+    "model = js.model.reduce(\n",
+    "    model=full_model,\n",
+    "    considered_joints=joints_list\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create the data object \n",
+    "\n",
+    "The data object is never changed by reference. Anytime you call a method aimed at modifying data, like `reset_base_position`, a new data object will be returned with the updated attributes while the original data will not be changed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create the data of a single model.\n",
+    "data_zero = js.data.JaxSimModelData.zero(model=model)\n",
+    "base_position = jnp.array([0.0, 0.0, 1.0])\n",
+    "data = data_zero.reset_base_position(base_position=base_position)  # Note that the reset position returns the updated data object"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Simulation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a random JAX key.\n",
+    "\n",
+    "key = jax.random.PRNGKey(seed=0)\n",
+    "\n",
+    "# Initialize the simulated time.\n",
+    "T = jnp.arange(start=0, stop=0.3, step=model.time_step)\n",
+    "\n",
+    "# Simulate\n",
+    "for _t in T:\n",
+    "    data, _ = js.model.step(\n",
+    "        model=model,\n",
+    "        data=data,\n",
+    "        link_forces=None,\n",
+    "        joint_force_references=None,\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Vectorized simulation \n",
+    "\n",
+    "We will now vectorize the simulation on batched data using `jax.vmap`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# first we have to vmap the function\n",
+    "\n",
+    "import functools\n",
+    "from typing import Any\n",
+    "\n",
+    "\n",
+    "@jax.jit\n",
+    "def step_single(\n",
+    "    model: js.model.JaxSimModel,\n",
+    "    data: js.data.JaxSimModelData,\n",
+    ") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
+    "\n",
+    "    # Close step over static arguments.\n",
+    "    return js.model.step(\n",
+    "        model=model,\n",
+    "        data=data,\n",
+    "        link_forces=None,\n",
+    "        joint_force_references=None,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "@jax.jit\n",
+    "@functools.partial(jax.vmap, in_axes=(None, 0))\n",
+    "def step_parallel(\n",
+    "    model: js.model.JaxSimModel,\n",
+    "    data: js.data.JaxSimModelData,\n",
+    ") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
+    "\n",
+    "    return step_single(\n",
+    "        model=model, data=data\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "# Then we have to create the vector of initial state\n",
+    "batch_size = 5\n",
+    "data_batch_t0 = jax.vmap(\n",
+    "    lambda pos: data_zero.reset_base_position(base_position=pos)\n",
+    ")(jnp.tile(jnp.array([0.0, 0.0, 1.0]), (batch_size, 1)))\n",
+    "\n",
+    "data = data_batch_t0\n",
+    "for _t in T:\n",
+    "    data, _ = step_parallel(model, data)"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuClass": "premium",
+   "private_outputs": true,
+   "provenance": [],
+   "toc_visible": true
+  },
+  "kernelspec": {
+   "display_name": "jaxsim",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}