jaxsim 0.2.1.dev123__tar.gz → 0.2.1.dev155__tar.gz

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/.github/workflows/ci_cd.yml

@@ -118,8 +118,7 @@ jobs:
               - *src
               - *tests
 
-      # https://gazebosim.org/docs/harmonic/install_ubuntu
-      - name: Install Gazebo Sim
+      - name: Install 'gz sdf' system command
         if: |
           contains(matrix.os, 'ubuntu') &&
           (github.event_name != 'pull_request' ||
@@ -130,7 +129,7 @@ jobs:
           sudo wget https://packages.osrfoundation.org/gazebo.gpg -O /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 gz-harmonic
+          sudo apt-get install --no-install-recommends libsdformat13 gz-tools2
 
       - name: Run the Python tests
         if: |

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/.gitignore

@@ -140,5 +140,10 @@ src/jaxsim/_version.py
 
 # ruff
 .ruff_cache/
+
 # pixi environments
 .pixi
+
+# data
+.mp4
+.png

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev123
+Version: 0.2.1.dev155
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it
@@ -68,7 +68,7 @@ Requires-Dist: mujoco>=3.0.0; extra == "all"
 
 JaxSim is a **differentiable physics engine** and **multibody dynamics library** designed for applications in control and robot learning, implemented with JAX.
 
-Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence. 
+Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence.
 
 ## Features
 
@@ -91,7 +91,7 @@ Its design facilitates research and accelerates prototyping in the intersection
 
 ### JaxSim as a multibody dynamics library
 
-- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians. 
+- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians.
 - Provides all the quantities included in the Euler-Poincarè formulation of the equations of motion.
 - Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].
 - Exposes all the necessary quantities to develop controllers in centroidal coordinates.
@@ -198,10 +198,10 @@ The main differences between MJX/Brax and JaxSim are as follows:
 
 - JaxSim supports out-of-the-box all SDF models with [Pose Frame Semantics][PFS].
 - JaxSim only supports collisions between points rigidly attached to bodies and a compliant ground surface.
-  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous  
+  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous
   function of the state $(\mathbf{q}, \boldsymbol{\nu})$, it doesn't require running any optimization algorithm
   when stepping the simulation forward.
-- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators. 
+- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators.
 
 [brax]: https://github.com/google/brax
 [mjx]: https://mujoco.readthedocs.io/en/3.0.0/mjx.html

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/README.md

@@ -2,7 +2,7 @@
 
 JaxSim is a **differentiable physics engine** and **multibody dynamics library** designed for applications in control and robot learning, implemented with JAX.
 
-Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence. 
+Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence.
 
 ## Features
 
@@ -25,7 +25,7 @@ Its design facilitates research and accelerates prototyping in the intersection
 
 ### JaxSim as a multibody dynamics library
 
-- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians. 
+- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians.
 - Provides all the quantities included in the Euler-Poincarè formulation of the equations of motion.
 - Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].
 - Exposes all the necessary quantities to develop controllers in centroidal coordinates.
@@ -132,10 +132,10 @@ The main differences between MJX/Brax and JaxSim are as follows:
 
 - JaxSim supports out-of-the-box all SDF models with [Pose Frame Semantics][PFS].
 - JaxSim only supports collisions between points rigidly attached to bodies and a compliant ground surface.
-  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous  
+  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous
   function of the state $(\mathbf{q}, \boldsymbol{\nu})$, it doesn't require running any optimization algorithm
   when stepping the simulation forward.
-- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators. 
+- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators.
 
 [brax]: https://github.com/google/brax
 [mjx]: https://mujoco.readthedocs.io/en/3.0.0/mjx.html

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155/examples}/.gitattributes

@@ -1,3 +1,2 @@
 # GitHub syntax highlighting
 pixi.lock linguist-language=YAML
-

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/__init__.py

@@ -8,8 +8,9 @@ def _jnp_options() -> None:
 
     import jax
 
-    # Enable by default
-    if not ("JAX_ENABLE_X64" in os.environ and os.environ["JAX_ENABLE_X64"] == "0"):
+    # Enable by default 64bit precision in JAX.
+    if os.environ.get("JAX_ENABLE_X64", "1") != "0":
+
         logging.info("Enabling JAX to use 64bit precision")
         jax.config.update("jax_enable_x64", True)
 
@@ -27,6 +28,7 @@ def _np_options() -> None:
 
 
 def _is_editable() -> bool:
+
     import importlib.util
     import pathlib
     import site
@@ -45,11 +47,40 @@ def _is_editable() -> bool:
     return jaxsim_package_dir not in site.getsitepackages()
 
 
-# Initialize the logging verbosity
-if _is_editable():
-    logging.configure(level=logging.LoggingLevel.DEBUG)
-else:
-    logging.configure(level=logging.LoggingLevel.WARNING)
+def _get_default_logging_level(env_var: str) -> logging.LoggingLevel:
+    """
+    Get the default logging level.
+
+    Args:
+        env_var: The environment variable to check.
+
+    Returns:
+        The logging level to set.
+    """
+
+    import os
+
+    # Define the default logging level depending on the installation mode.
+    default_logging_level = (
+        logging.LoggingLevel.DEBUG
+        if _is_editable()  # noqa: F821
+        else logging.LoggingLevel.WARNING
+    )
+
+    # Allow to override the default logging level with an environment variable.
+    try:
+        return logging.LoggingLevel[
+            os.environ.get(env_var, default_logging_level.name).upper()
+        ]
+
+    except KeyError as exc:
+        msg = f"Invalid logging level defined in {env_var}='{os.environ[env_var]}'"
+        raise RuntimeError(msg) from exc
+
+
+# Configure the logger with the default logging level.
+logging.configure(level=_get_default_logging_level(env_var="JAXSIM_LOGGING_LEVEL"))
+
 
 # Configure JAX
 _jnp_options()
@@ -59,6 +90,7 @@ _np_options()
 
 del _jnp_options
 del _np_options
+del _get_default_logging_level
 del _is_editable
 
 from . import terrain  # isort:skip

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1.dev123'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev123')
+__version__ = version = '0.2.1.dev155'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev155')

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/api/contact.py

@@ -365,20 +365,20 @@ def jacobian(
 
             W_H_C = transforms(model=model, data=data)
 
-            def jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
+            def body_jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
                 C_X_W = jaxsim.math.Adjoint.from_transform(
                     transform=W_H_C, inverse=True
                 )
                 C_J_WC = C_X_W @ W_J_WC
                 return C_J_WC
 
-            O_J_WC = jax.vmap(jacobian)(W_H_C, W_J_WC)
+            O_J_WC = jax.vmap(body_jacobian)(W_H_C, W_J_WC)
 
         case VelRepr.Mixed:
 
             W_H_C = transforms(model=model, data=data)
 
-            def jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
+            def mixed_jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
 
                 W_H_CW = W_H_C.at[0:3, 0:3].set(jnp.eye(3))
 
@@ -389,7 +389,7 @@ def jacobian(
                 CW_J_WC = CW_X_W @ W_J_WC
                 return CW_J_WC
 
-            O_J_WC = jax.vmap(jacobian)(W_H_C, W_J_WC)
+            O_J_WC = jax.vmap(mixed_jacobian)(W_H_C, W_J_WC)
 
         case _:
             raise ValueError(output_vel_repr)

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/api/kin_dyn_parameters.py

@@ -6,6 +6,7 @@ import jax.lax
 import jax.numpy as jnp
 import jax_dataclasses
 import jaxlie
+import numpy as np
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
@@ -220,7 +221,9 @@ class KynDynParameters(JaxsimDataclass):
             (
                 hash(self.number_of_links()),
                 hash(self.number_of_joints()),
-                hash(tuple(jnp.atleast_1d(self.parent_array).flatten().tolist())),
+                hash(tuple(np.atleast_1d(self.parent_array).flatten().tolist())),
+                hash(self._parent_array),
+                hash(self._support_body_array_bool),
             )
         )
 

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/api/link.py

@@ -241,8 +241,8 @@ def jacobian(
     )
 
     # Compute the actual doubly-left free-floating jacobian of the link.
-    κ = model.kin_dyn_parameters.support_body_array_bool[link_index]
-    B_J_WL_B = jnp.hstack([jnp.ones(5), κ]) * B_J_full_WX_B
+    κb = model.kin_dyn_parameters.support_body_array_bool[link_index]
+    B_J_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J_full_WX_B
 
     # Adjust the input representation such that `J_WL_I @ I_ν`.
     match data.velocity_representation:

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/api/model.py

@@ -16,7 +16,8 @@ from jax_dataclasses import Static
 import jaxsim.api as js
 import jaxsim.parsers.descriptions
 import jaxsim.typing as jtp
-from jaxsim.utils import HashlessObject, JaxsimDataclass, Mutability
+from jaxsim.math import Cross
+from jaxsim.utils import JaxsimDataclass, Mutability
 
 from .common import VelRepr
 
@@ -32,6 +33,7 @@ class JaxSimModel(JaxsimDataclass):
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
         default=jaxsim.terrain.FlatTerrain(), repr=False, compare=False, hash=False
     )
+
     kin_dyn_parameters: js.kin_dyn_parameters.KynDynParameters | None = (
         dataclasses.field(default=None, repr=False, compare=False, hash=False)
     )
@@ -40,13 +42,9 @@ class JaxSimModel(JaxsimDataclass):
         default=None, repr=False, compare=False, hash=False
     )
 
-    _description: Static[
-        HashlessObject[jaxsim.parsers.descriptions.ModelDescription | None]
-    ] = dataclasses.field(default=None, repr=False, compare=False, hash=False)
-
-    @property
-    def description(self) -> jaxsim.parsers.descriptions.ModelDescription:
-        return self._description.get()
+    description: Static[jaxsim.parsers.descriptions.ModelDescription | None] = (
+        dataclasses.field(default=None, repr=False, compare=False, hash=False)
+    )
 
     def __eq__(self, other: JaxSimModel) -> bool:
 
@@ -60,6 +58,7 @@ class JaxSimModel(JaxsimDataclass):
         return hash(
             (
                 hash(self.model_name),
+                hash(self.description),
                 hash(self.kin_dyn_parameters),
             )
         )
@@ -156,7 +155,7 @@ class JaxSimModel(JaxsimDataclass):
         # Build the model
         model = JaxSimModel(
             model_name=model_name,
-            _description=HashlessObject(obj=model_description),
+            description=model_description,
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
@@ -301,7 +300,7 @@ def reduce(
         locked_joint_positions:
             A dictionary containing the positions of the joints to be considered
             in the reduction process. The removed joints in the reduced model
-            will have their position locked to their value in this dictionary.
+            will have their position locked to their value of this dictionary.
             If a joint is not part of the dictionary, its position is set to zero.
     """
 
@@ -314,10 +313,9 @@ def reduce(
         new_joints = set(model.joint_names()) - set(locked_joint_positions)
         raise ValueError(f"Passed joints not existing in the model: {new_joints}")
 
-    # Copy the model description with a deep copy of the joints.
-    intermediate_description = dataclasses.replace(
-        model.description, joints=copy.deepcopy(model.description.joints)
-    )
+    # Operate on a deep copy of the model description in order to prevent problems
+    # when mutable attributes are updated.
+    intermediate_description = copy.deepcopy(model.description)
 
     # Update the initial position of the joints.
     # This is necessary to compute the correct pose of the link pairs connected
@@ -685,8 +683,6 @@ def forward_dynamics_aba(
         another representation C_v̇_WB expressed in a generic frame C.
         """
 
-        from jaxsim.math import Cross
-
         # In Mixed representation, we need to include a cross product in ℝ⁶.
         # In Inertial and Body representations, the cross product is always zero.
         C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
@@ -871,6 +867,126 @@ def free_floating_mass_matrix(
             raise ValueError(data.velocity_representation)
 
 
+@jax.jit
+def free_floating_coriolis_matrix(
+    model: JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    """
+    Compute the free-floating Coriolis matrix of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The free-floating Coriolis matrix of the model.
+
+    Note:
+        This function, contrarily to other quantities of the equations of motion,
+        does not exploit any iterative algorithm. Therefore, the computation of
+        the Coriolis matrix may be much slower than other quantities.
+    """
+
+    # We perform all the calculation in body-fixed.
+    # The Coriolis matrix computed in this representation is converted later
+    # to the active representation stored in data.
+    with data.switch_velocity_representation(VelRepr.Body):
+
+        B_ν = data.generalized_velocity()
+
+        # Doubly-left free-floating Jacobian.
+        L_J_WL_B = generalized_free_floating_jacobian(model=model, data=data)
+
+        # Doubly-left free-floating Jacobian derivative.
+        L_J̇_WL_B = jax.vmap(
+            lambda link_index: js.link.jacobian_derivative(
+                model=model, data=data, link_index=link_index
+            )
+        )(js.link.names_to_idxs(model=model, link_names=model.link_names()))
+
+    L_M_L = link_spatial_inertia_matrices(model=model)
+
+    # Body-fixed link velocities.
+    # Note: we could have called link.velocity() instead of computing it ourselves,
+    # but since we need the link Jacobians later, we can save a double calculation.
+    L_v_WL = jax.vmap(lambda J: J @ B_ν)(L_J_WL_B)
+
+    # Compute the contribution of each link to the Coriolis matrix.
+    def compute_link_contribution(M, v, J, J̇) -> jtp.Array:
+
+        return J.T @ ((Cross.vx_star(v) @ M + M @ Cross.vx(v)) @ J + M @ J̇)
+
+    C_B_links = jax.vmap(compute_link_contribution)(
+        L_M_L,
+        L_v_WL,
+        L_J_WL_B,
+        L_J̇_WL_B,
+    )
+
+    # We need to adjust the Coriolis matrix for fixed-base models.
+    # In this case, the base link does not contribute to the matrix, and we need to zero
+    # the off-diagonal terms mapping joint quantities onto the base configuration.
+    if model.floating_base():
+        C_B = C_B_links.sum(axis=0)
+    else:
+        C_B = C_B_links[1:].sum(axis=0)
+        C_B = C_B.at[0:6, 6:].set(0.0)
+        C_B = C_B.at[6:, 0:6].set(0.0)
+
+    # Adjust the representation of the Coriolis matrix.
+    # Refer to https://github.com/traversaro/traversaro-phd-thesis, Section 3.6.
+    match data.velocity_representation:
+
+        case VelRepr.Body:
+            return C_B
+
+        case VelRepr.Inertial:
+
+            n = model.dofs()
+            W_H_B = data.base_transform()
+            B_X_W = jaxsim.math.Adjoint.from_transform(W_H_B, inverse=True)
+            B_T_W = jax.scipy.linalg.block_diag(B_X_W, jnp.eye(n))
+
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                W_v_WB = data.base_velocity()
+                B_Ẋ_W = -B_X_W @ jaxsim.math.Cross.vx(W_v_WB)
+
+            B_Ṫ_W = jax.scipy.linalg.block_diag(B_Ẋ_W, jnp.zeros(shape=(n, n)))
+
+            with data.switch_velocity_representation(VelRepr.Body):
+                M = free_floating_mass_matrix(model=model, data=data)
+
+            C = B_T_W.T @ (M @ B_Ṫ_W + C_B @ B_T_W)
+
+            return C
+
+        case VelRepr.Mixed:
+
+            n = model.dofs()
+            BW_H_B = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
+            B_X_BW = jaxsim.math.Adjoint.from_transform(transform=BW_H_B, inverse=True)
+            B_T_BW = jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(n))
+
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                BW_v_WB = data.base_velocity()
+                BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
+
+            BW_v_BW_B = BW_v_WB - BW_v_W_BW
+            B_Ẋ_BW = -B_X_BW @ jaxsim.math.Cross.vx(BW_v_BW_B)
+
+            B_Ṫ_BW = jax.scipy.linalg.block_diag(B_Ẋ_BW, jnp.zeros(shape=(n, n)))
+
+            with data.switch_velocity_representation(VelRepr.Body):
+                M = free_floating_mass_matrix(model=model, data=data)
+
+            C = B_T_BW.T @ (M @ B_Ṫ_BW + C_B @ B_T_BW)
+
+            return C
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+
 @jax.jit
 def inverse_dynamics(
     model: JaxSimModel,
@@ -931,8 +1047,6 @@ def inverse_dynamics(
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
-        from jaxsim.math import Cross
-
         W_X_C = jaxlie.SE3.from_matrix(W_H_C).adjoint()
         C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
         C_v_WC = C_X_W @ W_v_WC
@@ -1364,12 +1478,7 @@ def link_bias_accelerations(
     # ================================================
 
     # Compute the base transform.
-    W_H_B = jaxlie.SE3.from_rotation_and_translation(
-        rotation=jaxlie.SO3.from_quaternion_xyzw(
-            xyzw=jaxsim.math.Quaternion.to_xyzw(wxyz=data.base_orientation())
-        ),
-        translation=data.base_position(),
-    ).as_matrix()
+    W_H_B = data.base_transform()
 
     def other_representation_to_inertial(
         C_v̇_WB: jtp.Vector, C_v_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WC: jtp.Vector
@@ -1410,9 +1519,12 @@ def link_bias_accelerations(
             W_H_C = W_H_BW
             with data.switch_velocity_representation(VelRepr.Mixed):
                 W_ṗ_B = data.base_velocity()[0:3]
-                W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)
+                BW_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)
+                W_X_BW = jaxsim.math.Adjoint.from_transform(transform=W_H_BW)
+                W_v_WC = W_v_W_BW = W_X_BW @ BW_v_W_BW
             with data.switch_velocity_representation(VelRepr.Mixed):
                 C_v_WB = BW_v_WB = data.base_velocity()
+
         case _:
             raise ValueError(data.velocity_representation)
 

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/api/ode.py

@@ -223,7 +223,9 @@ def system_velocity_dynamics(
 
 @jax.jit
 def system_position_dynamics(
-    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[jtp.Vector, jtp.Vector, jtp.Vector]:
     """
     Compute the dynamics of the system position.
@@ -231,6 +233,8 @@ def system_position_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
+        baumgarte_quaternion_regularization:
+            The Baumgarte regularization coefficient for adjusting the quaternion norm.
 
     Returns:
         A tuple containing the derivative of the base position, the derivative of the
@@ -250,6 +254,7 @@ def system_position_dynamics(
         quaternion=W_Q_B,
         omega=W_ω_WB,
         omega_in_body_fixed=False,
+        K=baumgarte_quaternion_regularization,
     ).squeeze()
 
     return W_ṗ_B, W_Q̇_B, ṡ
@@ -262,6 +267,7 @@ def system_dynamics(
     *,
     joint_forces: jtp.Vector | None = None,
     link_forces: jtp.Vector | None = None,
+    baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[ODEState, dict[str, Any]]:
     """
     Compute the dynamics of the system.
@@ -271,6 +277,9 @@ def system_dynamics(
         data: The data of the considered model.
         joint_forces: The joint forces to apply.
         link_forces: The 6D forces to apply to the links.
+        baumgarte_quaternion_regularization:
+            The Baumgarte regularization coefficient used to adjust the norm of the
+            quaternion (only used in integrators not operating on the SO(3) manifold).
 
     Returns:
         A tuple with an `ODEState` object storing in each of its attributes the
@@ -287,7 +296,11 @@ def system_dynamics(
     )
 
     # Extract the velocities.
-    W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(model=model, data=data)
+    W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
+        model=model,
+        data=data,
+        baumgarte_quaternion_regularization=baumgarte_quaternion_regularization,
+    )
 
     # Create an ODEState object populated with the derivative of each leaf.
     # Our integrators, operating on generic pytrees, will be able to handle it

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/math/joint_model.py

@@ -34,8 +34,8 @@ class JointModel:
         already in a vectorized form. In other words, it cannot be created using vmap.
     """
 
-    λ_H_pre: jax.Array
-    suc_H_i: jax.Array
+    λ_H_pre: jtp.Array
+    suc_H_i: jtp.Array
 
     joint_dofs: Static[tuple[int, ...]]
     joint_names: Static[tuple[str, ...]]

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/mujoco/loaders.py

@@ -188,10 +188,9 @@ class RodModelToMjcf:
         )
 
         # If considered joints are passed, make sure that they are all part of the model.
-        if considered_joints - set([j.name for j in rod_model.joints()]):
-            extra_joints = set(considered_joints) - set(
-                [j.name for j in rod_model.joints()]
-            )
+        if considered_joints - {j.name for j in rod_model.joints()}:
+            extra_joints = set(considered_joints) - {j.name for j in rod_model.joints()}
+
             msg = f"Couldn't find the following joints in the model: '{extra_joints}'"
             raise ValueError(msg)
 
@@ -352,7 +351,7 @@ class RodModelToMjcf:
         # Set alpha=0 to the color of all collision elements
         for geometry_element in mujoco_element.findall(".//geom[@rgba]"):
             if geometry_element.attrib.get("name") in collision_names:
-                r, g, b, a = geometry_element.attrib["rgba"].split(" ")
+                r, g, b, _ = geometry_element.attrib["rgba"].split(" ")
                 geometry_element.set("rgba", f"{r} {g} {b} 0")
 
         # -----------------------

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/mujoco/model.py

@@ -73,7 +73,7 @@ class MujocoModelHelper:
             new_hfield = generate_hfield(heightmap, (nrow, ncol))
             model.hfield_data = new_hfield
 
-        return MujocoModelHelper(model=model, data=mj.MjData(model))
+        return MujocoModelHelper(model=model, data=data)
 
     def time(self) -> float:
         """Return the simulation time."""

{jaxsim-0.2.1.dev123 → jaxsim-0.2.1.dev155}/src/jaxsim/mujoco/visualizer.py

@@ -173,4 +173,4 @@ class MujocoVisualizer:
         try:
             yield handle
         finally:
-            handle.close() if close_on_exit else None
+            _ = handle.close() if close_on_exit else None