diff --git a/examples/Rigid_body_dynamics.ipynb b/examples/Rigid_body_dynamics.ipynb
index bb14b1ba5..4cf0d76bf 100644
--- a/examples/Rigid_body_dynamics.ipynb
+++ b/examples/Rigid_body_dynamics.ipynb
@@ -1,578 +1,1138 @@
{
- "cells": [
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "# `JAXsim` Showcase: Rigid Body Dynamics\n",
- "\n",
- "First, we install the necessary packages and import them.\n",
- "\n",
- "\n",
- " ![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
\n",
- ""
- ]
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "DpLq0-lltwZ1"
+ },
+ "source": [
+ "# `JaxSim` as a Multibody Dynamics Library\n",
+ "\n",
+ "\n",
+ "JaxSim was initially developed as a **hardware-accelerated physics engine**. Over time, it has evolved, adding new features to become a comprehensive **JAX-based multibody dynamics library**.\n",
+ "\n",
+ "In this notebook, you'll explore the main APIs for loading robot models and computing key quantities for applications such as control, planning, and more.\n",
+ "\n",
+ "A key advantage of JaxSim is its ability to create fully differentiable closed-loop systems, enabling end-to-end optimization. Combined with the flexibility to parameterize model kinematics and dynamics, JaxSim can serve as an excellent playground for robot learning applications.\n",
+ "\n",
+ "\n",
+ " \n",
+ "\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "rcEwprINtwZ3"
+ },
+ "source": [
+ "## Prepare environment\n",
+ "\n",
+ "First, we need to install the necessary packages and import their resources."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "cellView": "form",
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "u4xL7dbBtwZ3",
+ "outputId": "b30c4790-f0f4-4ce6-f439-db662176640e"
+ },
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Running on [CpuDevice(id=0)]\n"
+ ]
+ }
+ ],
+ "source": [
+ "# @title Imports and setup\n",
+ "from IPython.display import clear_output\n",
+ "import os\n",
+ "import pathlib\n",
+ "import sys\n",
+ "\n",
+ "IS_COLAB = \"google.colab\" in sys.modules\n",
+ "\n",
+ "# Install JAX, sdformat, and other notebook dependencies.\n",
+ "if IS_COLAB:\n",
+ " !{sys.executable} -m pip install -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",
+ " !pip install robot_descriptions\n",
+ "\n",
+ " clear_output()\n",
+ "\n",
+ "\n",
+ "import jax\n",
+ "import jax.numpy as jnp\n",
+ "from jaxsim import logging\n",
+ "import robot_descriptions\n",
+ "\n",
+ "import jaxsim.api as js\n",
+ "import jaxsim.math\n",
+ "from jaxsim import VelRepr\n",
+ "\n",
+ "logging.set_logging_level(logging.LoggingLevel.INFO)\n",
+ "logging.info(f\"Running on {jax.devices()}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "fN8Xg4QgtwZ4"
+ },
+ "source": [
+ "## Robot model\n",
+ "\n",
+ "JaxSim allows loading robot descriptions from both [SDF][sdformat] and [URDF][urdf] files.\n",
+ "\n",
+ "In this example, we will use the [ErgoCub][ergocub] humanoid robot model. If you have a URDF/SDF file for your robot that is compatible with [`gazebosim/sdformat`][sdformat_github][1], it should work out-of-the-box with JaxSim.\n",
+ "\n",
+ "[sdformat]: http://sdformat.org/\n",
+ "[urdf]: http://wiki.ros.org/urdf/\n",
+ "[ergocub]: https://ergocub.eu/\n",
+ "[sdformat_github]: https://github.com/gazebosim/sdformat\n",
+ "\n",
+ "---\n",
+ "\n",
+ "[1]: JaxSim validates robot descriptions using the command `gz sdf -p /path/to/file.urdf`. Ensure this command runs successfully on your file.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {
+ "cellView": "form",
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "rB0BFxyPtwZ5",
+ "outputId": "78d99087-7dd3-44dc-f9f4-875472946942"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Cloning https://github.com/icub-tech-iit/ergocub-software.git...\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "100%|██████████| 1900.0/1900.0 [00:06<00:00, 286.06it/s]\n"
+ ]
+ }
+ ],
+ "source": [
+ "# @title Fetch the URDF file\n",
+ "try:\n",
+ " os.environ[\"ROBOT_DESCRIPTION_COMMIT\"] = \"v0.7.1\"\n",
+ "\n",
+ " import robot_descriptions.ergocub_description\n",
+ "\n",
+ "finally:\n",
+ " _ = os.environ.pop(\"ROBOT_DESCRIPTION_COMMIT\", None)\n",
+ "\n",
+ "model_description_path = pathlib.Path(\n",
+ " robot_descriptions.ergocub_description.URDF_PATH.replace(\n",
+ " \"ergoCubSN000\", \"ergoCubSN001\"\n",
+ " )\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "jeTUZic8twZ5"
+ },
+ "source": [
+ "### Create the model and its data\n",
+ "\n",
+ "The dynamics of a generic floating-base model are governed by the following equations of motion:\n",
+ "\n",
+ "$$\n",
+ "M(\\mathbf{q}) \\dot{\\boldsymbol{\\nu}} + \\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = B \\boldsymbol{\\tau} + \\sum_{L \\in \\mathcal{L}} J_{W,L}^\\top(\\mathbf{q}) \\: \\mathbf{f}_L\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "Here, the system state is represented by:\n",
+ "\n",
+ "- $\\mathbf{q} = ({}^W \\mathbf{p}_B, \\, \\mathbf{s}) \\in \\text{SE}(3) \\times \\mathbb{R}^n$ is the generalized position.\n",
+ "- $\\boldsymbol{\\nu} = (\\boldsymbol{v}_{W,B}, \\, \\boldsymbol{\\omega}_{W,B}, \\, \\dot{\\mathbf{s}}) \\in \\mathbb{R}^{6+n}$ is the generalized velocity.\n",
+ "\n",
+ "The inputs to the system are:\n",
+ "\n",
+ "- $\\boldsymbol{\\tau} \\in \\mathbb{R}^n$ are the joint torques.\n",
+ "- $\\mathbf{f}_L \\in \\mathbb{R}^6$ is the 6D force applied to the frame associated to link $L$.\n",
+ "\n",
+ "JaxSim exposes functional APIs to operate over the following two main data structures:\n",
+ "\n",
+ "- **`JaxSimModel`** stores all the constant information parsed from the model description.\n",
+ "- **`JaxSimModelData`** holds the state of model.\n",
+ "\n",
+ "Additionally, JaxSim includes a utility class, **`JaxSimModelReferences`**, for managing and manipulating system inputs.\n",
+ "\n",
+ "---\n",
+ "\n",
+ "This notebook uses the notation summarized in the following report. Please refer to this document if you have any questions or if something is unclear.\n",
+ "\n",
+ "> Traversaro and Saccon, **Multibody dynamics notation**, 2019, [URL](https://pure.tue.nl/ws/portalfiles/portal/139293126/A_Multibody_Dynamics_Notation_Revision_2_.pdf)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "WYgBAxU0twZ6",
+ "outputId": "994b1cd5-d742-4735-a713-ed64a2ad1afa"
+ },
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_foot_rear->(r_foot_rear_ft_sensor)->r_ankle_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_foot_front->(r_foot_front_ft_sensor)->r_ankle_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_foot_rear->(l_foot_rear_ft_sensor)->l_ankle_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_foot_front->(l_foot_front_ft_sensor)->l_ankle_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_shoulder_3->(r_arm_ft_sensor)->r_shoulder_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_shoulder_3->(l_arm_ft_sensor)->l_shoulder_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hip_3->(r_leg_ft_sensor)->r_hip_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hip_3->(l_leg_ft_sensor)->l_hip_2\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Create the model from the model description.\n",
+ "# JaxSim removes all fixed joints by lumping together their parent and child links.\n",
+ "full_model = js.model.JaxSimModel.build_from_model_description(\n",
+ " model_description=model_description_path\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "DdaETmDStwZ6"
+ },
+ "source": [
+ "It is often useful to work with only a subset of joints, referred to as the _considered joints_. JaxSim allows to reduce a model so that the computation of the rigid body dynamics quantities is simplified.\n",
+ "\n",
+ "By default, the positions of the removed joints are considered to be zero. If this is not the case, the `reduce` function accepts a dictionary `dict[str, float]` to specify custom joint positions."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "QuhG7Zv5twZ7",
+ "outputId": "0a7fd714-218c-4a35-ce78-5d5a5c2bbbd0"
+ },
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_thumb_3->(r_thumb_dist)->r_hand_thumb_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_index_3->(r_index_dist)->r_hand_index_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_thumb_3->(l_thumb_dist)->l_hand_thumb_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_index_3->(l_index_dist)->l_hand_index_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_thumb_2->(r_thumb_prox)->r_hand_thumb_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_ring_2->(r_ring_dist)->r_hand_ring_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_pinkie_2->(r_pinkie_dist)->r_hand_pinkie_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_middle_2->(r_middle_dist)->r_hand_middle_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_index_2->(r_index_prox)->r_hand_index_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_thumb_2->(l_thumb_prox)->l_hand_thumb_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_ring_2->(l_ring_dist)->l_hand_ring_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_pinkie_2->(l_pinkie_dist)->l_hand_pinkie_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_middle_2->(l_middle_dist)->l_hand_middle_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_index_2->(l_index_prox)->l_hand_index_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_thumb_1->(r_thumb_add)->r_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_ring_1->(r_ring_prox)->r_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_pinkie_1->(r_pinkie_prox)->r_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_middle_1->(r_middle_prox)->r_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_index_1->(r_index_add)->r_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_thumb_1->(l_thumb_add)->l_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_ring_1->(l_ring_prox)->l_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_pinkie_1->(l_pinkie_prox)->l_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_middle_1->(l_middle_prox)->l_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_index_1->(l_index_add)->l_hand_palm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_hand_palm->(r_wrist_pitch)->r_wrist_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_hand_palm->(l_wrist_pitch)->l_wrist_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_wrist_2->(r_wrist_roll)->r_wrist_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_wrist_2->(l_wrist_roll)->l_wrist_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_wrist_1->(r_wrist_yaw)->r_forearm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_wrist_1->(l_wrist_yaw)->l_forearm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_forearm->(r_elbow)->r_upper_arm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: realsense->(camera_tilt)->head\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_forearm->(l_elbow)->l_upper_arm\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_upper_arm->(r_shoulder_yaw)->r_shoulder_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: head->(neck_yaw)->neck_3\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_upper_arm->(l_shoulder_yaw)->l_shoulder_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_shoulder_2->(r_shoulder_roll)->r_shoulder_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: neck_3->(neck_roll)->neck_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_shoulder_2->(l_shoulder_roll)->l_shoulder_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: r_shoulder_1->(r_shoulder_pitch)->chest\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: neck_2->(neck_pitch)->chest\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: l_shoulder_1->(l_shoulder_pitch)->chest\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: chest->(torso_yaw)->torso_2\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: torso_2->(torso_pitch)->torso_1\n",
+ "\u001b[34mjaxsim[401]\u001b[0m \u001b[1;30mINFO\u001b[0m Lumping chain: torso_1->(torso_roll)->root_link\n"
+ ]
+ }
+ ],
+ "source": [
+ "model = js.model.reduce(\n",
+ " model=full_model,\n",
+ " considered_joints=tuple(\n",
+ " j\n",
+ " for j in full_model.joint_names()\n",
+ " # Remove sensor joints.\n",
+ " if \"camera\" not in j\n",
+ " # Remove head and hands.\n",
+ " and \"neck\" not in j\n",
+ " and \"wrist\" not in j\n",
+ " and \"thumb\" not in j\n",
+ " and \"index\" not in j\n",
+ " and \"middle\" not in j\n",
+ " and \"ring\" not in j\n",
+ " and \"pinkie\" not in j\n",
+ " # Remove upper body.\n",
+ " and \"torso\" not in j and \"elbow\" not in j and \"shoulder\" not in j\n",
+ " ),\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "RLvAit_i2ZiA",
+ "outputId": "f6a77f31-e341-4350-c057-b1edabfe8f0e"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Model name: ergoCub\n",
+ "Number of links: 13\n",
+ "Number of joints: 12\n",
+ "\n",
+ "Links:\n",
+ "('root_link', 'l_hip_1', 'r_hip_1', 'l_hip_2', 'r_hip_2', 'l_upper_leg', 'r_upper_leg', 'l_lower_leg', 'r_lower_leg', 'l_ankle_1', 'r_ankle_1', 'l_ankle_2', 'r_ankle_2')\n",
+ "\n",
+ "Joints:\n",
+ "('l_hip_pitch', 'r_hip_pitch', 'l_hip_roll', 'r_hip_roll', 'l_hip_yaw', 'r_hip_yaw', 'l_knee', 'r_knee', 'l_ankle_pitch', 'r_ankle_pitch', 'l_ankle_roll', 'r_ankle_roll')\n",
+ "\n",
+ "Frames:\n",
+ "('base_link', 'base_link_fixed_joint', 'chest', 'head', 'head_imu_0', 'head_imu_0_fixed_joint', 'head_laser_frame', 'head_laser_frame_fixed_joint', 'imu_frame', 'imu_frame_fixed_joint', 'l_arm_ft', 'l_arm_ft_fixed_joint', 'l_foot_front', 'l_foot_front_ft', 'l_foot_front_ft_fixed_joint', 'l_foot_rear', 'l_foot_rear_ft', 'l_foot_rear_ft_fixed_joint', 'l_forearm', 'l_hand_index_1', 'l_hand_index_2', 'l_hand_index_3', 'l_hand_index_skin_0', 'l_hand_index_skin_0_fixed_joint', 'l_hand_index_skin_1', 'l_hand_index_skin_10', 'l_hand_index_skin_10_fixed_joint', 'l_hand_index_skin_11', 'l_hand_index_skin_11_fixed_joint', 'l_hand_index_skin_1_fixed_joint', 'l_hand_index_skin_2', 'l_hand_index_skin_2_fixed_joint', 'l_hand_index_skin_3', 'l_hand_index_skin_3_fixed_joint', 'l_hand_index_skin_4', 'l_hand_index_skin_4_fixed_joint', 'l_hand_index_skin_5', 'l_hand_index_skin_5_fixed_joint', 'l_hand_index_skin_6', 'l_hand_index_skin_6_fixed_joint', 'l_hand_index_skin_7', 'l_hand_index_skin_7_fixed_joint', 'l_hand_index_skin_8', 'l_hand_index_skin_8_fixed_joint', 'l_hand_index_skin_9', 'l_hand_index_skin_9_fixed_joint', 'l_hand_middle_1', 'l_hand_middle_2', 'l_hand_middle_skin_0', 'l_hand_middle_skin_0_fixed_joint', 'l_hand_middle_skin_1', 'l_hand_middle_skin_10', 'l_hand_middle_skin_10_fixed_joint', 'l_hand_middle_skin_11', 'l_hand_middle_skin_11_fixed_joint', 'l_hand_middle_skin_1_fixed_joint', 'l_hand_middle_skin_2', 'l_hand_middle_skin_2_fixed_joint', 'l_hand_middle_skin_3', 'l_hand_middle_skin_3_fixed_joint', 'l_hand_middle_skin_4', 'l_hand_middle_skin_4_fixed_joint', 'l_hand_middle_skin_5', 'l_hand_middle_skin_5_fixed_joint', 'l_hand_middle_skin_6', 'l_hand_middle_skin_6_fixed_joint', 'l_hand_middle_skin_7', 'l_hand_middle_skin_7_fixed_joint', 'l_hand_middle_skin_8', 'l_hand_middle_skin_8_fixed_joint', 'l_hand_middle_skin_9', 'l_hand_middle_skin_9_fixed_joint', 'l_hand_palm', 'l_hand_pinkie_1', 'l_hand_pinkie_2', 'l_hand_pinkie_skin_0', 'l_hand_pinkie_skin_0_fixed_joint', 'l_hand_pinkie_skin_1', 'l_hand_pinkie_skin_10', 'l_hand_pinkie_skin_10_fixed_joint', 'l_hand_pinkie_skin_11', 'l_hand_pinkie_skin_11_fixed_joint', 'l_hand_pinkie_skin_1_fixed_joint', 'l_hand_pinkie_skin_2', 'l_hand_pinkie_skin_2_fixed_joint', 'l_hand_pinkie_skin_3', 'l_hand_pinkie_skin_3_fixed_joint', 'l_hand_pinkie_skin_4', 'l_hand_pinkie_skin_4_fixed_joint', 'l_hand_pinkie_skin_5', 'l_hand_pinkie_skin_5_fixed_joint', 'l_hand_pinkie_skin_6', 'l_hand_pinkie_skin_6_fixed_joint', 'l_hand_pinkie_skin_7', 'l_hand_pinkie_skin_7_fixed_joint', 'l_hand_pinkie_skin_8', 'l_hand_pinkie_skin_8_fixed_joint', 'l_hand_pinkie_skin_9', 'l_hand_pinkie_skin_9_fixed_joint', 'l_hand_ring_1', 'l_hand_ring_2', 'l_hand_ring_skin_0', 'l_hand_ring_skin_0_fixed_joint', 'l_hand_ring_skin_1', 'l_hand_ring_skin_10', 'l_hand_ring_skin_10_fixed_joint', 'l_hand_ring_skin_11', 'l_hand_ring_skin_11_fixed_joint', 'l_hand_ring_skin_1_fixed_joint', 'l_hand_ring_skin_2', 'l_hand_ring_skin_2_fixed_joint', 'l_hand_ring_skin_3', 'l_hand_ring_skin_3_fixed_joint', 'l_hand_ring_skin_4', 'l_hand_ring_skin_4_fixed_joint', 'l_hand_ring_skin_5', 'l_hand_ring_skin_5_fixed_joint', 'l_hand_ring_skin_6', 'l_hand_ring_skin_6_fixed_joint', 'l_hand_ring_skin_7', 'l_hand_ring_skin_7_fixed_joint', 'l_hand_ring_skin_8', 'l_hand_ring_skin_8_fixed_joint', 'l_hand_ring_skin_9', 'l_hand_ring_skin_9_fixed_joint', 'l_hand_thumb_1', 'l_hand_thumb_2', 'l_hand_thumb_3', 'l_hand_thumb_skin_0', 'l_hand_thumb_skin_0_fixed_joint', 'l_hand_thumb_skin_1', 'l_hand_thumb_skin_10', 'l_hand_thumb_skin_10_fixed_joint', 'l_hand_thumb_skin_11', 'l_hand_thumb_skin_11_fixed_joint', 'l_hand_thumb_skin_1_fixed_joint', 'l_hand_thumb_skin_2', 'l_hand_thumb_skin_2_fixed_joint', 'l_hand_thumb_skin_3', 'l_hand_thumb_skin_3_fixed_joint', 'l_hand_thumb_skin_4', 'l_hand_thumb_skin_4_fixed_joint', 'l_hand_thumb_skin_5', 'l_hand_thumb_skin_5_fixed_joint', 'l_hand_thumb_skin_6', 'l_hand_thumb_skin_6_fixed_joint', 'l_hand_thumb_skin_7', 'l_hand_thumb_skin_7_fixed_joint', 'l_hand_thumb_skin_8', 'l_hand_thumb_skin_8_fixed_joint', 'l_hand_thumb_skin_9', 'l_hand_thumb_skin_9_fixed_joint', 'l_hip_3', 'l_leg_ft', 'l_leg_ft_fixed_joint', 'l_shoulder_1', 'l_shoulder_2', 'l_shoulder_3', 'l_sole', 'l_sole_fixed_joint', 'l_upper_arm', 'l_wrist_1', 'l_wrist_2', 'neck_2', 'neck_3', 'r_arm_ft', 'r_arm_ft_fixed_joint', 'r_foot_front', 'r_foot_front_ft', 'r_foot_front_ft_fixed_joint', 'r_foot_rear', 'r_foot_rear_ft', 'r_foot_rear_ft_fixed_joint', 'r_forearm', 'r_hand_index_1', 'r_hand_index_2', 'r_hand_index_3', 'r_hand_index_skin_0', 'r_hand_index_skin_0_fixed_joint', 'r_hand_index_skin_1', 'r_hand_index_skin_10', 'r_hand_index_skin_10_fixed_joint', 'r_hand_index_skin_11', 'r_hand_index_skin_11_fixed_joint', 'r_hand_index_skin_1_fixed_joint', 'r_hand_index_skin_2', 'r_hand_index_skin_2_fixed_joint', 'r_hand_index_skin_3', 'r_hand_index_skin_3_fixed_joint', 'r_hand_index_skin_4', 'r_hand_index_skin_4_fixed_joint', 'r_hand_index_skin_5', 'r_hand_index_skin_5_fixed_joint', 'r_hand_index_skin_6', 'r_hand_index_skin_6_fixed_joint', 'r_hand_index_skin_7', 'r_hand_index_skin_7_fixed_joint', 'r_hand_index_skin_8', 'r_hand_index_skin_8_fixed_joint', 'r_hand_index_skin_9', 'r_hand_index_skin_9_fixed_joint', 'r_hand_middle_1', 'r_hand_middle_2', 'r_hand_middle_skin_0', 'r_hand_middle_skin_0_fixed_joint', 'r_hand_middle_skin_1', 'r_hand_middle_skin_10', 'r_hand_middle_skin_10_fixed_joint', 'r_hand_middle_skin_11', 'r_hand_middle_skin_11_fixed_joint', 'r_hand_middle_skin_1_fixed_joint', 'r_hand_middle_skin_2', 'r_hand_middle_skin_2_fixed_joint', 'r_hand_middle_skin_3', 'r_hand_middle_skin_3_fixed_joint', 'r_hand_middle_skin_4', 'r_hand_middle_skin_4_fixed_joint', 'r_hand_middle_skin_5', 'r_hand_middle_skin_5_fixed_joint', 'r_hand_middle_skin_6', 'r_hand_middle_skin_6_fixed_joint', 'r_hand_middle_skin_7', 'r_hand_middle_skin_7_fixed_joint', 'r_hand_middle_skin_8', 'r_hand_middle_skin_8_fixed_joint', 'r_hand_middle_skin_9', 'r_hand_middle_skin_9_fixed_joint', 'r_hand_palm', 'r_hand_pinkie_1', 'r_hand_pinkie_2', 'r_hand_pinkie_skin_0', 'r_hand_pinkie_skin_0_fixed_joint', 'r_hand_pinkie_skin_1', 'r_hand_pinkie_skin_10', 'r_hand_pinkie_skin_10_fixed_joint', 'r_hand_pinkie_skin_11', 'r_hand_pinkie_skin_11_fixed_joint', 'r_hand_pinkie_skin_1_fixed_joint', 'r_hand_pinkie_skin_2', 'r_hand_pinkie_skin_2_fixed_joint', 'r_hand_pinkie_skin_3', 'r_hand_pinkie_skin_3_fixed_joint', 'r_hand_pinkie_skin_4', 'r_hand_pinkie_skin_4_fixed_joint', 'r_hand_pinkie_skin_5', 'r_hand_pinkie_skin_5_fixed_joint', 'r_hand_pinkie_skin_6', 'r_hand_pinkie_skin_6_fixed_joint', 'r_hand_pinkie_skin_7', 'r_hand_pinkie_skin_7_fixed_joint', 'r_hand_pinkie_skin_8', 'r_hand_pinkie_skin_8_fixed_joint', 'r_hand_pinkie_skin_9', 'r_hand_pinkie_skin_9_fixed_joint', 'r_hand_ring_1', 'r_hand_ring_2', 'r_hand_ring_skin_0', 'r_hand_ring_skin_0_fixed_joint', 'r_hand_ring_skin_1', 'r_hand_ring_skin_10', 'r_hand_ring_skin_10_fixed_joint', 'r_hand_ring_skin_11', 'r_hand_ring_skin_11_fixed_joint', 'r_hand_ring_skin_1_fixed_joint', 'r_hand_ring_skin_2', 'r_hand_ring_skin_2_fixed_joint', 'r_hand_ring_skin_3', 'r_hand_ring_skin_3_fixed_joint', 'r_hand_ring_skin_4', 'r_hand_ring_skin_4_fixed_joint', 'r_hand_ring_skin_5', 'r_hand_ring_skin_5_fixed_joint', 'r_hand_ring_skin_6', 'r_hand_ring_skin_6_fixed_joint', 'r_hand_ring_skin_7', 'r_hand_ring_skin_7_fixed_joint', 'r_hand_ring_skin_8', 'r_hand_ring_skin_8_fixed_joint', 'r_hand_ring_skin_9', 'r_hand_ring_skin_9_fixed_joint', 'r_hand_thumb_1', 'r_hand_thumb_2', 'r_hand_thumb_3', 'r_hand_thumb_skin_0', 'r_hand_thumb_skin_0_fixed_joint', 'r_hand_thumb_skin_1', 'r_hand_thumb_skin_10', 'r_hand_thumb_skin_10_fixed_joint', 'r_hand_thumb_skin_11', 'r_hand_thumb_skin_11_fixed_joint', 'r_hand_thumb_skin_1_fixed_joint', 'r_hand_thumb_skin_2', 'r_hand_thumb_skin_2_fixed_joint', 'r_hand_thumb_skin_3', 'r_hand_thumb_skin_3_fixed_joint', 'r_hand_thumb_skin_4', 'r_hand_thumb_skin_4_fixed_joint', 'r_hand_thumb_skin_5', 'r_hand_thumb_skin_5_fixed_joint', 'r_hand_thumb_skin_6', 'r_hand_thumb_skin_6_fixed_joint', 'r_hand_thumb_skin_7', 'r_hand_thumb_skin_7_fixed_joint', 'r_hand_thumb_skin_8', 'r_hand_thumb_skin_8_fixed_joint', 'r_hand_thumb_skin_9', 'r_hand_thumb_skin_9_fixed_joint', 'r_hip_3', 'r_leg_ft', 'r_leg_ft_fixed_joint', 'r_shoulder_1', 'r_shoulder_2', 'r_shoulder_3', 'r_sole', 'r_sole_fixed_joint', 'r_upper_arm', 'r_wrist_1', 'r_wrist_2', 'realsense', 'realsense_depth_frame', 'realsense_depth_frame_fixed_joint', 'realsense_imu_0', 'realsense_imu_0_fixed_joint', 'realsense_rgb_frame', 'realsense_rgb_frame_fixed_joint', 'torso_1', 'torso_2', 'waist_imu_0', 'waist_imu_0_fixed_joint')\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Print model quantities.\n",
+ "print(f\"Model name: {model.name()}\")\n",
+ "print(f\"Number of links: {model.number_of_links()}\")\n",
+ "print(f\"Number of joints: {model.number_of_joints()}\")\n",
+ "\n",
+ "print()\n",
+ "print(f\"Links:\\n{model.link_names()}\")\n",
+ "\n",
+ "print()\n",
+ "print(f\"Joints:\\n{model.joint_names()}\")\n",
+ "\n",
+ "print()\n",
+ "print(f\"Frames:\\n{model.frame_names()}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "Xp8V5on5twZ8",
+ "outputId": "5e1c50e3-267d-42aa-fd1c-6cac3e246790"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "W_H_B: shape=(4, 4)\n",
+ "[[ 0.60959 0.05722 -0.79065 -0.05375]\n",
+ " [-0.78842 0.14744 -0.59721 0.62771]\n",
+ " [ 0.0824 0.98741 0.13499 0.79038]\n",
+ " [ 0. 0. 0. 1. ]]\n",
+ "\n",
+ "s: shape=(12,)\n",
+ "[ 0.98172 1.02608 0.25454 0.98976 -0.5168 -1.19545 -0.28428 -0.01129 -0.68968 0.59422 -0.19612 -0.02112]\n",
+ "\n",
+ "ν: shape=(18,)\n",
+ "[-0.63478 0.653 -0.52829 0.70579 0.50915 0.62625 0.51191 0.6864 -0.13997 -0.17041 0.52786 -0.28581 -0.00081 0.54389 -0.22844 0.44337\n",
+ " 0.05165 0.55334]\n",
+ "\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Create a random data object from the reduced model.\n",
+ "data = js.data.random_model_data(model=model)\n",
+ "\n",
+ "# Print the default state.\n",
+ "W_H_B, s = data.generalized_position()\n",
+ "ν = data.generalized_velocity()\n",
+ "\n",
+ "print(f\"W_H_B: shape={W_H_B.shape}\\n{W_H_B}\\n\")\n",
+ "print(f\"s: shape={s.shape}\\n{s}\\n\")\n",
+ "print(f\"ν: shape={ν.shape}\\n{ν}\\n\") # noqa: RUF001"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "XLx3sv9VtwZ9",
+ "outputId": "7f83bcd0-1099-44b1-e407-2c719aca6895"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "link_forces: shape=(13, 6)\n",
+ "joint_force_references: shape=(12,)\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Create a random link forces matrix.\n",
+ "link_forces = jax.random.uniform(\n",
+ " minval=-10.0,\n",
+ " maxval=10.0,\n",
+ " shape=(model.number_of_links(), 6),\n",
+ " key=jax.random.PRNGKey(0),\n",
+ ")\n",
+ "\n",
+ "# Create a random joint forces references vector.\n",
+ "# Note that these are called references because the actual joint forces that\n",
+ "# are actuated might differ due to effects like joint friction.\n",
+ "joint_force_references = jax.random.uniform(\n",
+ " minval=-10.0, maxval=10.0, shape=(model.dofs(),), key=jax.random.PRNGKey(0)\n",
+ ")\n",
+ "\n",
+ "# Create the references object.\n",
+ "references = js.references.JaxSimModelReferences.build(\n",
+ " model=model,\n",
+ " data=data,\n",
+ " link_forces=link_forces,\n",
+ " joint_force_references=joint_force_references,\n",
+ ")\n",
+ "\n",
+ "print(f\"link_forces: shape={references.link_forces(model=model, data=data).shape}\")\n",
+ "print(f\"joint_force_references: shape={references.joint_force_references(model=model).shape}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "AaG817vP4LfT"
+ },
+ "source": [
+ "## Robot Kinematics\n",
+ "\n",
+ "JaxSim offers functional APIs for computing kinematic quantities:\n",
+ "\n",
+ "- **`jaxsim.api.model`**: vectorized functions operating on the whole model.\n",
+ "- **`jaxsim.api.link`**: functions operating on individual links.\n",
+ "\n",
+ "Due to JAX limitations on vectorizable data types, many APIs operate on indices instead of names. Since using indices can be error prone, JaxSim provides conversion functions:\n",
+ "\n",
+ "- **jaxsim.api.link.names_to_idxs()**\n",
+ "- **jaxsim.api.link.idxs_to_names()**\n",
+ "\n",
+ "We recommend using names whenever possible to avoid hard-to-trace errors.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "metadata": {
+ "id": "QxImwfZz7pz-"
+ },
+ "outputs": [],
+ "source": [
+ "# Find the index of a link.\n",
+ "link_name = \"l_ankle_2\"\n",
+ "link_index = js.link.name_to_idx(model=model, link_name=link_name)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "C22Iqu2i4G-I",
+ "outputId": "cf1614dc-6c3c-442e-adf8-d3344a331d4a"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Transform of 'l_ankle_2': shape=(4, 4)\n",
+ "[[-0.53839 -0.15819 -0.82771 0.65199]\n",
+ " [-0.73956 -0.38215 0.55409 0.53659]\n",
+ " [-0.40396 0.91046 0.08876 1.07258]\n",
+ " [ 0. 0. 0. 1. ]]\n"
+ ]
+ }
+ ],
+ "source": [
+ "# @title Pose\n",
+ "\n",
+ "# Compute its pose w.r.t. the world frame through forward kinematics.\n",
+ "W_H_L = js.link.transform(model=model, data=data, link_index=link_index)\n",
+ "\n",
+ "# Alternatively, the same transform can be extracted from the matrix of\n",
+ "# all link transforms.\n",
+ "assert jnp.allclose(\n",
+ " W_H_L, js.model.forward_kinematics(model=model, data=data)[link_index]\n",
+ ")\n",
+ "\n",
+ "print(f\"Transform of '{link_name}': shape={W_H_L.shape}\\n{W_H_L}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "DnSpE_f97RkX",
+ "outputId": "47765e57-3fc1-40b6-ac20-04392aba85ef"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Body-fixed velocity L_v_WL=[ 0.50782 -0.43394 0.27792 -0.58574 -0.55775 0.15607]\n",
+ "Mixed velocity: LW_v_WL=[-0.4348 -0.05574 -0.57556 0.27441 0.73281 -0.25734]\n",
+ "Inertial-fixed velocity: W_v_WL=[-1.35889 0.40637 -0.24502 0.27441 0.73281 -0.25734]\n"
+ ]
+ }
+ ],
+ "source": [
+ "# @title 6D Velocity\n",
+ "\n",
+ "# JaxSim allows to select the so-called reprentation of the frame velocity:\n",
+ "L_v_WL = js.link.velocity(model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Body)\n",
+ "LW_v_WL = js.link.velocity(model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Mixed)\n",
+ "W_v_WL = js.link.velocity(model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Inertial)\n",
+ "\n",
+ "print(f\"Body-fixed velocity L_v_WL={L_v_WL}\")\n",
+ "print(f\"Mixed velocity: LW_v_WL={LW_v_WL}\")\n",
+ "print(f\"Inertial-fixed velocity: W_v_WL={W_v_WL}\")\n",
+ "\n",
+ "# These can also be computed passing through the link free-floating Jacobian.\n",
+ "# This type of Jacobian has a input velocity representation that corresponds\n",
+ "# the velocity representation of ν, and an output velocity representation that\n",
+ "# corresponds to the velocity representation of the desired 6D velocity.\n",
+ "\n",
+ "# You can use the following context manager to easily switch between representations:\n",
+ "with data.switch_velocity_representation(VelRepr.Body):\n",
+ "\n",
+ " # Body-fixed generalized velocity.\n",
+ " B_ν = data.generalized_velocity()\n",
+ "\n",
+ " # Free-floating Jacobian accepting a body-fixed generalized velocity and\n",
+ " # returning an inertial-fixed link velocity.\n",
+ " W_J_WL_B = js.link.jacobian(\n",
+ " model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Inertial\n",
+ " )\n",
+ "\n",
+ "# Now the following relation should hold.\n",
+ "assert jnp.allclose(W_v_WL, W_J_WL_B @ B_ν)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "SSoziCShtwZ9"
+ },
+ "source": [
+ "## Robot Dynamics\n",
+ "\n",
+ "JaxSim provides all the quantities involved in the equations of motion, restated here:\n",
+ "\n",
+ "$$\n",
+ "M(\\mathbf{q}) \\dot{\\boldsymbol{\\nu}} + \\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = B \\boldsymbol{\\tau} + \\sum_{L \\in \\mathcal{L}} J_{W,L}^\\top(\\mathbf{q}) \\: \\mathbf{f}_L\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "Specifically, it can compute:\n",
+ "\n",
+ "- $M(\\mathbf{q}) \\in \\mathbb{R}^{(6+n)\\times(6+n)}$: the mass matrix.\n",
+ "- $\\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) \\in \\mathbb{R}^{6+n}$: the vector of bias forces.\n",
+ "- $B \\in \\mathbb{R}^{(6+n) \\times n}$ the joint selector matrix.\n",
+ "- $J_{W,L} \\in \\mathbb{R}^{6 \\times (6+n)}$ the Jacobian of link $L$.\n",
+ "\n",
+ "Often, for convenience, link Jacobians are stacked together. Since JaxSim efficiently computes the Jacobians for all links, using the stacked version is recommended when needed:\n",
+ "\n",
+ "$$\n",
+ "M(\\mathbf{q}) \\dot{\\boldsymbol{\\nu}} + \\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = B \\boldsymbol{\\tau} + J_{W,\\mathcal{L}}^\\top(\\mathbf{q}) \\: \\mathbf{f}_\\mathcal{L}\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "Furthermore, there are applications that require unpacking the vector of bias forces as follow:\n",
+ "\n",
+ "$$\n",
+ "\\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = C(\\mathbf{q}, \\boldsymbol{\\nu}) \\boldsymbol{\\nu} + \\mathbf{g}(\\mathbf{q})\n",
+ ",\n",
+ "$$\n",
+ "\n",
+ "where:\n",
+ "\n",
+ "- $\\mathbf{g}(\\mathbf{q}) \\in \\mathbb{R}^{6+n}$: the vector of gravity forces.\n",
+ "- $C(\\mathbf{q}, \\boldsymbol{\\nu}) \\in \\mathbb{R}^{(6+n)\\times(6+n)}$: the Coriolis matrix.\n",
+ "\n",
+ "Here below we report the functions to compute all these quantities. Note that all quantities depend on the active velocity representation of `data`. As it was done for the link velocity, it is possible to change the representation associated to all the computed quantities by operating within the corresponding context manager. Here below we consider the default representation of data."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "fVp_xP_1twZ9",
+ "outputId": "b06f20a0-6c01-435c-e353-502045234487"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Velocity representation of data: VelRepr.Inertial \n",
+ "\n",
+ "M: shape=(18, 18)\n",
+ "h: shape=(18,)\n",
+ "g: shape=(18,)\n",
+ "C: shape=(18, 18)\n",
+ "B: shape=(18, 12)\n",
+ "J: shape=(13, 6, 18)\n",
+ "τ: shape=(12,)\n",
+ "f_L: shape=(13, 6)\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(\"Velocity representation of data:\", data.velocity_representation, \"\\n\")\n",
+ "\n",
+ "# Compute the mass matrix.\n",
+ "M = js.model.free_floating_mass_matrix(model=model, data=data)\n",
+ "print(f\"M: shape={M.shape}\")\n",
+ "\n",
+ "# Compute the vector of bias forces.\n",
+ "h = js.model.free_floating_bias_forces(model=model, data=data)\n",
+ "print(f\"h: shape={h.shape}\")\n",
+ "\n",
+ "# Compute the vector of gravity forces.\n",
+ "g = js.model.free_floating_gravity_forces(model=model, data=data)\n",
+ "print(f\"g: shape={g.shape}\")\n",
+ "\n",
+ "# Compute the Coriolis matrix.\n",
+ "C = js.model.free_floating_coriolis_matrix(model=model, data=data)\n",
+ "print(f\"C: shape={C.shape}\")\n",
+ "\n",
+ "# Create a the joint selector matrix.\n",
+ "B = jnp.block([jnp.zeros(shape=(model.dofs(), 6)), jnp.eye(model.dofs())]).T\n",
+ "print(f\"B: shape={B.shape}\")\n",
+ "\n",
+ "# Compute the stacked tensor of link Jacobians.\n",
+ "J = js.model.generalized_free_floating_jacobian(model=model, data=data)\n",
+ "print(f\"J: shape={J.shape}\")\n",
+ "\n",
+ "# Extract the joint forces from the references object.\n",
+ "τ = references.joint_force_references(model=model)\n",
+ "print(f\"τ: shape={τ.shape}\")\n",
+ "\n",
+ "# Extract the link forces from the references object.\n",
+ "f_L = references.link_forces(model=model, data=data)\n",
+ "print(f\"f_L: shape={f_L.shape}\")\n",
+ "\n",
+ "# The following relation should hold.\n",
+ "assert jnp.allclose(h, C @ ν + g)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "QqaqxneEFYiW"
+ },
+ "source": [
+ "### Forward Dynamics\n",
+ "\n",
+ "$$\n",
+ "\\dot{\\boldsymbol{\\nu}} = \\text{FD}(\\mathbf{q}, \\boldsymbol{\\nu}, \\boldsymbol{\\tau}, \\mathbf{f}_{\\mathcal{L}})\n",
+ "$$\n",
+ "\n",
+ "JaxSim provides two alternative methods to compute the forward dynamics:\n",
+ "\n",
+ "1. Operate on the quantities of the equations of motion.\n",
+ "2. Call the recursive Articulated Body Algorithm (ABA).\n",
+ "\n",
+ "The physics engine provided by JaxSim exploits the efficient calculation of the forward dynamics with ABA for simulating the trajectories of the system dynamics."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "d_vp6D74GoVZ",
+ "outputId": "ae1c1b59-f5ce-4539-bca0-5d34d62995c8"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "ν̇: shape=(18,)\n"
+ ]
+ }
+ ],
+ "source": [
+ "ν̇_eom = jnp.linalg.pinv(M) @ (B @ τ - h + jnp.einsum(\"l6g,l6->g\", J, f_L))\n",
+ "\n",
+ "v̇_WB, s̈ = js.model.forward_dynamics_aba(\n",
+ " model=model, data=data, link_forces=f_L, joint_forces=joint_force_references\n",
+ ")\n",
+ "\n",
+ "ν̇_aba = jnp.hstack([v̇_WB, s̈])\n",
+ "print(f\"ν̇: shape={ν̇_aba.shape}\") # noqa: RUF001\n",
+ "\n",
+ "# The following relation should hold.\n",
+ "assert jnp.allclose(ν̇_eom, ν̇_aba)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "oOKJOVfsH4Ki"
+ },
+ "source": [
+ "### Inverse Dynamics\n",
+ "\n",
+ "$$\n",
+ "(\\boldsymbol{\\tau}, \\, \\mathbf{f}_B) = \\text{ID}(\\mathbf{q}, \\boldsymbol{\\nu}, \\dot{\\boldsymbol{\\nu}}, \\mathbf{f}_{\\mathcal{L}})\n",
+ "$$\n",
+ "\n",
+ "JaxSim offers two methods to compute inverse dynamics:\n",
+ "\n",
+ "- Directly use the quantities from the equations of motion.\n",
+ "- Use the Recursive Newton-Euler Algorithm (RNEA).\n",
+ "\n",
+ "Unlike many other implementations, JaxSim's RNEA for floating-base systems is the true inverse of $\\text{FD}$. It also computes the 6D force applied to the base link that generates the base acceleration."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "FlNo8dNWKKtu",
+ "outputId": "867e7568-016d-4418-a2bd-ab7d75fa69a7"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "f_B: shape=(6,)\n",
+ "τ_rnea: shape=(12,)\n"
+ ]
+ }
+ ],
+ "source": [
+ "f_B, τ_rnea = js.model.inverse_dynamics(\n",
+ " model=model,\n",
+ " data=data,\n",
+ " base_acceleration=v̇_WB,\n",
+ " joint_accelerations=s̈,\n",
+ " # To check that f_B works, let's remove the force applied\n",
+ " # to the base link from the link forces.\n",
+ " link_forces=f_L.at[0].set(jnp.zeros(6))\n",
+ ")\n",
+ "\n",
+ "print(f\"f_B: shape={f_B.shape}\")\n",
+ "print(f\"τ_rnea: shape={τ_rnea.shape}\")\n",
+ "\n",
+ "# The following relations should hold.\n",
+ "assert jnp.allclose(τ_rnea, τ)\n",
+ "assert jnp.allclose(f_B, link_forces[0])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "LXARuRu1Ly1K"
+ },
+ "source": [
+ "### Centroidal Dynamics\n",
+ "\n",
+ "Centroidal dynamics is a useful simplification often employed in planning and control applications. It represents the dynamics projected onto a mixed frame associated with the center of mass (CoM):\n",
+ "\n",
+ "$$\n",
+ "G = G[W] = ({}^W \\mathbf{p}_{\\text{CoM}}, [W])\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "The governing equations for centroidal dynamics take into account the 6D total momentum:\n",
+ "\n",
+ "$$\n",
+ "{}_G \\mathbf{h} =\n",
+ "\\begin{bmatrix}\n",
+ "{}_G \\mathbf{h}^l \\\\ {}_G \\mathbf{h}^\\omega\n",
+ "\\end{bmatrix} =\n",
+ "\\begin{bmatrix}\n",
+ "m \\, {}^W \\dot{\\mathbf{p}}_\\text{CoM} \\\\ {}_G \\mathbf{h}^\\omega\n",
+ "\\end{bmatrix}\n",
+ "\\in \\mathbb{R}^6\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "The equations of centroidal dynamics can be expressed as:\n",
+ "\n",
+ "$$\n",
+ "{}_G \\dot{\\mathbf{h}} =\n",
+ "m \\,\n",
+ "\\begin{bmatrix}\n",
+ "{}^W \\mathbf{g} \\\\ \\mathbf{0}_3\n",
+ "\\end{bmatrix} +\n",
+ "\\sum_{C_i \\in \\mathcal{C}} {}_G \\mathbf{X}^{C_i} \\, {}_{C_i} \\mathbf{f}_i\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "While centroidal dynamics can function independently by considering the total mass $m \\in \\mathbb{R}$ of the robot and the transformations for 6D contact forces ${}_G \\mathbf{X}^{C_i}$ corresponding to the pose ${}^G \\mathbf{H}_{C_i} \\in \\text{SE}(3)$ of the contact frames, advanced kino-dynamic methods require a relationship between full kinematics and centroidal dynamics. This is typically achieved through the _Centroidal Momentum Matrix_ (also known as the _centroidal momentum Jacobian_):\n",
+ "\n",
+ "$$\n",
+ "{}_G \\mathbf{h} = J_\\text{CMM}(\\mathbf{q}) \\, \\boldsymbol{\\nu}\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "JaxSim offers APIs to compute all these quantities (and many more) in the `jaxsim.api.com` package."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 15,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "g5GOYXDnLySU",
+ "outputId": "b6e5df8f-1d56-4208-a5b0-87dfa1737ca4"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Number of contact points: 32 \n",
+ "\n",
+ "G_h: shape=(6,)\n",
+ "J_CMM: shape=(6, 18)\n",
+ "G_Xf_C: shape=(32, 6, 6)\n",
+ "G_ḣ: shape=(6,)\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Number of contact points.\n",
+ "n_cp = len(model.kin_dyn_parameters.contact_parameters.body)\n",
+ "print(\"Number of contact points:\", n_cp, \"\\n\")\n",
+ "\n",
+ "# Compute the centroidal momentum.\n",
+ "J_CMM = js.com.centroidal_momentum_jacobian(model=model, data=data)\n",
+ "G_h = J_CMM @ ν\n",
+ "print(f\"G_h: shape={G_h.shape}\")\n",
+ "print(f\"J_CMM: shape={J_CMM.shape}\")\n",
+ "\n",
+ "# The following relation should hold.\n",
+ "assert jnp.allclose(G_h, js.com.centroidal_momentum(model=model, data=data))\n",
+ "\n",
+ "# If we consider all contact points of the model as active\n",
+ "# (discourages since they might be too many), the 6D transforms of\n",
+ "# collidable points can be computed as follows:\n",
+ "W_H_C = js.contact.transforms(model=model, data=data)\n",
+ "\n",
+ "# Compute the pose of the G frame.\n",
+ "W_p_CoM = js.com.com_position(model=model, data=data)\n",
+ "G_H_W = jaxsim.math.Transform.inverse(jnp.eye(4).at[0:3, 3].set(W_p_CoM))\n",
+ "\n",
+ "# Convert from SE(3) to the transforms for 6D forces.\n",
+ "G_Xf_C = jax.vmap(\n",
+ " lambda W_H_Ci: jaxsim.math.Adjoint.from_transform(\n",
+ " transform=G_H_W @ W_H_Ci, inverse=True\n",
+ " )\n",
+ ")(W_H_C)\n",
+ "print(f\"G_Xf_C: shape={G_Xf_C.shape}\")\n",
+ "\n",
+ "# Let's create random 3D linear forces applied to the contact points.\n",
+ "C_fl = jax.random.uniform(\n",
+ " minval=-10.0,\n",
+ " maxval=10.0,\n",
+ " shape=(n_cp, 3),\n",
+ " key=jax.random.PRNGKey(0),\n",
+ ")\n",
+ "\n",
+ "# Compute the 3D gravity vector and the total mass of the robot.\n",
+ "W_g = data.gravity\n",
+ "m = js.model.total_mass(model=model)\n",
+ "\n",
+ "# The centroidal dynamics can be computed as follows.\n",
+ "G_ḣ = 0\n",
+ "G_ḣ += m * jnp.hstack([W_g, jnp.zeros(3)])\n",
+ "G_ḣ += jnp.einsum(\"c66,c6->6\", G_Xf_C, jnp.hstack([C_fl, jnp.zeros_like(C_fl)]))\n",
+ "print(f\"G_ḣ: shape={G_ḣ.shape}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "UTae5MjhaP2H"
+ },
+ "source": [
+ "## Contact Frames\n",
+ "\n",
+ "Many control and planning applications require projecting the floating-base dynamics into the contact space or computing quantities related to active contact points, such as enforcing holonomic constraints.\n",
+ "\n",
+ "The underlying theory for these applications becomes clearer in a mixed representation. Specifically, the position, linear velocity, and linear acceleration of contact points in their corresponding mixed frame align with the numerical derivatives of their coordinate vectors.\n",
+ "\n",
+ "Key methodologies in this area may involve the Delassus matrix:\n",
+ "\n",
+ "$$\n",
+ "\\Psi(\\mathbf{q}) = J_{W,C}^T(\\mathbf{q}) \\, M(\\mathbf{q})^{-1} \\, J_{W,C}^T(\\mathbf{q})\n",
+ "$$\n",
+ "\n",
+ "or the linear acceleration of a contact point:\n",
+ "\n",
+ "$$\n",
+ "{}^W \\ddot{\\mathbf{p}}_C = \\frac{\\text{d} (J^l_{W,C} \\boldsymbol{\\nu})}{\\text{d}t}\n",
+ "= \\dot{J}^l_{W,C} \\boldsymbol{\\nu} + J^l_{W,C} \\dot{\\boldsymbol{\\nu}}\n",
+ ".\n",
+ "$$\n",
+ "\n",
+ "JaxSim offers APIs to compute all these quantities (and many more) in the `jaxsim.api.contact` package."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "gYZ1jK1Neg1H",
+ "outputId": "274a628d-efad-4738-b500-bb2927f3e826"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Ψ: shape=(96, 96)\n",
+ "W_H_C: shape=(32, 4, 4)\n",
+ "W_ṗ_B: shape=(32, 3)\n",
+ "W_p̈_C: shape=(32, 3)\n"
+ ]
+ }
+ ],
+ "source": [
+ "with (\n",
+ " data.switch_velocity_representation(VelRepr.Mixed),\n",
+ " references.switch_velocity_representation(VelRepr.Mixed),\n",
+ "):\n",
+ "\n",
+ " # Compute the mixed generalized velocity.\n",
+ " BW_ν = data.generalized_velocity()\n",
+ "\n",
+ " # Compute the mixed generalized acceleration.\n",
+ " BW_ν̇ = jnp.hstack(\n",
+ " js.model.forward_dynamics(\n",
+ " model=model,\n",
+ " data=data,\n",
+ " link_forces=references.link_forces(model=model, data=data),\n",
+ " joint_forces=references.joint_force_references(model=model),\n",
+ " )\n",
+ " )\n",
+ "\n",
+ " # Compute the mass matrix in mixed representation.\n",
+ " BW_M = js.model.free_floating_mass_matrix(model=model, data=data)\n",
+ "\n",
+ " # Compute the contact Jacobian and its derivative.\n",
+ " Jl_WC = js.contact.jacobian(model=model, data=data)[:, 0:3, :]\n",
+ " J̇l_WC = js.contact.jacobian_derivative(model=model, data=data)[:, 0:3, :]\n",
+ "\n",
+ "# Compute the Delassus matrix.\n",
+ "Ψ = jnp.vstack(Jl_WC) @ jnp.linalg.lstsq(BW_M, jnp.vstack(J̇l_WC).T)[0]\n",
+ "print(f\"Ψ: shape={Ψ.shape}\")\n",
+ "\n",
+ "# Compute the transforms of the mixed frames implicitly associated\n",
+ "# to each collidable point.\n",
+ "W_H_C = js.contact.transforms(model=model, data=data)\n",
+ "print(f\"W_H_C: shape={W_H_C.shape}\")\n",
+ "\n",
+ "# Compute the linear velocity of the collidable points.\n",
+ "with data.switch_velocity_representation(VelRepr.Mixed):\n",
+ " W_ṗ_B = js.contact.collidable_point_velocities(model=model, data=data)[:, 0:3]\n",
+ " print(f\"W_ṗ_B: shape={W_ṗ_B.shape}\")\n",
+ "\n",
+ "# Compute the linear acceleration of the collidable points.\n",
+ "W_p̈_C = 0\n",
+ "W_p̈_C += jnp.einsum(\"c3g,g->c3\", J̇l_WC, BW_ν̇)\n",
+ "W_p̈_C += jnp.einsum(\"c3g,g->c3\", Jl_WC, BW_ν)\n",
+ "print(f\"W_p̈_C: shape={W_p̈_C.shape}\")\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "rrSfxp8lh9YZ"
+ },
+ "source": [
+ "### Contact Forces\n",
+ "\n",
+ "To conclude, JaxSim also provides different contact models to compute the contact forces. However, these forces are mainly related to the usage of JaxSim as a simulator rather than as a multibody dynamics library. This use case is shown in more details in other examples.\n",
+ "\n",
+ "Here below we report the function to call to compute the contact forces using the default contact model stored in the `JaxSimModel`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 422
+ },
+ "id": "Ot6HePB_twaE",
+ "outputId": "5d1d717c-4523-402d-9763-1d485720aa3b"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Active contact model: \n",
+ "\n"
+ ]
+ },
+ {
+ "ename": "TypeError",
+ "evalue": "collidable_point_forces() got an unexpected keyword argument 'joint_force_references'",
+ "output_type": "error",
+ "traceback": [
+ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+ "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
+ "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[0;31m# Compute the contact forces.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m W_f_C = js.contact.collidable_point_forces(\n\u001b[0m\u001b[1;32m 5\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ " \u001b[0;31m[... skipping hidden 10 frame]\u001b[0m\n",
+ "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/linear_util.py\u001b[0m in \u001b[0;36mcall_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 191\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 192\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 193\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 194\u001b[0m \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 195\u001b[0m \u001b[0;31m# Some transformations yield from inside context managers, so we have to\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;31mTypeError\u001b[0m: collidable_point_forces() got an unexpected keyword argument 'joint_force_references'"
+ ]
+ }
+ ],
+ "source": [
+ "print(f\"Active contact model: {type(model.contact_model)}\\n\")\n",
+ "\n",
+ "# Compute the contact forces.\n",
+ "W_f_C = js.contact.collidable_point_forces(\n",
+ " model=model,\n",
+ " data=data,\n",
+ " link_forces=references.link_forces(model=model, data=data),\n",
+ " joint_force_references=references.joint_force_references(model=model),\n",
+ ")\n",
+ "print(f\"W_f_C: shape={W_f_C.shape}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "LITRC3STliKR"
+ },
+ "source": [
+ "## Conclusions\n",
+ "\n",
+ "This notebook provided an overview of the main APIs in JaxSim for its use as a multibody dynamics library. Here are a few key points to remember:\n",
+ "\n",
+ "- Explore all the modules in the `jaxsim.api` package to discover the full range of APIs available. Many more functionalities exist beyond what was covered in this notebook.\n",
+ "- All APIs follow a functional approach, consistent with the JAX programming style.\n",
+ "- This functional design allows for easy application of `jax.vmap` to execute functions in parallel on hardware accelerators.\n",
+ "- Since the entire multibody dynamics library is built with JAX, it natively supports `jax.grad`, `jax.jacfwd`, and `jax.jacrev` transformations, enabling automatic differentiation through complex logic without additional effort.\n",
+ "\n",
+ "Have fun!"
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "provenance": [],
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "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.10.15"
+ }
},
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# @title Imports and setup\n",
- "from IPython.display import clear_output\n",
- "import os\n",
- "import pathlib\n",
- "import sys\n",
- "import robot_descriptions\n",
- "\n",
- "IS_COLAB = \"google.colab\" in sys.modules\n",
- "\n",
- "# Install JAX and Gazebo SDF\n",
- "if IS_COLAB:\n",
- " !{sys.executable} -m pip install -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",
- "\n",
- "import jax\n",
- "import jax.numpy as jnp\n",
- "from jaxsim import logging\n",
- "\n",
- "logging.set_logging_level(logging.LoggingLevel.INFO)\n",
- "logging.info(f\"Running on {jax.devices()}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "We will use a model of the [ErgoCub](https://ergocub.eu/) humanoid robot."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# @title Fetch the URDF file\n",
- "try:\n",
- " os.environ[\"ROBOT_DESCRIPTION_COMMIT\"] = \"v0.7.1\"\n",
- "\n",
- " import robot_descriptions.ergocub_description\n",
- "\n",
- "finally:\n",
- " _ = os.environ.pop(\"ROBOT_DESCRIPTION_COMMIT\", None)\n",
- "\n",
- "model_urdf_path = pathlib.Path(\n",
- " robot_descriptions.ergocub_description.URDF_PATH.replace(\n",
- " \"ergoCubSN000\", \"ergoCubSN001\"\n",
- " )\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "JAXsim offers a simple functional API in order to interact in a memory-efficient way with the simulation. Four main elements are used to define a simulation:\n",
- "\n",
- "- `model`: an object that defines the dynamics of the system.\n",
- "- `data`: an object that contains the state of the system.\n",
- "- `integrator`: an object that defines the integration method.\n",
- "- `integrator_state`: an object that contains the state of the integrator.\n",
- "\n",
- "The last two elements are not going to be used in this example."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "import jaxsim.api as js\n",
- "\n",
- "# Create the model from the URDF file.\n",
- "model = js.model.JaxSimModel.build_from_model_description(\n",
- " model_description=model_urdf_path\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "It is possible to reduce the model by lumping some links together in order to simplify the computation. "
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "model = js.model.reduce(\n",
- " model=model,\n",
- " considered_joints=tuple(\n",
- " j\n",
- " for j in model.joint_names()\n",
- " if \"camera\" not in j\n",
- " # Remove head and hands.\n",
- " and \"neck\" not in j\n",
- " and \"wrist\" not in j\n",
- " and \"thumb\" not in j\n",
- " and \"index\" not in j\n",
- " and \"middle\" not in j\n",
- " and \"ring\" not in j\n",
- " and \"pinkie\" not in j\n",
- " # Remove upper body.\n",
- " and \"torso\" not in j and \"elbow\" not in j and \"shoulder\" not in j\n",
- " ),\n",
- ")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Create the data object from the reduced model.\n",
- "data = js.data.random_model_data(model=model)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "Let's now create a random `references` object. This will be useful to manage the link forces and the torques applied to the joints."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Create a random link forces matrix.\n",
- "link_forces = jax.random.uniform(\n",
- " minval=-10.0,\n",
- " maxval=10.0,\n",
- " shape=(model.number_of_links(), 6),\n",
- " key=jax.random.PRNGKey(0),\n",
- ")\n",
- "\n",
- "# Create a random joint forces references vector.\n",
- "joint_force_references = jax.random.uniform(\n",
- " minval=-10.0, maxval=10.0, shape=(model.dofs(),), key=jax.random.PRNGKey(0)\n",
- ")\n",
- "\n",
- "# Create the references object.\n",
- "references = js.references.JaxSimModelReferences.build(\n",
- " model=model,\n",
- " data=data,\n",
- " link_forces=link_forces,\n",
- " joint_force_references=joint_force_references,\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "**Note:**\n",
- "- The argument `joint_force_references` is called `joint_force_references` because the actual joint forces will take into account the effect of the joint frictions."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Forward Dynamics\n",
- "\n",
- "We can compute the forward dynamics of the model by inverting the free floating mass matrix, e.g. starting from:\n",
- "\n",
- "$$\n",
- "M(q) \\dot{\\nu} + h(q, \\nu) = B\\tau + J^T(q) f\n",
- "$$\n",
- "\n",
- "we can compute the acceleration $\\dot{\\nu}$ as:\n",
- "\n",
- "$$\n",
- "\\dot{\\nu} = M(q)^{-1} (B\\tau - h(q, \\nu) + J^T(q) f)\n",
- "$$\n",
- "\n",
- "where:\n",
- "- $M(q) \\in \\mathbb{R} ^{6+n \\times 6+n}$ is the mass matrix\n",
- "- $h(q, \\nu) \\in \\mathbb{R} ^{6+n}$ is the bias forces term\n",
- "- $B \\in \\mathbb{R}^{6+n \\times n}$ is the input matrix\n",
- "- $\\tau \\in \\mathbb{R}^n$ is the input torque\n",
- "- $J(q) \\in \\mathbb{R}^{6+n \\times n \\times 6}$ is the Jacobian of the contact points\n",
- "- $f \\in \\mathbb{R}^{6+n}$ is the contact force.\n",
- "\n",
- "Let's now compute the necessary quantities for the forward dynamics of ErgoCub model."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Compute the free floating mass matrix.\n",
- "M = js.model.free_floating_mass_matrix(model=model, data=data)\n",
- "\n",
- "# Compute the generalized free floating Jacobian matrix.\n",
- "J = js.model.generalized_free_floating_jacobian(model=model, data=data)\n",
- "\n",
- "# Compute the free floating bias forces vector.\n",
- "h = js.model.free_floating_bias_forces(model=model, data=data)\n",
- "\n",
- "# Extract the link forces from the reference object.\n",
- "f_ext = references.link_forces(model=model, data=data)\n",
- "\n",
- "# Extract the random joint forces from the reference object.\n",
- "τ = references.joint_force_references(model=model)\n",
- "\n",
- "# Create a dummy selection matrix.\n",
- "B = jnp.block([jnp.zeros(shape=(model.dofs(), 6)), jnp.eye(model.dofs())]).T"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "We can now finally compute the acceleration of the robot:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "ν̇ = jnp.linalg.inv(M) @ (B @ τ - h + jnp.einsum(\"l6g,l6->g\", J, f_ext))\n",
- "\n",
- "print(f\"Free floating acceleration: \\n{ν̇ }\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "We can obtain the same result by using the `jaxsim.api.model` module:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "ν̇ = jnp.concatenate(\n",
- " js.model.forward_dynamics(\n",
- " model=model,\n",
- " data=data,\n",
- " link_forces=references.link_forces(model=model, data=data),\n",
- " joint_forces=references.joint_force_references(model=model),\n",
- " )\n",
- ")\n",
- "\n",
- "print(f\"Free floating acceleration: \\n{ν̇ }\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "**Note:** \n",
- "- `js.model.forward_dynamics` returns a tuple of the base acceleration and the joint accelerations, that's why we use `jnp.concatenate` to concatenate them.\n",
- "- We used `jnp.einsum` to compute $J ^\\top f$ instead of flattening the Jacobian and using `jnp.dot`."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "Other interesting quantities that may be computed are the standalone Coriolis matrix, the centroidal Jacobian matrix and the generalized free floating Jacobian derivative."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Compute the standalone Coriolis matrix.\n",
- "C = js.model.free_floating_coriolis_matrix(model=model, data=data)\n",
- "\n",
- "# Compute the average centroidal velocity Jacobian.\n",
- "G_J = js.model.average_velocity_jacobian(model=model, data=data)\n",
- "\n",
- "# Compute the generalized free floating Jacobian derivative.\n",
- "J̇ = js.model.generalized_free_floating_jacobian_derivative(model=model, data=data)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "**Note:**\n",
- "- These last methods are particularly computationally expensive and should be used with caution."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Forward Kinematics"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "It is possible to compute the *forward kinematics* of the robot and then extract a link's transform by using its index:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Find the link index\n",
- "foot_index = js.link.name_to_idx(model=model, link_name=\"l_ankle_2\")\n",
- "\n",
- "# Extract the link's transform\n",
- "H_L = js.model.forward_kinematics(model=model, data=data)[foot_index]\n",
- "\n",
- "print(f\"Link transform: \\n{H_L}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "or by directly using the `jaxsim.api.link` module:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "H_L = js.link.transform(\n",
- " model=model,\n",
- " data=data,\n",
- " link_index=foot_index,\n",
- ")\n",
- "\n",
- "print(f\"Link transform: \\n{H_L}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "### Computing the link's velocity\n",
- "\n",
- "It is possible to compute the link's velocity by leveraging:\n",
- "\n",
- "$$\n",
- "v _L = J_L \\nu\n",
- "$$\n",
- " \n",
- "where $J_L$ is the Jacobian of the link and $\\nu$ is the generalized velocity.\n",
- "\n",
- "For the link's Jacobian, we can either extract it from the generalized Jacobian as:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Extract the link's Jacobian from the generalized Jacobian\n",
- "J_L = J[foot_index]\n",
- "\n",
- "print(f\"Link Jacobian: \\n{J_L}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "or by using the `jaxsim.api.link` module:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Compute the link's Jacobian by using the `js.link` module\n",
- "J_L = js.link.jacobian(model=model, data=data, link_index=foot_index)\n",
- "\n",
- "print(f\"Link Jacobian: \\n{J_L}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "We can then compute the link's velocity as:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Extract the floating-base velocity from the data object\n",
- "ν = data.generalized_velocity()\n",
- "\n",
- "# Compute the link's velocity from the generalized velocity\n",
- "v_L = J_L @ ν\n",
- "\n",
- "print(f\"Link Velocity: \\n{v_L}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "or by directly using the `jaxsim.api.link`:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Compute the link's velocity by using the `js.link` module\n",
- "v_L = js.link.velocity(model=model, data=data, link_index=foot_index)\n",
- "\n",
- "print(f\"Link Velocity: \\n{v_L}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Contact Dynamics\n",
- "\n",
- "### Contact Points Acceleration\n",
- "\n",
- "The contact point acceleration can be calculated as the time derivative of their velocity:\n",
- "\n",
- "$$\n",
- "\\frac{\\partial}{\\partial t} \\dot{p} = \\frac{\\partial}{\\partial t} \\left( J _C \\nu \\right) = \\frac{\\partial J _C}{\\partial t} \\nu + J _C \\frac{\\partial \\nu}{\\partial t}\n",
- "$$\n",
- "\n",
- "which when brings to the following formulation:\n",
- "\n",
- "$$\n",
- "\\ddot{p} = \\dot{J _C}\\nu + J _C \\dot{\\nu}\n",
- "$$\n",
- "\n",
- "The required quantities can be easily computed using the `jaxsim.api.contact` module."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Compute the Jacobian of the contact points\n",
- "J_C = js.contact.jacobian(model=model, data=data)\n",
- "\n",
- "# Compute the Jacobian derivative of the contact points\n",
- "J̇_C = js.contact.jacobian_derivative(model=model, data=data)\n",
- "\n",
- "# Extract the floating-base velocity from the data object\n",
- "ν = data.generalized_velocity()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "We already computed $\\nu$ and $\\dot{\\nu}$ from the previous example, so we can now compute the contact points' acceleration:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "a = J̇_C @ ν + J_C @ ν̇\n",
- "\n",
- "print(f\"Contact points acceleration: \\n{a}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "Other useful quantities that can be computed by using the `jaxsim.api.contact` module."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Compute the contact forces.\n",
- "F_C, _ = js.contact.collidable_point_dynamics(\n",
- " model=model,\n",
- " data=data,\n",
- " link_forces=references.link_forces(model=model, data=data),\n",
- " joint_force_references=references.joint_force_references(model=model),\n",
- ")\n",
- "\n",
- "# Compute the contact points positions.\n",
- "p_C = js.contact.collidable_point_positions(model=model, data=data)\n",
- "\n",
- "# Compute the contact points velocities.\n",
- "v_C = js.contact.collidable_point_velocities(model=model, data=data)"
- ]
- }
- ],
- "metadata": {
- "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.12.7"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat": 4,
+ "nbformat_minor": 0
}
|