Merge pull request #366 from ami-iit/sprint/dev

refactoring of  integrators, contacts models, and data caching  to improve speed and readability

.github/workflows/ci_cd.yml

@@ -16,7 +16,7 @@ jobs:
 
   package:
     name: Package the project
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-latest
 
     steps:
 
@@ -58,7 +58,7 @@ jobs:
       fail-fast: false
       matrix:
         os:
-          - ubuntu-22.04
+          - ubuntu-latest
           - macos-latest
           - windows-latest
         python:
@@ -124,11 +124,11 @@ jobs:
            steps.changes.outputs.all == 'true')
         run: |
           sudo apt-get update
-          sudo apt-get install lsb-release wget gnupg
-          sudo wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg
+          sudo apt-get install lsb-release gnupg
+          sudo curl https://packages.osrfoundation.org/gazebo.gpg --output /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg
           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
           sudo apt-get update
-          sudo apt-get install --no-install-recommends libsdformat13 gz-tools2
+          sudo apt-get install gz-ionic sdf
 
       - name: Run the Python tests
         if: |
@@ -144,7 +144,7 @@ jobs:
   publish:
     name: Publish to PyPI
     needs: test
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-latest
     permissions:
         id-token: write
 

README.md

@@ -6,9 +6,8 @@
 <br/>
 <table>
   <tr>
-    <th><img src="https://github.com/user-attachments/assets/115b1c1c-6ae5-4c59-92e0-1be13ba954db" width="250"></th>
-    <th><img src="https://github.com/user-attachments/assets/f9661fae-9a85-41dd-9a58-218758ec8c9c" width="250"></th>
-    <th><img src="https://github.com/user-attachments/assets/ae8adadf-3bca-47b8-97ca-3a9273633d60" width="250"></th>
+    <th><img src="https://github.com/user-attachments/assets/f9661fae-9a85-41dd-9a58-218758ec8c9c" width="500"></th>
+    <th><img src="https://github.com/user-attachments/assets/62b88b9d-45ea-4d22-99d2-f24fc842dd29" width="500"></th>
   </tr>
 </table>
 <br/>
@@ -51,15 +50,14 @@ 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])
-)
+# Note that the default data representation is mixed velocity representation
+data = js.data.JaxSimModelData.build(model=model,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)
+    data = js.model.step(model=model, data=data, link_forces=None, joint_force_references=tau)
 
 ```
 

docs/guide/configuration.rst

@@ -61,7 +61,7 @@ The logging and exceptions configurations is controlled by the following environ
 
   *Default:* ``DEBUG`` for development, ``WARNING`` for production.
 
-- ``JAXSIM_DISABLE_EXCEPTIONS``: Disables the runtime checks and exceptions.
+- ``JAXSIM_ENABLE_EXCEPTIONS``: Enables the runtime checks and exceptions. Note that enabling exceptions might lead to device-to-host transfer of data, increasing the computational time required.
 
   *Default:* ``False``.
 

environment.yml

@@ -11,7 +11,7 @@ dependencies:
   - jaxlib >= 0.4.26
   - jaxlie >= 1.3.0
   - jax-dataclasses >= 1.4.0
-  - optax >= 0.2.3
+  - optax == 0.2.3
   - pptree
   - qpax
   - rod >= 0.3.3

examples/jaxsim_as_multibody_dynamics_library.ipynb

@@ -261,8 +261,8 @@
     "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",
+    "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",
@@ -365,12 +365,6 @@
     "# 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 tensor 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}\")"
    ]
   },
@@ -406,7 +400,7 @@
     "with data.switch_velocity_representation(VelRepr.Body):\n",
     "\n",
     "    # Body-fixed generalized velocity.\n",
-    "    B_ν = data.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",
@@ -479,7 +473,7 @@
     "with data.switch_velocity_representation(VelRepr.Body):\n",
     "\n",
     "    # Body-fixed generalized velocity.\n",
-    "    B_ν = data.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",
@@ -782,12 +776,11 @@
     ")\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_ḣ += m * jnp.hstack([0, 0, model.gravity, 0, 0, 0])\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}\")"
    ]
@@ -839,7 +832,7 @@
     "):\n",
     "\n",
     "    # Compute the mixed generalized velocity.\n",
-    "    BW_ν = data.generalized_velocity()\n",
+    "    BW_ν = data.generalized_velocity\n",
     "\n",
     "    # Compute the mixed generalized acceleration.\n",
     "    BW_ν̇ = jnp.hstack(\n",
@@ -876,44 +869,7 @@
     "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": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Ot6HePB_twaE",
-    "outputId": "02a6abae-257e-45ee-e9de-6a607cdbeb9a"
-   },
-   "outputs": [],
-   "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}\")"
+    "print(f\"W_p̈_C: shape={W_p̈_C.shape}\")"
    ]
   },
   {
@@ -944,7 +900,7 @@
    "toc_visible": true
   },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "comodo_jaxsim",
    "language": "python",
    "name": "python3"
   },
@@ -958,7 +914,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.15"
+   "version": "3.12.8"
   }
  },
  "nbformat": 4,

examples/jaxsim_as_physics_engine.ipynb

@@ -52,6 +52,7 @@
     "# Set environment variable to avoid GPU out of memory errors\n",
     "%env XLA_PYTHON_CLIENT_MEM_PREALLOCATE=false\n",
     "\n",
+    "\n",
     "# ================\n",
     "# Notebook imports\n",
     "# ================\n",
@@ -148,9 +149,7 @@
    "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"
+    "data = js.data.JaxSimModelData.build(model=model, base_position=jnp.array([0.0, 0.0, 1.0]))"
    ]
   },
   {
@@ -175,7 +174,7 @@
     "\n",
     "# Simulate\n",
     "for _t in T:\n",
-    "    data, _ = js.model.step(\n",
+    "    data = js.model.step(\n",
     "        model=model,\n",
     "        data=data,\n",
     "        link_forces=None,\n",
@@ -234,12 +233,11 @@
     "# 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",
+    "    lambda pos:  js.data.JaxSimModelData.build(model=model, base_position=pos))(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)"
+    "    data = step_parallel(model, data)"
    ]
   }
  ],
@@ -266,7 +264,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.0"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,