Finalize minor changes for v0.4 release

examples/PD_controller.ipynb

@@ -20,7 +20,7 @@
    "outputs": [],
    "source": [
     "# @title Imports and setup\n",
-    "from IPython.display import clear_output, HTML, display\n",
+    "from IPython.display import clear_output\n",
     "import sys\n",
     "\n",
     "IS_COLAB = \"google.colab\" in sys.modules\n",
@@ -100,6 +100,8 @@
     "from jaxsim import integrators\n",
     "\n",
     "dt = 0.01\n",
+    "integration_time = 5.0\n",
+    "num_steps = int(integration_time / dt)\n",
     "\n",
     "model = js.model.JaxSimModel.build_from_model_description(\n",
     "    model_description=model_urdf_string, is_urdf=True\n",
@@ -150,7 +152,7 @@
    "source": [
     "# @title Set up MuJoCo renderer\n",
     "\n",
-    "from jaxsim.mujoco import RodModelToMjcf, MujocoModelHelper, MujocoVideoRecorder\n",
+    "from jaxsim.mujoco import MujocoModelHelper, MujocoVideoRecorder\n",
     "from jaxsim.mujoco.loaders import UrdfToMjcf\n",
     "\n",
     "import os\n",
@@ -227,7 +229,7 @@
    "source": [
     "import mediapy as media\n",
     "\n",
-    "for _ in range(500):\n",
+    "for _ in range(num_steps):\n",
     "    data, integrator_state = js.model.step(\n",
     "        dt=dt,\n",
     "        model=model,\n",
@@ -299,7 +301,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "for _ in range(500):\n",
+    "for _ in range(num_steps):\n",
     "    control_torques = pd_controller(\n",
     "        data=data,\n",
     "        q_d=jnp.array([0.0, 0.0]),\n",

examples/Parallel_computing.ipynb

@@ -23,8 +23,6 @@
     "# @title Imports and setup\n",
     "import sys\n",
     "\n",
-    "from IPython.display import HTML, clear_output, display\n",
-    "\n",
     "IS_COLAB = \"google.colab\" in sys.modules\n",
     "\n",
     "# Install JAX and Gazebo\n",
@@ -40,15 +38,12 @@
     "%env XLA_PYTHON_CLIENT_MEM_PREALLOCATE=false\n",
     "\n",
     "import time\n",
-    "from typing import Dict, Tuple\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "import jax_dataclasses\n",
     "import rod\n",
     "from rod.builder.primitives import SphereBuilder\n",
     "\n",
-    "import jaxsim.typing as jtp\n",
     "from jaxsim import logging\n",
     "\n",
     "logging.set_logging_level(logging.LoggingLevel.INFO)\n",
@@ -105,7 +100,7 @@
     "from jaxsim import integrators\n",
     "\n",
     "dt = 0.001\n",
-    "integration_time = 1500\n",
+    "integration_time = 1.5  # seconds\n",
     "\n",
     "model = js.model.JaxSimModel.build_from_model_description(\n",
     "    model_description=model_sdf_string\n",
@@ -129,7 +124,9 @@
     "\n",
     "By default, in JaxSim a sphere primitive has 250 collision points. This can be modified by setting the `JAXSIM_COLLISION_SPHERE_POINTS` environment variable.\n",
     "\n",
-    "Given that at its steady-state the sphere will act on two or three points, we can estimate the ground parameters by explicitly setting the number of active points to these values."
+    "Given that at its steady-state the sphere will act on two or three points, we can estimate the ground parameters by explicitly setting the number of active points to these values. \n",
+    "\n",
+    "Eventually, you can specify the maximum penetration depth of the sphere into the terrain by passing `max_penetraion` to the `estimate_good_soft_contacts_parameters` function."
    ]
   },
   {
@@ -139,8 +136,10 @@
    "outputs": [],
    "source": [
     "data = data.replace(\n",
-    "    soft_contacts_params=js.contact.estimate_good_soft_contacts_parameters(\n",
-    "        model, number_of_active_collidable_points_steady_state=3\n",
+    "    contacts_params=js.contact.estimate_good_soft_contacts_parameters(\n",
+    "        model=model,\n",
+    "        number_of_active_collidable_points_steady_state=3,\n",
+    "        max_penetration=None,\n",
     "    )\n",
     ")"
    ]
@@ -149,7 +148,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Let's create a position vector for a 3x3 grid. Every sphere will be placed at a different height."
+    "Let's create a position vector for a 4x4 grid. Every sphere will be placed at a different height."
    ]
   },
   {
@@ -169,14 +168,13 @@
     "def grid(edge_len, envs_per_row):\n",
     "    edge = jnp.linspace(-edge_len, edge_len, envs_per_row)\n",
     "    xx, yy = jnp.meshgrid(edge, edge)\n",
-    "\n",
-    "    poses = [\n",
-    "        [[xx[i, j], yy[i, j], 0.2 + 0.1 * (i * envs_per_row + j)], [0, 0, 0]]\n",
-    "        for i in range(xx.shape[0])\n",
-    "        for j in range(yy.shape[0])\n",
-    "    ]\n",
-    "\n",
-    "    return jnp.array(poses)\n",
+    "    zz = 0.2 + 0.1 * (\n",
+    "        jnp.arange(envs_per_row**2) % envs_per_row\n",
+    "        + jnp.arange(envs_per_row**2) // envs_per_row\n",
+    "    )\n",
+    "    zz = zz.reshape(envs_per_row, envs_per_row)\n",
+    "    poses = jnp.stack([xx, yy, zz], axis=-1).reshape(envs_per_row**2, 3)\n",
+    "    return poses\n",
     "\n",
     "\n",
     "logging.info(f\"Simulating {envs_per_row**2} environments\")\n",
@@ -200,11 +198,13 @@
     "def simulate(\n",
     "    data: js.data.JaxSimModelData, integrator_state: dict, pose: jnp.array\n",
     ") -> tuple:\n",
-    "\n",
+    "    # Set the base position to the initial pose\n",
     "    data = data.reset_base_position(base_position=pose)\n",
+    "\n",
+    "    # Create a list to store the base position over time\n",
     "    x_t_i = []\n",
     "\n",
-    "    for _ in range(integration_time):\n",
+    "    for _ in range(int(integration_time // dt)):\n",
     "        data, integrator_state = js.model.step(\n",
     "            dt=dt,\n",
     "            model=model,\n",
@@ -243,15 +243,15 @@
     "# Run and time the simulation\n",
     "now = time.perf_counter()\n",
     "\n",
-    "x_t = simulate_vectorized(data, integrator_state, poses[:, 0])\n",
+    "x_t = simulate_vectorized(data, integrator_state, poses)\n",
     "\n",
     "comp_time = time.perf_counter() - now\n",
     "\n",
     "logging.info(\n",
     "    f\"Running simulation with {envs_per_row**2} models took {comp_time} seconds.\"\n",
     ")\n",
     "logging.info(\n",
-    "    f\"This corresponds to an RTF (Real Time Factor) of {(envs_per_row**2 *integration_time/comp_time):.2f}\"\n",
+    "    f\"This corresponds to an RTF (Real Time Factor) of {(envs_per_row**2 * integration_time / comp_time):.2f}\"\n",
     ")"
    ]
   },

pyproject.toml

@@ -41,25 +41,26 @@ preview = true
 
 [tool.ruff.lint]
 # https://docs.astral.sh/ruff/rules/
+select = [
+  "B",
+  "E",
+  "F",
+  "I",
+  "W",
+  "RUF",
+  "YTT",
+]
+
 ignore = [
   "B008", # Function call in default argument
   "B024", # Abstract base class without abstract methods
-  "B904", # Raise without from inside exception
-  "B905", # Zip without explicit strict
   "E402", # Module level import not at top of file
   "E501", # Line too long
   "E731", # Do not assign a `lambda` expression, use a `def`
   "E741", # Ambiguous variable name
   "F841", # Local variable is assigned to but never used
   "I001", # Import block is unsorted or unformatted
-]
-select = [
-  "B",
-  "E",
-  "F",
-  "I",
-  "W",
-  "YTT",
+  "RUF003", # Ambigous unicode character in comment
 ]
 
 [tool.ruff.lint.per-file-ignores]

src/jaxsim/__init__.py

@@ -33,17 +33,17 @@ def _is_editable() -> bool:
     import pathlib
     import site
 
-    # Get the ModuleSpec of jaxsim
+    # Get the ModuleSpec of jaxsim.
     jaxsim_spec = importlib.util.find_spec(name="jaxsim")
 
     # This can be None. If it's None, assume non-editable installation.
     if jaxsim_spec.origin is None:
         return False
 
-    # Get the folder containing the jaxsim package
+    # Get the folder containing the jaxsim package.
     jaxsim_package_dir = str(pathlib.Path(jaxsim_spec.origin).parent.parent)
 
-    # The installation is editable if the package dir is not in any {site|dist}-packages
+    # The installation is editable if the package dir is not in any {site|dist}-packages.
     return jaxsim_package_dir not in site.getsitepackages()
 
 
@@ -82,10 +82,10 @@ def _get_default_logging_level(env_var: str) -> logging.LoggingLevel:
 logging.configure(level=_get_default_logging_level(env_var="JAXSIM_LOGGING_LEVEL"))
 
 
-# Configure JAX
+# Configure JAX.
 _jnp_options()
 
-# Initialize the numpy print options
+# Initialize the numpy print options.
 _np_options()
 
 del _jnp_options

src/jaxsim/api/com.py

@@ -1,6 +1,5 @@
 import jax
 import jax.numpy as jnp
-import jaxlie
 
 import jaxsim.api as js
 import jaxsim.math
@@ -28,7 +27,7 @@ def com_position(
 
     W_H_L = js.model.forward_kinematics(model=model, data=data)
     W_H_B = data.base_transform()
-    B_H_W = jaxlie.SE3.from_matrix(W_H_B).inverse().as_matrix()
+    B_H_W = jaxsim.math.Transform.inverse(transform=W_H_B)
 
     def B_p̃_LCoM(i) -> jtp.Vector:
         m = js.link.mass(model=model, link_index=i)
@@ -179,9 +178,9 @@ def locked_centroidal_spatial_inertia(
         case _:
             raise ValueError(data.velocity_representation)
 
-    B_H_G = jaxlie.SE3.from_matrix(jaxsim.math.Transform.inverse(W_H_B) @ W_H_G)
+    B_H_G = jaxsim.math.Transform.inverse(W_H_B) @ W_H_G
 
-    B_Xv_G = B_H_G.adjoint()
+    B_Xv_G = jaxsim.math.Adjoint.from_transform(transform=B_H_G)
     G_Xf_B = B_Xv_G.transpose()
 
     return G_Xf_B @ B_Mbb_B @ B_Xv_G

src/jaxsim/api/common.py

@@ -8,10 +8,10 @@
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
+from jaxsim.math import Adjoint
 from jaxsim.utils import JaxsimDataclass, Mutability
 
 try:
@@ -59,7 +59,7 @@ def switch_velocity_representation(
 
         try:
 
-            # First, we replace the velocity representation
+            # First, we replace the velocity representation.
             with self.mutable_context(
                 mutability=Mutability.MUTABLE_NO_VALIDATION,
                 restore_after_exception=True,
@@ -97,7 +97,7 @@ def inertial_to_other_representation(
             array: The 6D quantity to convert.
             other_representation: The representation to convert to.
             transform:
-                The `math:W \mathbf{H}_O` transform, where `math:O` is the
+                The :math:`W \mathbf{H}_O` transform, where :math:`O` is the
                 reference frame of the other representation.
             is_force: Whether the quantity is a 6D force or a 6D velocity.
 
@@ -122,11 +122,11 @@ def inertial_to_other_representation(
             case VelRepr.Body:
 
                 if not is_force:
-                    O_Xv_W = jaxlie.SE3.from_matrix(W_H_O).inverse().adjoint()
+                    O_Xv_W = Adjoint.from_transform(transform=W_H_O, inverse=True)
                     O_array = O_Xv_W @ W_array
 
                 else:
-                    O_Xf_W = jaxlie.SE3.from_matrix(W_H_O).adjoint().T
+                    O_Xf_W = Adjoint.from_transform(transform=W_H_O).T
                     O_array = O_Xf_W @ W_array
 
                 return O_array
@@ -136,11 +136,11 @@ def inertial_to_other_representation(
                 W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
 
                 if not is_force:
-                    OW_Xv_W = jaxlie.SE3.from_matrix(W_H_OW).inverse().adjoint()
+                    OW_Xv_W = Adjoint.from_transform(transform=W_H_OW, inverse=True)
                     OW_array = OW_Xv_W @ W_array
 
                 else:
-                    OW_Xf_W = jaxlie.SE3.from_matrix(W_H_OW).adjoint().transpose()
+                    OW_Xf_W = Adjoint.from_transform(transform=W_H_OW).T
                     OW_array = OW_Xf_W @ W_array
 
                 return OW_array
@@ -190,11 +190,11 @@ def other_representation_to_inertial(
                 O_array = array
 
                 if not is_force:
-                    W_Xv_O: jtp.Array = jaxlie.SE3.from_matrix(W_H_O).adjoint()
+                    W_Xv_O: jtp.Array = Adjoint.from_transform(W_H_O)
                     W_array = W_Xv_O @ O_array
 
                 else:
-                    W_Xf_O = jaxlie.SE3.from_matrix(W_H_O).inverse().adjoint().T
+                    W_Xf_O = Adjoint.from_transform(transform=W_H_O, inverse=True).T
                     W_array = W_Xf_O @ O_array
 
                 return W_array
@@ -205,11 +205,11 @@ def other_representation_to_inertial(
                 W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
 
                 if not is_force:
-                    W_Xv_BW: jtp.Array = jaxlie.SE3.from_matrix(W_H_OW).adjoint()
+                    W_Xv_BW: jtp.Array = Adjoint.from_transform(W_H_OW)
                     W_array = W_Xv_BW @ BW_array
 
                 else:
-                    W_Xf_BW = jaxlie.SE3.from_matrix(W_H_OW).inverse().adjoint().T
+                    W_Xf_BW = Adjoint.from_transform(transform=W_H_OW, inverse=True).T
                     W_array = W_Xf_BW @ BW_array
 
                 return W_array

src/jaxsim/api/contact.py

@@ -8,7 +8,7 @@
 import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
-from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsParams
+from jaxsim.rbda.contacts.soft import SoftContactsParams
 
 from .common import VelRepr
 
@@ -137,9 +137,17 @@ def collidable_point_dynamics(
     # all collidable points belonging to the robot.
     W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
 
+    # Import privately the soft contacts classes.
+    from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
+
     # Build the soft contact model.
     match model.contact_model:
-        case s if isinstance(s, SoftContacts):
+
+        case SoftContacts():
+
+            assert isinstance(model.contact_model, SoftContacts)
+            assert isinstance(data.state.contact, SoftContactsState)
+
             # Build the contact model.
             soft_contacts = SoftContacts(
                 parameters=data.contacts_params, terrain=model.terrain
@@ -337,7 +345,7 @@ def jacobian(
             The output velocity representation of the free-floating jacobian.
 
     Returns:
-        The stacked 6×(6+n) free-floating jacobians of the frames associated to the
+        The stacked :math:`6 \times (6+n)` free-floating jacobians of the frames associated to the
         collidable points.
 
     Note: