jaxsim 0.4.3.dev139__py3-none-any.whl → 0.4.3.dev155__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/contact.py

@@ -144,7 +144,8 @@ def collidable_point_dynamics(
             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:
+        The 6D force applied to each collidable point and additional data based
+        on the contact model configured:
         - Soft: the material deformation rate.
         - Rigid: no additional data.
         - QuasiRigid: no additional data.
@@ -156,21 +157,13 @@ def collidable_point_dynamics(
     """
 
     # Import privately the contacts classes.
-    from jaxsim.rbda.contacts import (
-        RelaxedRigidContacts,
-        RelaxedRigidContactsState,
-        RigidContacts,
-        RigidContactsState,
-        SoftContacts,
-        SoftContactsState,
-    )
+    from jaxsim.rbda.contacts import RelaxedRigidContacts, RigidContacts, SoftContacts
 
     # Build the soft contact model.
     match model.contact_model:
 
         case SoftContacts():
             assert isinstance(model.contact_model, SoftContacts)
-            assert isinstance(data.state.contact, SoftContactsState)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point, and the corresponding material deformation rate.
@@ -187,7 +180,6 @@ def collidable_point_dynamics(
 
         case RigidContacts():
             assert isinstance(model.contact_model, RigidContacts)
-            assert isinstance(data.state.contact, RigidContactsState)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.
@@ -203,7 +195,6 @@ def collidable_point_dynamics(
 
         case RelaxedRigidContacts():
             assert isinstance(model.contact_model, RelaxedRigidContacts)
-            assert isinstance(data.state.contact, RelaxedRigidContactsState)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.

jaxsim/api/data.py

@@ -13,7 +13,6 @@ import jaxsim.api as js
 import jaxsim.math
 import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.rbda.contacts import SoftContacts
 from jaxsim.utils import Mutability
 from jaxsim.utils.tracing import not_tracing
 
@@ -107,17 +106,17 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     @staticmethod
     def build(
         model: js.model.JaxSimModel,
-        base_position: jtp.Vector | None = None,
-        base_quaternion: jtp.Vector | None = None,
-        joint_positions: jtp.Vector | None = None,
-        base_linear_velocity: jtp.Vector | None = None,
-        base_angular_velocity: jtp.Vector | None = None,
-        joint_velocities: jtp.Vector | None = None,
+        base_position: jtp.VectorLike | None = None,
+        base_quaternion: jtp.VectorLike | None = None,
+        joint_positions: jtp.VectorLike | None = None,
+        base_linear_velocity: jtp.VectorLike | None = None,
+        base_angular_velocity: jtp.VectorLike | None = None,
+        joint_velocities: jtp.VectorLike | None = None,
         standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
-        contact: jaxsim.rbda.contacts.ContactsState | None = None,
         contacts_params: jaxsim.rbda.contacts.ContactsParams | None = None,
         velocity_representation: VelRepr = VelRepr.Inertial,
         time: jtp.FloatLike | None = None,
+        extended_ode_state: dict[str, jtp.PyTree] | None = None,
     ) -> JaxSimModelData:
         """
         Create a `JaxSimModelData` object with the given state.
@@ -133,56 +132,73 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 The base angular velocity in the selected representation.
             joint_velocities: The joint velocities.
             standard_gravity: The standard gravity constant.
-            contact: The state of the soft contacts.
             contacts_params: The parameters of the soft contacts.
             velocity_representation: The velocity representation to use.
             time: The time at which the state is created.
+            extended_ode_state:
+                Additional user-defined state variables that are not part of the
+                standard `ODEState` object. Useful to extend the system dynamics
+                considered by default in JaxSim.
 
         Returns:
-            A `JaxSimModelData` object with the given state.
+            A `JaxSimModelData` initialized with the given state.
         """
 
         base_position = jnp.array(
-            base_position if base_position is not None else jnp.zeros(3)
+            base_position if base_position is not None else jnp.zeros(3),
+            dtype=float,
         ).squeeze()
 
         base_quaternion = jnp.array(
-            base_quaternion
-            if base_quaternion is not None
-            else jnp.array([1.0, 0, 0, 0])
+            (
+                base_quaternion
+                if base_quaternion is not None
+                else jnp.array([1.0, 0, 0, 0])
+            ),
+            dtype=float,
         ).squeeze()
 
         base_linear_velocity = jnp.array(
-            base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3)
+            base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3),
+            dtype=float,
         ).squeeze()
 
         base_angular_velocity = jnp.array(
-            base_angular_velocity if base_angular_velocity is not None else jnp.zeros(3)
+            (
+                base_angular_velocity
+                if base_angular_velocity is not None
+                else jnp.zeros(3)
+            ),
+            dtype=float,
         ).squeeze()
 
         gravity = jnp.zeros(3).at[2].set(-standard_gravity)
 
         joint_positions = jnp.atleast_1d(
-            joint_positions.squeeze()
-            if joint_positions is not None
-            else jnp.zeros(model.dofs())
+            jnp.array(
+                (
+                    joint_positions
+                    if joint_positions is not None
+                    else jnp.zeros(model.dofs())
+                ),
+                dtype=float,
+            ).squeeze()
         )
 
         joint_velocities = jnp.atleast_1d(
-            joint_velocities.squeeze()
-            if joint_velocities is not None
-            else jnp.zeros(model.dofs())
+            jnp.array(
+                (
+                    joint_velocities
+                    if joint_velocities is not None
+                    else jnp.zeros(model.dofs())
+                ),
+                dtype=float,
+            ).squeeze()
         )
 
-        time_ns = (
-            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 if jax.config.read("jax_enable_x64") else jnp.uint32
-            )
+        time_ns = jnp.array(
+            time * 1e9 if time is not None else 0.0,
+            dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32,
         )
 
         W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
@@ -194,21 +210,22 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             other_representation=velocity_representation,
             transform=W_H_B,
             is_force=False,
-        )
+        ).astype(float)
 
         ode_state = ODEState.build_from_jaxsim_model(
             model=model,
-            base_position=base_position.astype(float),
-            base_quaternion=base_quaternion.astype(float),
-            joint_positions=joint_positions.astype(float),
-            base_linear_velocity=v_WB[0:3].astype(float),
-            base_angular_velocity=v_WB[3:6].astype(float),
-            joint_velocities=joint_velocities.astype(float),
-            tangential_deformation=(
-                contact.tangential_deformation
-                if contact is not None and isinstance(model.contact_model, SoftContacts)
-                else None
-            ),
+            base_position=base_position,
+            base_quaternion=base_quaternion,
+            joint_positions=joint_positions,
+            base_linear_velocity=v_WB[0:3],
+            base_angular_velocity=v_WB[3:6],
+            joint_velocities=joint_velocities,
+            # Unpack all the additional ODE states. If the contact model requires an
+            # additional state that is not explicitly passed to this builder, ODEState
+            # automatically populates that state with zeroed variables.
+            # This is not true for any other custom state that the user might want to
+            # pass to the integrator.
+            **(extended_ode_state if extended_ode_state else {}),
         )
 
         if not ode_state.valid(model=model):
@@ -220,13 +237,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 contacts_params = js.contact.estimate_good_soft_contacts_parameters(
                     model=model, standard_gravity=standard_gravity
                 )
+
             else:
                 contacts_params = model.contact_model.parameters
 
         return JaxSimModelData(
             time_ns=time_ns,
             state=ode_state,
-            gravity=gravity.astype(float),
+            gravity=gravity,
             contacts_params=contacts_params,
             velocity_representation=velocity_representation,
         )

jaxsim/api/model.py

@@ -33,7 +33,7 @@ class JaxSimModel(JaxsimDataclass):
     model_name: Static[str]
 
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
-        default=jaxsim.terrain.FlatTerrain(), repr=False
+        default_factory=jaxsim.terrain.FlatTerrain.build, repr=False
     )
 
     contact_model: jaxsim.rbda.contacts.ContactModel | None = dataclasses.field(
@@ -101,13 +101,14 @@ class JaxSimModel(JaxsimDataclass):
                 A path to an SDF/URDF file, a string containing
                 its content, or a pre-parsed/pre-built rod model.
             model_name:
-                The optional name of the model that overrides the one in
-                the description.
-            terrain:
-                The optional terrain to consider.
+                The name of the model. If not specified, it is read from the description.
+            terrain: The terrain to consider (the default is a flat infinite plane).
+            contact_model:
+                The contact model to consider.
+                If not specified, a soft contacts model is used.
             is_urdf:
-                The optional flag to force the model description to be parsed as a
-                URDF or a SDF. This is otherwise automatically inferred.
+                The optional flag to force the model description to be parsed as a URDF.
+                This is usually automatically inferred.
             considered_joints:
                 The list of joints to consider. If None, all joints are considered.
 
@@ -120,7 +121,7 @@ class JaxSimModel(JaxsimDataclass):
         # Parse the input resource (either a path to file or a string with the URDF/SDF)
         # and build the -intermediate- model description.
         intermediate_description = jaxsim.parsers.rod.build_model_description(
-            model_description=model_description
+            model_description=model_description, is_urdf=is_urdf
         )
 
         # Lump links together if not all joints are considered.
@@ -160,11 +161,11 @@ class JaxSimModel(JaxsimDataclass):
                 The intermediate model description defining the kinematics and dynamics
                 of the model.
             model_name:
-                The optional name of the model overriding the physics model name.
-            terrain:
-                The optional terrain to consider.
+                The name of the model. If not specified, it is read from the description.
+            terrain: The terrain to consider (the default is a flat infinite plane).
             contact_model:
-                The optional contact model to consider. If None, the soft contact model is used.
+                The contact model to consider.
+                If not specified, a soft contacts model is used.
 
         Returns:
             The built Model object.
@@ -173,21 +174,31 @@ class JaxSimModel(JaxsimDataclass):
         # Set the model name (if not provided, use the one from the model description).
         model_name = model_name if model_name is not None else model_description.name
 
-        # Set the terrain (if not provided, use the default flat terrain).
-        terrain = terrain or JaxSimModel.__dataclass_fields__["terrain"].default
-        contact_model = contact_model or jaxsim.rbda.contacts.SoftContacts(
-            terrain=terrain
+        # Consider the default terrain (a flat infinite plane) if not specified.
+        terrain = (
+            terrain or JaxSimModel.__dataclass_fields__["terrain"].default_factory()
+        )
+
+        # Create the default contact model.
+        # It will be populated with an initial estimation of good parameters.
+        # While these might not be the best, they are a good starting point.
+        contact_model = contact_model or jaxsim.rbda.contacts.SoftContacts.build(
+            terrain=terrain, parameters=None
         )
 
         # Build the model.
         model = JaxSimModel(
             model_name=model_name,
-            _description=wrappers.HashlessObject(obj=model_description),
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
             terrain=terrain,
             contact_model=contact_model,
+            # The following is wrapped as hashless since it's a static argument, and we
+            # don't want to trigger recompilation if it changes. All relevant parameters
+            # needed to compute kinematics and dynamics quantities are stored in the
+            # kin_dyn_parameters attribute.
+            _description=wrappers.HashlessObject(obj=model_description),
         )
 
         return model
@@ -1907,8 +1918,8 @@ def step(
     dt: jtp.FloatLike,
     integrator: jaxsim.integrators.Integrator,
     integrator_state: dict[str, Any] | None = None,
-    joint_forces: jtp.VectorLike | None = None,
     link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
     **kwargs,
 ) -> tuple[js.data.JaxSimModelData, dict[str, Any]]:
     """
@@ -1920,10 +1931,10 @@ def step(
         dt: The time step to consider.
         integrator: The integrator to use.
         integrator_state: The state of the integrator.
-        joint_forces: The joint forces to consider.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
+        joint_force_references: The joint force references to consider.
         kwargs: Additional kwargs to pass to the integrator.
 
     Returns:
@@ -1953,7 +1964,7 @@ def step(
         params=integrator_state_x0,
         # Always inject the current (model, data) pair into the system dynamics
         # considered by the integrator, and include the input variables represented
-        # by the pair (joint_forces, link_forces).
+        # by the pair (joint_force_references, link_forces).
         # Note that the wrapper of the system dynamics will override (state_x0, t0)
         # inside the passed data even if it is not strictly needed. This logic is
         # necessary to re-use the jit-compiled step function of compatible pytrees
@@ -1962,7 +1973,7 @@ def step(
             dict(
                 model=model,
                 data=data,
-                joint_forces=joint_forces,
+                joint_force_references=joint_force_references,
                 link_forces=link_forces,
             )
             | integrator_kwargs

jaxsim/api/ode.py

@@ -86,8 +86,8 @@ def system_velocity_dynamics(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.Vector | None = None,
     link_forces: jtp.Vector | None = None,
+    joint_force_references: jtp.Vector | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector, dict[str, Any]]:
     """
     Compute the dynamics of the system velocity.
@@ -95,10 +95,10 @@ def system_velocity_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        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`.
+        joint_force_references: The joint force references to apply.
 
     Returns:
         A tuple containing the derivative of the base 6D velocity in inertial-fixed
@@ -120,7 +120,7 @@ def system_velocity_dynamics(
     references = js.references.JaxSimModelReferences.build(
         model=model,
         link_forces=O_f_L,
-        joint_force_references=joint_forces,
+        joint_force_references=joint_force_references,
         data=data,
         velocity_representation=data.velocity_representation,
     )
@@ -192,7 +192,10 @@ def system_velocity_dynamics(
         f_L_total = references.link_forces(model=model, data=data)
 
         v̇_WB, s̈ = system_acceleration(
-            model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
+            model=model,
+            data=data,
+            joint_force_references=joint_force_references,
+            link_forces=f_L_total,
         )
 
     return v̇_WB, s̈, aux_data
@@ -202,8 +205,8 @@ def system_acceleration(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.VectorLike | None = None,
     link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
     Compute the system acceleration in the active representation.
@@ -211,12 +214,13 @@ def system_acceleration(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint forces to apply.
         link_forces:
-            The 6D forces to apply to the links expressed in the same representation of data.
+            The 6D forces to apply to the links expressed in the same
+            velocity representation of data.
+        joint_force_references: The joint force references to apply.
 
     Returns:
-        A tuple containing the base 6D acceleration in in the active representation
+        A tuple containing the base 6D acceleration in the active representation
         and the joint accelerations.
     """
 
@@ -232,9 +236,9 @@ def system_acceleration(
     ).astype(float)
 
     # Build joint torques if not provided.
-    τ = (
-        jnp.atleast_1d(joint_forces.squeeze())
-        if joint_forces is not None
+    τ_references = (
+        jnp.atleast_1d(joint_force_references.squeeze())
+        if joint_force_references is not None
         else jnp.zeros_like(data.joint_positions())
     ).astype(float)
 
@@ -243,15 +247,16 @@ def system_acceleration(
     # ====================
 
     # TODO: enforce joint limits
-    τ_position_limit = jnp.zeros_like(τ).astype(float)
+    τ_position_limit = jnp.zeros_like(τ_references).astype(float)
 
     # ====================
     # Joint friction model
     # ====================
 
-    τ_friction = jnp.zeros_like(τ).astype(float)
+    τ_friction = jnp.zeros_like(τ_references).astype(float)
 
     if model.dofs() > 0:
+
         # Static and viscous joint friction parameters
         kc = jnp.array(
             model.kin_dyn_parameters.joint_parameters.friction_static
@@ -271,22 +276,27 @@ def system_acceleration(
     # ========================
 
     # Compute the total joint forces.
-    τ_total = τ + τ_friction + τ_position_limit
+    τ_total = τ_references + τ_friction + τ_position_limit
 
+    # Store the link forces in a references object.
     references = js.references.JaxSimModelReferences.build(
         model=model,
         data=data,
         velocity_representation=data.velocity_representation,
-        joint_force_references=τ_total,
         link_forces=f_L,
     )
 
+    # Compute forward dynamics.
+    #
     # - Joint accelerations: s̈ ∈ ℝⁿ
     # - Base acceleration: v̇_WB ∈ ℝ⁶
+    #
+    # Note that ABA returns the base acceleration in the velocity representation
+    # stored in the `data` object.
     v̇_WB, s̈ = js.model.forward_dynamics_aba(
         model=model,
         data=data,
-        joint_forces=references.joint_force_references(model=model),
+        joint_forces=τ_total,
         link_forces=references.link_forces(model=model, data=data),
     )
 
@@ -337,8 +347,8 @@ def system_dynamics(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.Vector | None = None,
     link_forces: jtp.Vector | None = None,
+    joint_force_references: jtp.Vector | None = None,
     baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[ODEState, dict[str, Any]]:
     """
@@ -347,10 +357,10 @@ def system_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint forces to apply.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
+        joint_force_references: The joint force references to apply.
         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).
@@ -360,29 +370,31 @@ 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
+
+    from jaxsim.rbda.contacts import RelaxedRigidContacts, RigidContacts, SoftContacts
 
     # Compute the accelerations and the material deformation rate.
     W_v̇_WB, s̈, aux_dict = system_velocity_dynamics(
         model=model,
         data=data,
-        joint_forces=joint_forces,
+        joint_force_references=joint_force_references,
         link_forces=link_forces,
     )
 
-    ode_state_kwargs = {}
+    # Initialize the dictionary storing the derivative of the additional state variables
+    # that extend the state vector of the integrated ODE system.
+    extended_ode_state = {}
 
     match model.contact_model:
+
         case SoftContacts():
-            ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
+            extended_ode_state["tangential_deformation"] = aux_dict["m_dot"]
 
         case RigidContacts() | RelaxedRigidContacts():
             pass
 
         case _:
-            raise ValueError("Unable to determine contact state class prefix.")
+            raise ValueError(f"Invalid contact model {model.contact_model}")
 
     # Extract the velocities.
     W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
@@ -402,7 +414,7 @@ def system_dynamics(
         base_linear_velocity=W_v̇_WB[0:3],
         base_angular_velocity=W_v̇_WB[3:6],
         joint_velocities=s̈,
-        **ode_state_kwargs,
+        **extended_ode_state,
     )
 
     return ode_state_derivative, aux_dict