jaxsim 0.4.3.dev68__py3-none-any.whl → 0.4.3.dev77__py3-none-any.whl

jaxsim/__init__.py

@@ -20,6 +20,11 @@ def _jnp_options() -> None:
         if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
             logging.warning("Failed to enable 64bit precision in JAX")
 
+    else:
+        logging.warning(
+            "Using 32bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
+        )
+
 
 def _np_options() -> None:
     import numpy as np

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev68'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev68')
+__version__ = version = '0.4.3.dev77'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev77')

jaxsim/api/contact.py

@@ -117,6 +117,7 @@ def collidable_point_dynamics(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
 ) -> tuple[jtp.Matrix, dict[str, jtp.Array]]:
     r"""
     Compute the 6D force applied to each collidable point.
@@ -127,11 +128,14 @@ def collidable_point_dynamics(
         link_forces:
             The 6D external forces to apply to the links expressed in the same
             representation of data.
+        joint_force_references:
+            The joint force references to apply to the joints.
 
     Returns:
         The 6D force applied to each collidable point and additional data based on the contact model configured:
         - Soft: the material deformation rate.
-        - Rigid: nothing.
+        - Rigid: no additional data.
+        - QuasiRigid: no additional data.
 
     Note:
         The material deformation rate is always returned in the mixed frame
@@ -144,6 +148,10 @@ def collidable_point_dynamics(
     W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
 
     # Import privately the contacts classes.
+    from jaxsim.rbda.contacts.relaxed_rigid import (
+        RelaxedRigidContacts,
+        RelaxedRigidContactsState,
+    )
     from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
     from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 
@@ -186,6 +194,29 @@ def collidable_point_dynamics(
                 model=model,
                 data=data,
                 link_forces=link_forces,
+                joint_force_references=joint_force_references,
+            )
+
+            aux_data = dict()
+
+        case RelaxedRigidContacts():
+            assert isinstance(model.contact_model, RelaxedRigidContacts)
+            assert isinstance(data.state.contact, RelaxedRigidContactsState)
+
+            # Build the contact model.
+            relaxed_rigid_contacts = RelaxedRigidContacts(
+                parameters=data.contacts_params, terrain=model.terrain
+            )
+
+            # Compute the 6D force expressed in the inertial frame and applied to each
+            # collidable point.
+            W_f_Ci, _ = relaxed_rigid_contacts.compute_contact_forces(
+                position=W_p_Ci,
+                velocity=W_ṗ_Ci,
+                model=model,
+                data=data,
+                link_forces=link_forces,
+                joint_force_references=joint_force_references,
             )
 
             aux_data = dict()

jaxsim/api/data.py

@@ -6,10 +6,11 @@ from collections.abc import Sequence
 
 import jax
 import jax.numpy as jnp
+import jax.scipy.spatial.transform
 import jax_dataclasses
-import jaxlie
 
 import jaxsim.api as js
+import jaxsim.math
 import jaxsim.rbda
 import jaxsim.typing as jtp
 from jaxsim.rbda.contacts.soft import SoftContacts
@@ -39,7 +40,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     contacts_params: jaxsim.rbda.ContactsParams = dataclasses.field(repr=False)
 
     time_ns: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(0, dtype=jnp.uint64)
+        default_factory=lambda: jnp.array(
+            0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
+        ),
     )
 
     def __hash__(self) -> int:
@@ -172,9 +175,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         )
 
         time_ns = (
-            jnp.array(time * 1e9, dtype=jnp.uint64)
+            jnp.array(
+                time * 1e9,
+                dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32,
+            )
             if time is not None
-            else jnp.array(0, dtype=jnp.uint64)
+            else jnp.array(
+                0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
+            )
         )
 
         if isinstance(model.contact_model, SoftContacts):
@@ -188,10 +196,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         else:
             contacts_params = model.contact_model.parameters
 
-        W_H_B = jaxlie.SE3.from_rotation_and_translation(
-            translation=base_position,
-            rotation=jaxlie.SO3(wxyz=base_quaternion),
-        ).as_matrix()
+        W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
+            translation=base_position, quaternion=base_quaternion
+        )
 
         v_WB = JaxSimModelData.other_representation_to_inertial(
             array=jnp.hstack([base_linear_velocity, base_angular_velocity]),
@@ -377,7 +384,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             on_false=W_Q_B / jnp.linalg.norm(W_Q_B),
         )
 
-        return (W_Q_B if not dcm else jaxlie.SO3(wxyz=W_Q_B).as_matrix()).astype(float)
+        return (W_Q_B if not dcm else jaxsim.math.Quaternion.to_dcm(W_Q_B)).astype(
+            float
+        )
 
     @jax.jit
     def base_transform(self) -> jtp.Matrix:
@@ -586,16 +595,18 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             The updated `JaxSimModelData` object.
         """
 
-        base_quaternion = jnp.array(base_quaternion)
+        W_Q_B = jnp.array(base_quaternion, dtype=float)
+
+        W_Q_B = jax.lax.select(
+            pred=jnp.allclose(jnp.linalg.norm(W_Q_B), 1.0, atol=1e-6, rtol=0.0),
+            on_true=W_Q_B,
+            on_false=W_Q_B / jnp.linalg.norm(W_Q_B),
+        )
 
         return self.replace(
             validate=True,
             state=self.state.replace(
-                physics_model=self.state.physics_model.replace(
-                    base_quaternion=jnp.atleast_1d(base_quaternion.squeeze()).astype(
-                        float
-                    )
-                )
+                physics_model=self.state.physics_model.replace(base_quaternion=W_Q_B)
             ),
         )
 
@@ -728,6 +739,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         )
 
 
+@functools.partial(jax.jit, static_argnames=["velocity_representation", "base_rpy_seq"])
 def random_model_data(
     model: js.model.JaxSimModel,
     *,
@@ -737,6 +749,18 @@ def random_model_data(
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
     ] = ((-1, -1, 0.5), 1.0),
+    base_rpy_bounds: tuple[
+        jtp.FloatLike | Sequence[jtp.FloatLike],
+        jtp.FloatLike | Sequence[jtp.FloatLike],
+    ] = (-jnp.pi, jnp.pi),
+    base_rpy_seq: str = "XYZ",
+    joint_pos_bounds: (
+        tuple[
+            jtp.FloatLike | Sequence[jtp.FloatLike],
+            jtp.FloatLike | Sequence[jtp.FloatLike],
+        ]
+        | None
+    ) = None,
     base_vel_lin_bounds: tuple[
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
@@ -762,6 +786,12 @@ def random_model_data(
         key: The random key.
         velocity_representation: The velocity representation to use.
         base_pos_bounds: The bounds for the base position.
+        base_rpy_bounds:
+            The bounds for the euler angles used to build the base orientation.
+        base_rpy_seq:
+            The sequence of axes for rotation (using `Rotation` from scipy).
+        joint_pos_bounds:
+            The bounds for the joint positions (reading the joint limits if None).
         base_vel_lin_bounds: The bounds for the base linear velocity.
         base_vel_ang_bounds: The bounds for the base angular velocity.
         joint_vel_bounds: The bounds for the joint velocities.
@@ -776,6 +806,8 @@ def random_model_data(
 
     p_min = jnp.array(base_pos_bounds[0], dtype=float)
     p_max = jnp.array(base_pos_bounds[1], dtype=float)
+    rpy_min = jnp.array(base_rpy_bounds[0], dtype=float)
+    rpy_max = jnp.array(base_rpy_bounds[1], dtype=float)
     v_min = jnp.array(base_vel_lin_bounds[0], dtype=float)
     v_max = jnp.array(base_vel_lin_bounds[1], dtype=float)
     ω_min = jnp.array(base_vel_ang_bounds[0], dtype=float)
@@ -801,13 +833,29 @@ def random_model_data(
             key=k1, shape=(3,), minval=p_min, maxval=p_max
         )
 
-        physics_model_state.base_quaternion = jaxlie.SO3.from_rpy_radians(
-            *jax.random.uniform(key=k2, shape=(3,), minval=0, maxval=2 * jnp.pi)
-        ).wxyz
+        physics_model_state.base_quaternion = jaxsim.math.Quaternion.to_wxyz(
+            xyzw=jax.scipy.spatial.transform.Rotation.from_euler(
+                seq=base_rpy_seq,
+                angles=jax.random.uniform(
+                    key=k2, shape=(3,), minval=rpy_min, maxval=rpy_max
+                ),
+            ).as_quat()
+        )
 
         if model.number_of_joints() > 0:
-            physics_model_state.joint_positions = js.joint.random_joint_positions(
-                model=model, key=k3
+
+            s_min, s_max = (
+                jnp.array(joint_pos_bounds, dtype=float)
+                if joint_pos_bounds is not None
+                else (None, None)
+            )
+
+            physics_model_state.joint_positions = (
+                js.joint.random_joint_positions(model=model, key=k3)
+                if (s_min is None or s_max is None)
+                else jax.random.uniform(
+                    key=k3, shape=(model.dofs(),), minval=s_min, maxval=s_max
+                )
             )
 
             physics_model_state.joint_velocities = jax.random.uniform(

jaxsim/api/joint.py

@@ -180,17 +180,77 @@ def random_joint_positions(
 
     Args:
         model: The model to consider.
-        joint_names: The names of the joints.
-        key: The random key.
+        joint_names: The names of the considered joints (all if None).
+        key: The random key (initialized from seed 0 if None).
+
+    Note:
+        If the joint range or revolute joints is larger than 2π, their joint positions
+        will be sampled from an interval of size 2π.
 
     Returns:
         The random joint positions.
     """
 
+    # Consider the key corresponding to a zero seed if it was not passed.
     key = key if key is not None else jax.random.PRNGKey(seed=0)
 
+    # Get the joint limits parsed from the model description.
     s_min, s_max = position_limits(model=model, joint_names=joint_names)
 
+    # Get the joint indices.
+    # Note that it will trigger an exception if the given `joint_names` are not valid.
+    joint_names = joint_names if joint_names is not None else model.joint_names()
+    joint_indices = names_to_idxs(model=model, joint_names=joint_names)
+
+    from jaxsim.parsers.descriptions.joint import JointType
+
+    # Filter for revolute joints.
+    is_revolute = jnp.where(
+        jnp.array(model.kin_dyn_parameters.joint_model.joint_types[1:])[joint_indices]
+        == JointType.Revolute,
+        True,
+        False,
+    )
+
+    # Shorthand for π.
+    π = jnp.pi
+
+    # Filter for revolute with full range (or continuous).
+    is_revolute_full_range = jnp.logical_and(is_revolute, s_max - s_min >= 2 * π)
+
+    # Clip the lower limit to -π if the joint range is larger than [-π, π].
+    s_min = jnp.where(
+        jnp.logical_and(
+            is_revolute_full_range, jnp.logical_and(s_min <= -π, s_max >= π)
+        ),
+        -π,
+        s_min,
+    )
+
+    # Clip the upper limit to +π if the joint range is larger than [-π, π].
+    s_max = jnp.where(
+        jnp.logical_and(
+            is_revolute_full_range, jnp.logical_and(s_min <= -π, s_max >= π)
+        ),
+        π,
+        s_max,
+    )
+
+    # Shift the lower limit if the upper limit is smaller than +π.
+    s_min = jnp.where(
+        jnp.logical_and(is_revolute_full_range, s_max < π),
+        s_max - 2 * π,
+        s_min,
+    )
+
+    # Shift the upper limit if the lower limit is larger than -π.
+    s_max = jnp.where(
+        jnp.logical_and(is_revolute_full_range, s_min > -π),
+        s_min + 2 * π,
+        s_max,
+    )
+
+    # Sample the joint positions.
     s_random = jax.random.uniform(
         minval=s_min,
         maxval=s_max,

jaxsim/api/model.py

@@ -1747,14 +1747,18 @@ def link_contact_forces(
         data: The data of the considered model.
 
     Returns:
-        A (nL, 6) array containing the stacked 6D contact forces of the links,
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
         expressed in the frame corresponding to the active representation.
     """
 
+    # Note: the following code should be kept in sync with the function
+    # `jaxsim.api.ode.system_velocity_dynamics`. We cannot merge them since
+    # there we need to get also aux_data.
+
     # Compute the 6D forces applied to each collidable point expressed in the
     # inertial frame.
     with data.switch_velocity_representation(VelRepr.Inertial):
-        W_f_Ci = js.contact.collidable_point_forces(model=model, data=data)
+        W_f_C = js.contact.collidable_point_forces(model=model, data=data)
 
     # Construct the vector defining the parent link index of each collidable point.
     # We use this vector to sum the 6D forces of all collidable points rigidly
@@ -1763,29 +1767,28 @@ def link_contact_forces(
         model.kin_dyn_parameters.contact_parameters.body, dtype=int
     )
 
+    # Create the mask that associate each collidable point to their parent link.
+    # We use this mask to sum the collidable points to the right link.
+    mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+        model.number_of_links()
+    )
+
     # Sum the forces of all collidable points rigidly attached to a body.
-    # Since the contact forces W_f_Ci are expressed in the world frame,
+    # Since the contact forces W_f_C are expressed in the world frame,
     # we don't need any coordinate transformation.
-    W_f_Li = jax.vmap(
-        lambda nc: (
-            jnp.vstack(
-                jnp.equal(parent_link_index_of_collidable_points, nc).astype(int)
-            )
-            * W_f_Ci
-        ).sum(axis=0)
-    )(jnp.arange(model.number_of_links()))
-
-    # Convert the 6D forces to the active representation.
-    f_Li = jax.vmap(
-        lambda W_f_L: data.inertial_to_other_representation(
-            array=W_f_L,
-            other_representation=data.velocity_representation,
-            transform=data.base_transform(),
-            is_force=True,
-        )
-    )(W_f_Li)
+    W_f_L = mask.T @ W_f_C
 
-    return f_Li
+    # Create a references object to store the link forces.
+    references = js.references.JaxSimModelReferences.build(
+        model=model, link_forces=W_f_L, velocity_representation=VelRepr.Inertial
+    )
+
+    # Use the references object to convert the link forces to the velocity
+    # representation of data.
+    with references.switch_velocity_representation(data.velocity_representation):
+        f_L = references.link_forces(model=model, data=data)
+
+    return f_L
 
 
 # ======
@@ -1931,11 +1934,22 @@ def step(
         ),
     )
 
+    tf_ns = t0_ns + jnp.array(dt * 1e9, dtype=t0_ns.dtype)
+    tf_ns = jnp.where(tf_ns >= t0_ns, tf_ns, jnp.array(0, dtype=t0_ns.dtype))
+
+    jax.lax.cond(
+        pred=tf_ns < t0_ns,
+        true_fun=lambda: jax.debug.print(
+            "The simulation time overflowed, resetting simulation time to 0."
+        ),
+        false_fun=lambda: None,
+    )
+
     data_tf = (
         # Store the new state of the model and the new time.
         data.replace(
             state=state_tf,
-            time_ns=t0_ns + jnp.array(dt * 1e9).astype(jnp.uint64),
+            time_ns=tf_ns,
         )
     )
 

jaxsim/api/ode.py

@@ -95,7 +95,7 @@ def system_velocity_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint forces to apply.
+        joint_forces: The joint force references to apply.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
@@ -120,6 +120,7 @@ def system_velocity_dynamics(
     references = js.references.JaxSimModelReferences.build(
         model=model,
         link_forces=O_f_L,
+        joint_force_references=joint_forces,
         data=data,
         velocity_representation=data.velocity_representation,
     )
@@ -132,9 +133,16 @@ def system_velocity_dynamics(
     # with the terrain.
     W_f_Li_terrain = jnp.zeros_like(O_f_L).astype(float)
 
+    # Initialize a dictionary of auxiliary data.
+    # This dictionary is used to store additional data computed by the contact model.
     aux_data = {}
+
     if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
 
+        # Note: the following code should be kept in sync with the function
+        # `jaxsim.api.model.link_contact_forces`. We cannot merge them since
+        # here we need to get also aux_data.
+
         # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
         # along with contact-specific auxiliary states.
         with data.switch_velocity_representation(VelRepr.Inertial):
@@ -142,6 +150,7 @@ def system_velocity_dynamics(
                 model=model,
                 data=data,
                 link_forces=references.link_forces(model=model, data=data),
+                joint_force_references=references.joint_force_references(model=model),
             )
 
         # Construct the vector defining the parent link index of each collidable point.
@@ -175,17 +184,15 @@ def system_velocity_dynamics(
             forces=W_f_Li_terrain,
             additive=True,
         )
-    # Get the link forces in the data representation
-    with references.switch_velocity_representation(data.velocity_representation):
+
+        # Get the link forces in inertial representation
         f_L_total = references.link_forces(model=model, data=data)
 
-    # The following method always returns the inertial-fixed acceleration, and expects
-    # the link_forces expressed in the inertial frame.
-    W_v̇_WB, s̈ = system_acceleration(
-        model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
-    )
+        v̇_WB, s̈ = system_acceleration(
+            model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
+        )
 
-    return W_v̇_WB, s̈, aux_data
+    return v̇_WB, s̈, aux_data
 
 
 def system_acceleration(
@@ -196,7 +203,7 @@ def system_acceleration(
     link_forces: jtp.MatrixLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
-    Compute the system acceleration in inertial-fixed representation.
+    Compute the system acceleration in the active representation.
 
     Args:
         model: The model to consider.
@@ -206,7 +213,7 @@ def system_acceleration(
             The 6D forces to apply to the links expressed in the same representation of data.
 
     Returns:
-        A tuple containing the base 6D acceleration in inertial-fixed representation
+        A tuple containing the base 6D acceleration in in the active representation
         and the joint accelerations.
     """
 
@@ -272,18 +279,15 @@ def system_acceleration(
     )
 
     # - Joint accelerations: s̈ ∈ ℝⁿ
-    # - Base inertial-fixed acceleration: W_v̇_WB = (W_p̈_B, W_ω̇_B) ∈ ℝ⁶
-    with (
-        data.switch_velocity_representation(velocity_representation=VelRepr.Inertial),
-        references.switch_velocity_representation(VelRepr.Inertial),
-    ):
-        W_v̇_WB, s̈ = js.model.forward_dynamics_aba(
-            model=model,
-            data=data,
-            joint_forces=references.joint_force_references(),
-            link_forces=references.link_forces(),
-        )
-    return W_v̇_WB, s̈
+    # - Base acceleration: v̇_WB ∈ ℝ⁶
+    v̇_WB, s̈ = js.model.forward_dynamics_aba(
+        model=model,
+        data=data,
+        joint_forces=references.joint_force_references(model=model),
+        link_forces=references.link_forces(model=model, data=data),
+    )
+
+    return v̇_WB, s̈
 
 
 @jax.jit
@@ -353,7 +357,7 @@ def system_dynamics(
         corresponding derivative, and the dictionary of auxiliary data returned
         by the system dynamics evaluation.
     """
-
+    from jaxsim.rbda.contacts.relaxed_rigid import RelaxedRigidContacts
     from jaxsim.rbda.contacts.rigid import RigidContacts
     from jaxsim.rbda.contacts.soft import SoftContacts
 
@@ -371,7 +375,7 @@ def system_dynamics(
         case SoftContacts():
             ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
 
-        case RigidContacts():
+        case RigidContacts() | RelaxedRigidContacts():
             pass
 
         case _:

jaxsim/api/ode_data.py

@@ -6,6 +6,10 @@ import jax_dataclasses
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim.rbda import ContactsState
+from jaxsim.rbda.contacts.relaxed_rigid import (
+    RelaxedRigidContacts,
+    RelaxedRigidContactsState,
+)
 from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
 from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 from jaxsim.utils import JaxsimDataclass
@@ -173,6 +177,10 @@ class ODEState(JaxsimDataclass):
                 )
             case RigidContacts():
                 contact = RigidContactsState.build()
+
+            case RelaxedRigidContacts():
+                contact = RelaxedRigidContactsState.build()
+
             case _:
                 raise ValueError("Unable to determine contact state class prefix.")
 
@@ -216,7 +224,9 @@ class ODEState(JaxsimDataclass):
 
         # Get the contact model from the `JaxSimModel`.
         match contact:
-            case SoftContactsState() | RigidContactsState():
+            case (
+                SoftContactsState() | RigidContactsState() | RelaxedRigidContactsState()
+            ):
                 pass
             case None:
                 contact = SoftContactsState.zero(model=model)

jaxsim/integrators/common.py

@@ -497,7 +497,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         b: jtp.Matrix,
         c: jtp.Vector,
         index_of_solution: jtp.IntLike = 0,
-    ) -> [bool, int | None]:
+    ) -> tuple[bool, int | None]:
         """
         Check if the Butcher tableau supports the FSAL (first-same-as-last) property.
 

jaxsim/math/inertia.py

@@ -45,7 +45,7 @@ class Inertia:
             M (jtp.Matrix): The 6x6 inertia matrix.
 
         Returns:
-            Tuple[jtp.Float, jtp.Vector, jtp.Matrix]: A tuple containing mass, center of mass (3D), and inertia matrix (3x3).
+            tuple[jtp.Float, jtp.Vector, jtp.Matrix]: A tuple containing mass, center of mass (3D), and inertia matrix (3x3).
 
         Raises:
             ValueError: If the input matrix M has an unexpected shape.

jaxsim/mujoco/loaders.py

@@ -211,7 +211,7 @@ class RodModelToMjcf:
         joints_dict = {j.name: j for j in rod_model.joints()}
 
         # Convert all the joints not considered to fixed joints.
-        for joint_name in set(j.name for j in rod_model.joints()) - considered_joints:
+        for joint_name in {j.name for j in rod_model.joints()} - considered_joints:
             joints_dict[joint_name].type = "fixed"
 
         # Convert the ROD model to URDF.
@@ -289,10 +289,10 @@ class RodModelToMjcf:
         mj_model = mj.MjModel.from_xml_string(xml=urdf_string, assets=assets)
 
         # Get the joint names.
-        mj_joint_names = set(
+        mj_joint_names = {
             mj.mj_id2name(mj_model, mj.mjtObj.mjOBJ_JOINT, idx)
             for idx in range(mj_model.njnt)
-        )
+        }
 
         # Check that the Mujoco model only has the considered joints.
         if mj_joint_names != considered_joints:

jaxsim/parsers/kinematic_graph.py

@@ -394,7 +394,7 @@ class KinematicGraph(Sequence[LinkDescription]):
             return copy.deepcopy(self)
 
         # Check if all considered joints are part of the full kinematic graph
-        if len(set(considered_joints) - set(j.name for j in full_graph.joints)) != 0:
+        if len(set(considered_joints) - {j.name for j in full_graph.joints}) != 0:
             extra_j = set(considered_joints) - {j.name for j in full_graph.joints}
             msg = f"Not all joints to consider are part of the graph ({{{extra_j}}})"
             raise ValueError(msg)
@@ -536,8 +536,8 @@ class KinematicGraph(Sequence[LinkDescription]):
             root_link_name=full_graph.root.name,
         )
 
-        assert set(f.name for f in self.frames).isdisjoint(
-            set(f.name for f in unconnected_frames + reduced_frames)
+        assert {f.name for f in self.frames}.isdisjoint(
+            {f.name for f in unconnected_frames + reduced_frames}
         )
 
         for link in unconnected_links: