jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev5__py3-none-any.whl

jaxsim/api/references.py

@@ -7,10 +7,11 @@ import jax.numpy as jnp
 import jax_dataclasses
 
 import jaxsim.api as js
-import jaxsim.physics.model.physics_model_state
 import jaxsim.typing as jtp
-from jaxsim import VelRepr
-from jaxsim.simulation.ode_data import ODEInput
+from jaxsim import exceptions
+from jaxsim.utils.tracing import not_tracing
+
+from .common import VelRepr
 
 try:
     from typing import Self
@@ -22,13 +23,19 @@ except ImportError:
 class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
     """
     Class containing the references for a `JaxSimModel` object.
+
+    Attributes:
+        _link_forces: The link 6D forces in inertial-fixed representation.
+        _joint_force_references: The joint force references.
     """
 
-    input: ODEInput
+    _link_forces: jtp.Matrix
+    _joint_force_references: jtp.Vector
 
     @staticmethod
     def zero(
         model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData | None = None,
         velocity_representation: VelRepr = VelRepr.Inertial,
     ) -> JaxSimModelReferences:
         """
@@ -36,6 +43,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         Args:
             model: The model for which to create the zero references.
+            data:
+                The data of the model, only needed if the velocity representation is
+                not inertial-fixed.
             velocity_representation: The velocity representation to use.
 
         Returns:
@@ -43,14 +53,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         """
 
         return JaxSimModelReferences.build(
-            model=model, velocity_representation=velocity_representation
+            model=model, data=data, velocity_representation=velocity_representation
         )
 
     @staticmethod
     def build(
         model: js.model.JaxSimModel,
-        joint_force_references: jtp.Vector | None = None,
-        link_forces: jtp.Matrix | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
+        link_forces: jtp.MatrixLike | None = None,
         data: js.data.JaxSimModelData | None = None,
         velocity_representation: VelRepr | None = None,
     ) -> JaxSimModelReferences:
@@ -72,14 +82,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # Create or adjust joint force references.
         joint_force_references = jnp.atleast_1d(
-            joint_force_references.squeeze()
+            jnp.array(joint_force_references, dtype=float).squeeze()
             if joint_force_references is not None
             else jnp.zeros(model.dofs())
         ).astype(float)
 
         # Create or adjust link forces.
         f_L = jnp.atleast_2d(
-            link_forces.squeeze()
+            jnp.array(link_forces, dtype=float).squeeze()
             if link_forces is not None
             else jnp.zeros((model.number_of_links(), 6))
         ).astype(float)
@@ -88,17 +98,21 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         velocity_representation = (
             velocity_representation
             if velocity_representation is not None
-            else (
-                data.velocity_representation if data is not None else VelRepr.Inertial
-            )
+            else getattr(data, "velocity_representation", VelRepr.Inertial)
         )
 
         # Create a zero references object.
         references = JaxSimModelReferences(
-            input=ODEInput.zero(physics_model=model.physics_model),
+            _link_forces=f_L,
+            _joint_force_references=joint_force_references,
             velocity_representation=velocity_representation,
         )
 
+        # If the velocity representation is inertial-fixed, we can return
+        # the references directly, as we store the link forces in this frame.
+        if velocity_representation is VelRepr.Inertial:
+            return references
+
         # Store the joint force references.
         references = references.set_joint_force_references(
             forces=joint_force_references,
@@ -129,17 +143,27 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             `False` otherwise.
         """
 
-        valid = True
+        if model is None:
+            return True
+
+        shape = self._joint_force_references.shape
+        expected_shape = (model.dofs(),)
+
+        if shape != expected_shape:
+            return False
+
+        shape = self._link_forces.shape
+        expected_shape = (model.number_of_links(), 6)
 
-        if model is not None:
-            valid = valid and self.input.valid(physics_model=model.physics_model)
+        if shape != expected_shape:
+            return False
 
-        return valid
+        return True
 
     # ==================
     # Extract quantities
     # ==================
-
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["link_names"])
     def link_forces(
         self,
@@ -172,7 +196,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             e.g. to the contact model and other kinematic constraints.
         """
 
-        W_f_L = self.input.physics_model.f_ext
+        W_f_L = self._link_forces
 
         # Return all link forces in inertial-fixed representation using the implicit
         # serialization.
@@ -184,11 +208,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             if link_names is not None:
                 raise ValueError("Link names cannot be provided without a model")
 
-            return self.input.physics_model.f_ext
+            return W_f_L
 
         # If we have the model, we can extract the link names, if not provided.
-        link_names = link_names if link_names is not None else model.link_names()
-        link_idxs = jaxsim.api.link.names_to_idxs(link_names=link_names, model=model)
+        link_idxs = (
+            js.link.names_to_idxs(link_names=link_names, model=model)
+            if link_names is not None
+            else jnp.arange(model.number_of_links())
+        )
 
         # In inertial-fixed representation, we already have the link forces.
         if self.velocity_representation is VelRepr.Inertial:
@@ -198,20 +225,25 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             msg = "Missing model data to use a representation different from {}"
             raise ValueError(msg.format(VelRepr.Inertial.name))
 
-        if not data.valid(model=model):
+        if not_tracing(self._link_forces) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
-        # Helper function to convert a single 6D force to the active representation.
-        def convert(f_L: jtp.Vector) -> jtp.Vector:
-            return JaxSimModelReferences.inertial_to_other_representation(
-                array=f_L,
-                other_representation=self.velocity_representation,
-                transform=data.base_transform(),
-                is_force=True,
-            )
+        # Helper function to convert a single 6D force to the active representation
+        # considering as body the link (i.e. L_f_L and LW_f_L).
+        def convert(W_f_L: jtp.MatrixLike, W_H_L: jtp.ArrayLike) -> jtp.Matrix:
 
-        # Convert to the desired representation.
-        f_L = jax.vmap(convert)(W_f_L[link_idxs, :])
+            return jax.vmap(
+                lambda W_f_L, W_H_L: JaxSimModelReferences.inertial_to_other_representation(
+                    array=W_f_L,
+                    other_representation=self.velocity_representation,
+                    transform=W_H_L,
+                    is_force=True,
+                )
+            )(W_f_L, W_H_L)
+
+        # The f_L output is either L_f_L or LW_f_L, depending on the representation.
+        W_H_L = js.model.forward_kinematics(model=model, data=data)
+        f_L = convert(W_f_L=W_f_L[link_idxs, :], W_H_L=W_H_L[link_idxs, :, :])
 
         return f_L
 
@@ -250,23 +282,26 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             if joint_names is not None:
                 raise ValueError("Joint names cannot be provided without a model")
 
-            return self.input.physics_model.tau
+            return self._joint_force_references
 
-        if not self.valid(model=model):
+        if not_tracing(self._joint_force_references) and not self.valid(model=model):
             msg = "The actuation object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
-        joint_idxs = js.joint.names_to_idxs(joint_names=joint_names, model=model)
+        joint_idxs = (
+            js.joint.names_to_idxs(joint_names=joint_names, model=model)
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
 
         return jnp.atleast_1d(
-            self.input.physics_model.tau[joint_idxs].squeeze()
+            self._joint_force_references[joint_idxs].squeeze()
         ).astype(float)
 
     # ================
     # Store quantities
     # ================
-
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def set_joint_force_references(
         self,
@@ -288,37 +323,37 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             A new `JaxSimModelReferences` object with the given joint force references.
         """
 
-        forces = jnp.array(forces)
+        forces = jnp.atleast_1d(jnp.array(forces, dtype=float).squeeze())
 
-        def replace(forces: jtp.VectorLike) -> JaxSimModelReferences:
+        def replace(forces: jtp.Vector) -> JaxSimModelReferences:
             return self.replace(
                 validate=True,
-                input=self.input.replace(
-                    physics_model=self.input.physics_model.replace(
-                        tau=jnp.atleast_1d(forces.squeeze()).astype(float)
-                    )
-                ),
+                _joint_force_references=jnp.atleast_1d(forces.squeeze()).astype(float),
             )
 
         if model is None:
             return replace(forces=forces)
 
-        if not self.valid(model=model):
+        if not_tracing(forces) and not self.valid(model=model):
             msg = "The references object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
-        joint_idxs = js.joint.names_to_idxs(joint_names=joint_names, model=model)
+        joint_idxs = (
+            js.joint.names_to_idxs(joint_names=joint_names, model=model)
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
 
-        return replace(forces=self.input.physics_model.tau.at[joint_idxs].set(forces))
+        return replace(forces=self._joint_force_references.at[joint_idxs].set(forces))
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["link_names", "additive"])
     def apply_link_forces(
         self,
         forces: jtp.MatrixLike,
         model: js.model.JaxSimModel | None = None,
         data: js.data.JaxSimModelData | None = None,
-        link_names: tuple[str, ...] | None = None,
+        link_names: tuple[str, ...] | str | None = None,
         additive: bool = False,
     ) -> Self:
         """
@@ -344,17 +379,13 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             Then, we always convert and store forces in inertial-fixed representation.
         """
 
-        f_L = jnp.array(forces)
+        f_L = jnp.atleast_2d(forces).astype(float)
 
         # Helper function to replace the link forces.
         def replace(forces: jtp.MatrixLike) -> JaxSimModelReferences:
             return self.replace(
                 validate=True,
-                input=self.input.replace(
-                    physics_model=self.input.physics_model.replace(
-                        f_ext=jnp.atleast_2d(forces.squeeze()).astype(float)
-                    )
-                ),
+                _link_forces=jnp.atleast_2d(forces.squeeze()).astype(float),
             )
 
         # In this case, we allow only to set the inertial 6D forces to all links
@@ -369,52 +400,157 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
             W_f_L = f_L
 
-            W_f0_L = (
-                jnp.zeros_like(W_f_L)
-                if not additive
-                else self.input.physics_model.f_ext
-            )
+            W_f0_L = jnp.zeros_like(W_f_L) if not additive else self._link_forces
 
             return replace(forces=W_f0_L + W_f_L)
 
-        # If we have the model, we can extract the link names if not provided.
-        link_names = link_names if link_names is not None else model.link_names()
-        link_idxs = jaxsim.api.link.names_to_idxs(link_names=link_names, model=model)
+        if link_names is not None and len(link_names) != f_L.shape[0]:
+            msg = "The number of link names ({}) must match the number of forces ({})"
+            raise ValueError(msg.format(len(link_names), f_L.shape[0]))
+
+        # Extract the link indices.
+        link_idxs = (
+            js.link.names_to_idxs(link_names=link_names, model=model)
+            if link_names is not None
+            else jnp.arange(model.number_of_links())
+        )
 
         # Compute the bias depending on whether we either set or add the link forces.
         W_f0_L = (
-            jnp.zeros_like(f_L)
-            if not additive
-            else self.input.physics_model.f_ext[link_idxs, :]
+            jnp.zeros_like(f_L) if not additive else self._link_forces[link_idxs, :]
         )
 
         # If inertial-fixed representation, we can directly store the link forces.
         if self.velocity_representation is VelRepr.Inertial:
             W_f_L = f_L
             return replace(
-                forces=self.input.physics_model.f_ext.at[link_idxs, :].set(
-                    W_f0_L + W_f_L
-                )
+                forces=self._link_forces.at[link_idxs, :].set(W_f0_L + W_f_L)
             )
 
         if data is None:
             msg = "Missing model data to use a representation different from {}"
             raise ValueError(msg.format(VelRepr.Inertial.name))
 
-        if not data.valid(model=model):
+        if not_tracing(forces) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
-        # Helper function to convert a single 6D force to the inertial representation.
-        def convert(f_L: jtp.Vector) -> jtp.Vector:
+        # Helper function to convert a single 6D force to the inertial representation
+        # considering as body the link (i.e. L_f_L and LW_f_L).
+        def convert_using_link_frame(
+            f_L: jtp.MatrixLike, W_H_L: jtp.ArrayLike
+        ) -> jtp.Matrix:
+
+            return jax.vmap(
+                lambda f_L, W_H_L: JaxSimModelReferences.other_representation_to_inertial(
+                    array=f_L,
+                    other_representation=self.velocity_representation,
+                    transform=W_H_L,
+                    is_force=True,
+                )
+            )(f_L, W_H_L)
+
+        # The f_L input is either L_f_L or LW_f_L, depending on the representation.
+        W_H_L = js.model.forward_kinematics(model=model, data=data)
+        W_f_L = convert_using_link_frame(f_L=f_L, W_H_L=W_H_L[link_idxs, :, :])
+
+        return replace(forces=self._link_forces.at[link_idxs, :].set(W_f0_L + W_f_L))
+
+    def apply_frame_forces(
+        self,
+        forces: jtp.MatrixLike,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        frame_names: tuple[str, ...] | str | None = None,
+        additive: bool = False,
+    ) -> Self:
+        """
+        Apply the frame forces.
+
+        Args:
+            forces: The frame 6D forces in the active representation.
+            model:
+                The model to consider, only needed if a frame serialization different
+                from the implicit one is used.
+            data:
+                The data of the considered model, only needed if the velocity
+                representation is not inertial-fixed.
+            frame_names: The names of the frames corresponding to the forces.
+            additive:
+                Whether to add the forces to the existing ones instead of replacing them.
+
+        Returns:
+            A new `JaxSimModelReferences` object with the given frame forces.
+
+        Note:
+            The frame forces must be expressed in the active representation.
+            Then, we always convert and store forces in inertial-fixed representation.
+        """
+
+        f_F = jnp.atleast_2d(forces).astype(float)
+
+        if len(frame_names) != f_F.shape[0]:
+            msg = "The number of frame names ({}) must match the number of forces ({})"
+            raise ValueError(msg.format(len(frame_names), f_F.shape[0]))
+
+        # Extract the frame indices.
+        frame_idxs = (
+            js.frame.names_to_idxs(frame_names=frame_names, model=model)
+            if frame_names is not None
+            else jnp.arange(len(model.frame_names()))
+        )
+
+        parent_link_idxs = jnp.array(model.kin_dyn_parameters.frame_parameters.body)[
+            frame_idxs - model.number_of_links()
+        ]
+
+        exceptions.raise_value_error_if(
+            condition=~data.valid(model=model),
+            msg="The provided data is not valid for the model",
+        )
+        W_H_Fi = jax.vmap(
+            lambda frame_idx: js.frame.transform(
+                model=model, data=data, frame_index=frame_idx
+            )
+        )(frame_idxs)
+
+        # Helper function to convert a single 6D force to the inertial representation
+        # considering as body the frame (i.e. L_f_F and LW_f_F).
+        def to_inertial(f_F: jtp.MatrixLike, W_H_F: jtp.MatrixLike) -> jtp.Matrix:
             return JaxSimModelReferences.other_representation_to_inertial(
-                array=f_L,
+                array=f_F,
                 other_representation=self.velocity_representation,
-                transform=data.base_transform(),
+                transform=W_H_F,
                 is_force=True,
             )
 
-        W_f_L = jax.vmap(convert)(f_L)
+        match self.velocity_representation:
+            case VelRepr.Inertial:
+                W_f_F = f_F
 
-        return replace(
-            forces=self.input.physics_model.f_ext.at[link_idxs, :].set(W_f0_L + W_f_L)
-        )
+            case VelRepr.Body | VelRepr.Mixed:
+                W_f_F = jax.vmap(to_inertial)(f_F, W_H_Fi)
+
+            case _:
+                raise ValueError("Invalid velocity representation.")
+
+        # Sum the forces on the parent links.
+        mask = parent_link_idxs[:, jnp.newaxis] == jnp.arange(model.number_of_links())
+        W_f_L = mask.T @ W_f_F
+
+        with self.switch_velocity_representation(
+            velocity_representation=VelRepr.Inertial
+        ):
+            references = self.apply_link_forces(
+                model=model,
+                data=data,
+                link_names=js.link.idxs_to_names(
+                    model=model, link_indices=parent_link_idxs
+                ),
+                forces=W_f_L,
+                additive=additive,
+            )
+
+        with references.switch_velocity_representation(
+            velocity_representation=self.velocity_representation
+        ):
+            return references

jaxsim/exceptions.py

@@ -0,0 +1,80 @@
+import os
+
+import jax
+
+
+def raise_if(
+    condition: bool | jax.Array, exception: type, msg: str, *args, **kwargs
+) -> None:
+    """
+    Raise a host-side exception if a condition is met. Useful in jit-compiled functions.
+
+    Args:
+        condition:
+            The boolean condition of the evaluated expression that triggers
+            the exception during runtime.
+        exception: The type of exception to raise.
+        msg:
+            The message to display when the exception is raised. The message can be a
+            format string (fmt), whose fields are filled with the args and kwargs.
+        *args: The arguments to fill the format string.
+        **kwargs: The keyword arguments to fill the format string
+    """
+
+    # Disable host callback if running on unsupported hardware or if the user
+    # explicitly disabled it.
+    if jax.devices()[0].platform in {"tpu", "METAL"} or os.environ.get(
+        "JAXSIM_DISABLE_EXCEPTIONS", 0
+    ):
+        return
+
+    # Check early that the format string is well-formed.
+    try:
+        _ = msg.format(*args, **kwargs)
+    except Exception as e:
+        msg = "Error in formatting exception message with args={} and kwargs={}"
+        raise ValueError(msg.format(args, kwargs)) from e
+
+    def _raise_exception(condition: bool, *args, **kwargs) -> None:
+        """The function called by the JAX callback."""
+
+        if condition:
+            raise exception(msg.format(*args, **kwargs))
+
+    def _callback(args, kwargs) -> None:
+        """The function that calls the JAX callback, executed only when needed."""
+
+        jax.debug.callback(_raise_exception, condition, *args, **kwargs)
+
+    # Since running a callable on the host is expensive, we prevent its execution
+    # if the condition is False with a low-level conditional expression.
+    def _run_callback_only_if_condition_is_true(*args, **kwargs) -> None:
+        return jax.lax.cond(
+            condition,
+            _callback,
+            lambda args, kwargs: None,
+            args,
+            kwargs,
+        )
+
+    return _run_callback_only_if_condition_is_true(*args, **kwargs)
+
+
+def raise_runtime_error_if(
+    condition: bool | jax.Array, msg: str, *args, **kwargs
+) -> None:
+    """
+    Raise a RuntimeError if a condition is met. Useful in jit-compiled functions.
+    """
+
+    return raise_if(condition, RuntimeError, msg, *args, **kwargs)
+
+
+def raise_value_error_if(
+    condition: bool | jax.Array, msg: str, *args, **kwargs
+) -> None:
+    """
+    Raise a ValueError if a condition is met. Useful in jit-compiled functions.
+    """
+
+    return raise_if(condition, ValueError, msg, *args, **kwargs)

jaxsim/integrators/__init__.py

@@ -1,2 +1,2 @@
-from . import fixed_step
-from .common import Integrator, Time, TimeStep
+from . import fixed_step, variable_step
+from .common import Integrator, SystemDynamics, Time, TimeStep