jaxsim 0.5.1.dev46__tar.gz → 0.5.1.dev56__tar.gz

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev46
+Version: 0.5.1.dev56
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

jaxsim-0.5.1.dev56/pixi.lock

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1516589d06af0d8a12e957d5fe907fee86117145d900c07f8fd99cbbb494b07f
+size 617307

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev46'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev46')
+__version__ = version = '0.5.1.dev56'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev56')

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/src/jaxsim/api/model.py

@@ -2198,7 +2198,9 @@ def step(
             isinstance(model.contact_model, jaxsim.rbda.contacts.ViscoElasticContacts)
             & (
                 ~jnp.allclose(dt, model.time_step)
-                | ~isinstance(integrator, jaxsim.integrators.fixed_step.ForwardEuler)
+                | ~int(
+                    isinstance(integrator, jaxsim.integrators.fixed_step.ForwardEuler)
+                )
             )
         ),
         msg=msg.format(module, name),

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/src/jaxsim/api/ode_data.py

@@ -10,108 +10,14 @@ import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass
 
-# =============================================================================
-# Define the input and state of the ODE system defining the integrated dynamics
-# =============================================================================
+# ===================================================================
+# Define the state of the ODE system defining the integrated dynamics
+# ===================================================================
 
 # Note: the ODE system is the combination of the floating-base dynamics and the
 #       soft-contacts dynamics.
 
 
-@jax_dataclasses.pytree_dataclass
-class ODEInput(JaxsimDataclass):
-    """
-    The input to the ODE system.
-
-    Attributes:
-        physics_model: The input to the physics model.
-    """
-
-    physics_model: PhysicsModelInput
-
-    @staticmethod
-    def build_from_jaxsim_model(
-        model: js.model.JaxSimModel | None = None,
-        link_forces: jtp.MatrixLike | None = None,
-        joint_force_references: jtp.VectorLike | None = None,
-    ) -> ODEInput:
-        """
-        Build an `ODEInput` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the ODE input.
-            link_forces: The matrix of external forces applied to the links.
-            joint_force_references: The vector of joint force references.
-
-        Returns:
-            The `ODEInput` built from the `JaxSimModel`.
-
-        Note:
-            If any of the input components are not provided, they are built from the
-            `JaxSimModel` and initialized to zero.
-        """
-
-        return ODEInput.build(
-            physics_model_input=PhysicsModelInput.build_from_jaxsim_model(
-                model=model,
-                link_forces=link_forces,
-                joint_force_references=joint_force_references,
-            ),
-            model=model,
-        )
-
-    @staticmethod
-    def build(
-        physics_model_input: PhysicsModelInput | None = None,
-        model: js.model.JaxSimModel | None = None,
-    ) -> ODEInput:
-        """
-        Build an `ODEInput` from a `PhysicsModelInput`.
-
-        Args:
-            physics_model_input: The `PhysicsModelInput` associated with the ODE input.
-            model: The `JaxSimModel` associated with the ODE input.
-
-        Returns:
-            A `ODEInput` instance.
-        """
-
-        physics_model_input = (
-            physics_model_input
-            if physics_model_input is not None
-            else PhysicsModelInput.zero(model=model)
-        )
-
-        return ODEInput(physics_model=physics_model_input)
-
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> ODEInput:
-        """
-        Build a zero `ODEInput` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the ODE input.
-
-        Returns:
-            A zero `ODEInput` instance.
-        """
-
-        return ODEInput.build(model=model)
-
-    def valid(self, model: js.model.JaxSimModel) -> bool:
-        """
-        Check if the `ODEInput` is valid for a given `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` to validate the `ODEInput` against.
-
-        Returns:
-            `True` if the ODE input is valid for the given model, `False` otherwise.
-        """
-
-        return self.physics_model.valid(model=model)
-
-
 @jax_dataclasses.pytree_dataclass
 class ODEState(JaxsimDataclass):
     """
@@ -493,123 +399,3 @@ class PhysicsModelState(JaxsimDataclass):
             return False
 
         return True
-
-
-@jax_dataclasses.pytree_dataclass
-class PhysicsModelInput(JaxsimDataclass):
-    """
-    Class storing the inputs of the physics model dynamics.
-
-    Attributes:
-        tau: The vector of joint forces.
-        f_ext: The matrix of external forces applied to the links.
-    """
-
-    tau: jtp.Vector
-    f_ext: jtp.Matrix
-
-    @staticmethod
-    def build_from_jaxsim_model(
-        model: js.model.JaxSimModel | None = None,
-        link_forces: jtp.MatrixLike | None = None,
-        joint_force_references: jtp.VectorLike | None = None,
-    ) -> PhysicsModelInput:
-        """
-        Build a `PhysicsModelInput` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the input.
-            link_forces: The matrix of external forces applied to the links.
-            joint_force_references: The vector of joint force references.
-
-        Returns:
-            A `PhysicsModelInput` instance.
-
-        Note:
-            If any of the input components are not provided, they are built from the
-            `JaxSimModel` and initialized to zero.
-        """
-
-        return PhysicsModelInput.build(
-            joint_force_references=joint_force_references,
-            link_forces=link_forces,
-            number_of_dofs=model.dofs(),
-            number_of_links=model.number_of_links(),
-        )
-
-    @staticmethod
-    def build(
-        link_forces: jtp.MatrixLike | None = None,
-        joint_force_references: jtp.VectorLike | None = None,
-        number_of_dofs: jtp.Int | None = None,
-        number_of_links: jtp.Int | None = None,
-    ) -> PhysicsModelInput:
-        """
-        Build a `PhysicsModelInput`.
-
-        Args:
-            link_forces: The matrix of external forces applied to the links.
-            joint_force_references: The vector of joint force references.
-            number_of_dofs: The number of degrees of freedom of the model.
-            number_of_links: The number of links of the model.
-
-        Returns:
-            A `PhysicsModelInput` instance.
-        """
-
-        joint_force_references = jnp.atleast_1d(
-            jnp.array(joint_force_references, dtype=float).squeeze()
-            if joint_force_references is not None
-            else jnp.zeros(number_of_dofs)
-        ).astype(float)
-
-        link_forces = jnp.atleast_2d(
-            jnp.array(link_forces, dtype=float).squeeze()
-            if link_forces is not None
-            else jnp.zeros(shape=(number_of_links, 6))
-        ).astype(float)
-
-        return PhysicsModelInput(
-            tau=joint_force_references,
-            f_ext=link_forces,
-        )
-
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> PhysicsModelInput:
-        """
-        Build a `PhysicsModelInput` with all components initialized to zero.
-
-        Args:
-            model: The `JaxSimModel` associated with the input.
-
-        Returns:
-            A `PhysicsModelInput` instance.
-        """
-
-        return PhysicsModelInput.build_from_jaxsim_model(model=model)
-
-    def valid(self, model: js.model.JaxSimModel) -> bool:
-        """
-        Check if the `PhysicsModelInput` is valid for a given `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` to validate the `PhysicsModelInput` against.
-
-        Returns:
-            `True` if the `PhysicsModelInput` is valid for the given model,
-            `False` otherwise.
-        """
-
-        shape = self.tau.shape
-        expected_shape = (model.dofs(),)
-
-        if shape != expected_shape:
-            return False
-
-        shape = self.f_ext.shape
-        expected_shape = (model.number_of_links(), 6)
-
-        if shape != expected_shape:
-            return False
-
-        return True

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/src/jaxsim/api/references.py

@@ -12,7 +12,6 @@ from jaxsim import exceptions
 from jaxsim.utils.tracing import not_tracing
 
 from .common import VelRepr
-from .ode_data import ODEInput
 
 try:
     from typing import Self
@@ -24,9 +23,14 @@ 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(
@@ -94,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(model=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,
@@ -135,12 +143,22 @@ 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
 
-        if model is not None:
-            valid = valid and self.input.valid(model=model)
+        shape = self._link_forces.shape
+        expected_shape = (model.number_of_links(), 6)
 
-        return valid
+        if shape != expected_shape:
+            return False
+
+        return True
 
     # ==================
     # Extract quantities
@@ -178,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.
@@ -190,7 +208,7 @@ 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_idxs = (
@@ -207,7 +225,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             msg = "Missing model data to use a representation different from {}"
             raise ValueError(msg.format(VelRepr.Inertial.name))
 
-        if not_tracing(self.input.physics_model.f_ext) and 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
@@ -264,9 +282,9 @@ 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_tracing(self.input.physics_model.tau) and 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)
 
@@ -277,7 +295,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         )
 
         return jnp.atleast_1d(
-            self.input.physics_model.tau[joint_idxs].squeeze()
+            self._joint_force_references[joint_idxs].squeeze()
         ).astype(float)
 
     # ================
@@ -310,11 +328,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         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:
@@ -330,7 +344,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             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))
 
     @functools.partial(jax.jit, static_argnames=["link_names", "additive"])
     def apply_link_forces(
@@ -370,11 +384,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         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
@@ -389,11 +399,7 @@ 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)
 
@@ -410,18 +416,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # 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:
@@ -450,9 +452,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         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.input.physics_model.f_ext.at[link_idxs, :].set(W_f0_L + W_f_L)
-        )
+        return replace(forces=self._link_forces.at[link_idxs, :].set(W_f0_L + W_f_L))
 
     def apply_frame_forces(
         self,

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/src/jaxsim/rbda/contacts/visco_elastic.py

@@ -851,7 +851,7 @@ class ViscoElasticContacts(common.ContactModel):
         W_f̅_L = (
             jnp.array(average_link_contact_forces_inertial)
             if average_link_contact_forces_inertial is not None
-            else jnp.zeros_like(references.input.physics_model.f_ext)
+            else jnp.zeros_like(references._link_forces)
         ).astype(float)
 
         LW_f̿_L = (

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev46
+Version: 0.5.1.dev56
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.5.1.dev46 → jaxsim-0.5.1.dev56}/tests/test_api_model.py

@@ -364,7 +364,7 @@ def test_model_jacobian(
     ):
 
         f = references.link_forces(model=model, data=data)
-        assert f == pytest.approx(references.input.physics_model.f_ext)
+        assert f == pytest.approx(references._link_forces)
 
         J = js.model.generalized_free_floating_jacobian(model=model, data=data)
         JTf_inertial = jnp.einsum("l6g,l6->g", J, f)

jaxsim-0.5.1.dev46/pixi.lock

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3ea898405d33dc9a7712d621c4539c523d766a7ef1c8e9e2e28942212b2aea97
-size 610957