jaxsim 0.4.3.dev31__tar.gz → 0.4.3.dev64__tar.gz

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev31
+Version: 0.4.3.dev64
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>
@@ -61,6 +61,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
 Requires-Dist: jax>=0.4.13
+Requires-Dist: jaxopt>=0.8.0
 Requires-Dist: jaxlib>=0.4.13
 Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax_dataclasses>=1.4.0

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/environment.yml

@@ -8,6 +8,7 @@ dependencies:
   - python >= 3.12.0
   - coloredlogs
   - jax >= 0.4.13
+  - jaxopt >= 0.8.0
   - jaxlib >= 0.4.13
   - jaxlie >= 1.3.0
   - jax-dataclasses >= 1.4.0

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/pyproject.toml

@@ -45,6 +45,7 @@ classifiers = [
 dependencies = [
     "coloredlogs",
     "jax >= 0.4.13",
+    "jaxopt >= 0.8.0",
     "jaxlib >= 0.4.13",
     "jaxlie >= 1.3.0",
     "jax_dataclasses >= 1.4.0",
@@ -181,6 +182,7 @@ platforms = ["linux-64", "osx-arm64", "osx-64"]
 coloredlogs = "*"
 jax = "*"
 jax-dataclasses = "*"
+jaxopt = "*"
 jaxlib = "*"
 jaxlie = "*"
 lxml = "*"

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/jaxsim/_version.py

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

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/jaxsim/api/contact.py

@@ -131,7 +131,8 @@ def collidable_point_dynamics(
     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 +145,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
 
@@ -190,6 +195,27 @@ def collidable_point_dynamics(
 
             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,
+            )
+
+            aux_data = dict()
+
         case _:
             raise ValueError(f"Invalid contact model {model.contact_model}")
 

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/jaxsim/api/data.py

@@ -593,16 +593,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)
             ),
         )
 
@@ -744,6 +746,13 @@ def random_model_data(
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
     ] = ((-1, -1, 0.5), 1.0),
+    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],
@@ -769,6 +778,8 @@ def random_model_data(
         key: The random key.
         velocity_representation: The velocity representation to use.
         base_pos_bounds: The bounds for the base position.
+        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.
@@ -813,8 +824,19 @@ def random_model_data(
         ).wxyz
 
         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-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/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-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/jaxsim/api/model.py

@@ -1935,7 +1935,7 @@ def step(
     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,
+        pred=tf_ns < t0_ns,
         true_fun=lambda: jax.debug.print(
             "The simulation time overflowed, resetting simulation time to 0."
         ),

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/jaxsim/api/ode.py

@@ -175,17 +175,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 +194,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 +204,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 +270,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 +348,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 +366,7 @@ def system_dynamics(
         case SoftContacts():
             ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
 
-        case RigidContacts():
+        case RigidContacts() | RelaxedRigidContacts():
             pass
 
         case _:

{jaxsim-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/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-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/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-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/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-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/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-0.4.3.dev31 → jaxsim-0.4.3.dev64}/src/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: