jaxsim 0.4.3.dev100__tar.gz → 0.4.3.dev115__tar.gz

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev100
+Version: 0.4.3.dev115
 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>
@@ -67,7 +67,7 @@ Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax_dataclasses>=1.4.0
 Requires-Dist: pptree
 Requires-Dist: qpax
-Requires-Dist: rod>=0.3.0
+Requires-Dist: rod>=0.3.3
 Requires-Dist: typing_extensions; python_version < "3.12"
 Provides-Extra: style
 Requires-Dist: black[jupyter]~=24.0; extra == "style"

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/environment.yml

@@ -14,7 +14,7 @@ dependencies:
   - jax-dataclasses >= 1.4.0
   - pptree
   - qpax
-  - rod >= 0.3.0
+  - rod >= 0.3.3
   - typing_extensions # python<3.12
   # ====================================
   # Optional dependencies from setup.cfg

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/examples/PD_controller.ipynb

@@ -100,7 +100,7 @@
     "num_steps = int(integration_time / dt)\n",
     "\n",
     "model = js.model.JaxSimModel.build_from_model_description(\n",
-    "    model_description=model_urdf_string, is_urdf=True\n",
+    "    model_description=model_urdf_string\n",
     ")\n",
     "data = js.data.JaxSimModelData.build(model=model)\n",
     "integrator = integrators.fixed_step.RungeKutta4SO3.build(\n",

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/pyproject.toml

@@ -51,7 +51,7 @@ dependencies = [
     "jax_dataclasses >= 1.4.0",
     "pptree",
     "qpax",
-    "rod >= 0.3.0",
+    "rod >= 0.3.3",
     "typing_extensions ; python_version < '3.12'",
 ]
 

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/_version.py

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

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/api/contact.py

@@ -186,7 +186,9 @@ def collidable_point_dynamics(
             # Note that the material deformation rate is always returned in the mixed frame
             # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
             W_f_Ci, (CW_ṁ,) = jax.vmap(soft_contacts.compute_contact_forces)(
-                W_p_Ci, W_ṗ_Ci, data.state.contact.tangential_deformation
+                position=W_p_Ci,
+                velocity=W_ṗ_Ci,
+                tangential_deformation=data.state.contact.tangential_deformation,
             )
             aux_data = dict(m_dot=CW_ṁ)
 

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/api/model.py

@@ -106,8 +106,8 @@ class JaxSimModel(JaxsimDataclass):
             terrain:
                 The optional terrain to consider.
             is_urdf:
-                Whether the model description is a URDF or an SDF. This is
-                automatically inferred if the model description is a path to a file.
+                The optional flag to force the model description to be parsed as a
+                URDF or a SDF. This is otherwise automatically inferred.
             considered_joints:
                 The list of joints to consider. If None, all joints are considered.
 
@@ -120,7 +120,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, is_urdf=is_urdf
+            model_description=model_description
         )
 
         # Lump links together if not all joints are considered.

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/mujoco/loaders.py

@@ -25,7 +25,7 @@ def load_rod_model(
 
     Args:
         model_description: The URDF/SDF file or ROD model to load.
-        is_urdf: Whether the model description is a URDF file.
+        is_urdf: Whether to force parsing the model description as a URDF file.
         model_name: The name of the model to load from the resource.
 
     Returns:

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/parsers/rod/parser.py

@@ -33,7 +33,7 @@ class SDFData(NamedTuple):
 
 
 def extract_model_data(
-    model_description: pathlib.Path | str | rod.Model,
+    model_description: pathlib.Path | str | rod.Model | rod.Sdf,
     model_name: str | None = None,
     is_urdf: bool | None = None,
 ) -> SDFData:
@@ -41,30 +41,37 @@ def extract_model_data(
     Extract data from an SDF/URDF resource useful to build a JaxSim model.
 
     Args:
-        model_description: A path to an SDF/URDF file, a string containing its content,
-          or a pre-parsed/pre-built rod model.
+        model_description:
+            A path to an SDF/URDF file, a string containing its content, or
+            a pre-parsed/pre-built rod model.
         model_name: The name of the model to extract from the SDF resource.
-        is_urdf: Whether the SDF resource is a URDF file. Needed only if model_description
-            is a URDF string.
+        is_urdf:
+            Whether to force parsing the resource as a URDF file. Automatically
+            detected if not provided.
 
     Returns:
         The extracted model data.
     """
 
-    if isinstance(model_description, rod.Model):
-        sdf_model = model_description
-    else:
-        # Parse the SDF resource.
-        sdf_element = rod.Sdf.load(sdf=model_description, is_urdf=is_urdf)
-
-        if len(sdf_element.models()) == 0:
-            raise RuntimeError("Failed to find any model in SDF resource")
-
-        # Assume the SDF resource has only one model, or the desired model name is given.
-        sdf_models = {m.name: m for m in sdf_element.models()}
-        sdf_model = (
-            sdf_element.models()[0] if len(sdf_models) == 1 else sdf_models[model_name]
-        )
+    match model_description:
+        case rod.Model():
+            sdf_model = model_description
+        case rod.Sdf() | str() | pathlib.Path():
+            sdf_element = (
+                model_description
+                if isinstance(model_description, rod.Sdf)
+                else rod.Sdf.load(sdf=model_description, is_urdf=is_urdf)
+            )
+            if not sdf_element.models():
+                raise RuntimeError("Failed to find any model in SDF resource")
+
+            # Assume the SDF resource has only one model, or the desired model name is given.
+            sdf_models = {m.name: m for m in sdf_element.models()}
+            sdf_model = (
+                sdf_element.models()[0]
+                if len(sdf_models) == 1
+                else sdf_models[model_name]
+            )
 
     # Log model name.
     logging.debug(msg=f"Found model '{sdf_model.name}' in SDF resource")
@@ -344,7 +351,7 @@ def extract_model_data(
 
 def build_model_description(
     model_description: pathlib.Path | str | rod.Model,
-    is_urdf: bool | None = False,
+    is_urdf: bool | None = None,
 ) -> descriptions.ModelDescription:
     """
     Builds a model description from an SDF/URDF resource.
@@ -352,8 +359,9 @@ def build_model_description(
     Args:
         model_description: A path to an SDF/URDF file, a string containing its content,
           or a pre-parsed/pre-built rod model.
-        is_urdf: Whether the SDF resource is a URDF file. Needed only if model_description
-            is a URDF string.
+        is_urdf: Whether the force parsing the resource as a URDF file. Automatically
+            detected if not provided.
+
     Returns:
         The parsed model description.
     """

jaxsim-0.4.3.dev115/src/jaxsim/rbda/contacts/__init__.py

@@ -0,0 +1,8 @@
+from .common import ContactModel, ContactsParams, ContactsState
+from .relaxed_rigid import (
+    RelaxedRigidContacts,
+    RelaxedRigidContactsParams,
+    RelaxedRigidContactsState,
+)
+from .rigid import RigidContacts, RigidContactsParams, RigidContactsState
+from .soft import SoftContacts, SoftContactsParams, SoftContactsState

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/rbda/contacts/common.py

@@ -7,15 +7,20 @@ import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass
 
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
 
-class ContactsState(abc.ABC):
+
+class ContactsState(JaxsimDataclass):
     """
     Abstract class storing the state of the contacts model.
     """
 
     @classmethod
     @abc.abstractmethod
-    def build(cls, **kwargs) -> ContactsState:
+    def build(cls: type[Self], **kwargs) -> Self:
         """
         Build the contact state object.
 
@@ -26,7 +31,7 @@ class ContactsState(abc.ABC):
 
     @classmethod
     @abc.abstractmethod
-    def zero(cls, **kwargs) -> ContactsState:
+    def zero(cls: type[Self], **kwargs) -> Self:
         """
         Build a zero contact state.
 
@@ -36,7 +41,7 @@ class ContactsState(abc.ABC):
         pass
 
     @abc.abstractmethod
-    def valid(self, **kwargs) -> bool:
+    def valid(self, **kwargs) -> jtp.BoolLike:
         """
         Check if the contacts state is valid.
         """
@@ -50,18 +55,20 @@ class ContactsParams(JaxsimDataclass):
 
     @classmethod
     @abc.abstractmethod
-    def build(cls) -> ContactsParams:
+    def build(cls: type[Self], **kwargs) -> Self:
         """
         Create a `ContactsParams` instance with specified parameters.
+
         Returns:
             The `ContactsParams` instance.
         """
         pass
 
     @abc.abstractmethod
-    def valid(self, *args, **kwargs) -> bool:
+    def valid(self, **kwargs) -> jtp.BoolLike:
         """
         Check if the parameters are valid.
+
         Returns:
             True if the parameters are valid, False otherwise.
         """
@@ -83,8 +90,8 @@ class ContactModel(JaxsimDataclass):
     @abc.abstractmethod
     def compute_contact_forces(
         self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
         **kwargs,
     ) -> tuple[jtp.Vector, tuple[Any, ...]]:
         """

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -16,6 +16,11 @@ from jaxsim.terrain.terrain import FlatTerrain, Terrain
 
 from .common import ContactModel, ContactsParams, ContactsState
 
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
 
 @jax_dataclasses.pytree_dataclass
 class RelaxedRigidContactsParams(ContactsParams):
@@ -106,7 +111,8 @@ class RelaxedRigidContactsParams(ContactsParams):
 
     @classmethod
     def build(
-        cls,
+        cls: type[Self],
+        *,
         time_constant: jtp.FloatLike | None = None,
         damping_coefficient: jtp.FloatLike | None = None,
         d_min: jtp.FloatLike | None = None,
@@ -119,7 +125,7 @@ class RelaxedRigidContactsParams(ContactsParams):
         mu: jtp.FloatLike | None = None,
         max_iterations: jtp.IntLike | None = None,
         tolerance: jtp.FloatLike | None = None,
-    ) -> RelaxedRigidContactsParams:
+    ) -> Self:
         """Create a `RelaxedRigidContactsParams` instance"""
 
         return cls(
@@ -132,7 +138,8 @@ class RelaxedRigidContactsParams(ContactsParams):
             }
         )
 
-    def valid(self) -> bool:
+    def valid(self) -> jtp.BoolLike:
+
         return bool(
             jnp.all(self.time_constant >= 0.0)
             and jnp.all(self.damping_coefficient > 0.0)
@@ -155,18 +162,19 @@ class RelaxedRigidContactsState(ContactsState):
     def __eq__(self, other: RelaxedRigidContactsState) -> bool:
         return hash(self) == hash(other)
 
-    @staticmethod
-    def build() -> RelaxedRigidContactsState:
+    @classmethod
+    def build(cls: type[Self]) -> Self:
         """Create a `RelaxedRigidContactsState` instance"""
 
-        return RelaxedRigidContactsState()
+        return cls()
 
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> RelaxedRigidContactsState:
+    @classmethod
+    def zero(cls: type[Self]) -> Self:
         """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
-        return RelaxedRigidContactsState.build()
 
-    def valid(self, model: js.model.JaxSimModel) -> bool:
+        return cls.build()
+
+    def valid(self, *, model: js.model.JaxSimModel) -> jtp.BoolLike:
         return True
 
 
@@ -182,10 +190,12 @@ class RelaxedRigidContacts(ContactModel):
         default_factory=FlatTerrain
     )
 
+    @jax.jit
     def compute_contact_forces(
         self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        *,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         link_forces: jtp.MatrixLike | None = None,

{jaxsim-0.4.3.dev100 → jaxsim-0.4.3.dev115}/src/jaxsim/rbda/contacts/rigid.py

@@ -14,6 +14,11 @@ from jaxsim.terrain import FlatTerrain, Terrain
 
 from .common import ContactModel, ContactsParams, ContactsState
 
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
 
 @jax_dataclasses.pytree_dataclass
 class RigidContactsParams(ContactsParams):
@@ -50,19 +55,22 @@ class RigidContactsParams(ContactsParams):
 
     @classmethod
     def build(
-        cls,
+        cls: type[Self],
+        *,
         mu: jtp.FloatLike | None = None,
         K: jtp.FloatLike | None = None,
         D: jtp.FloatLike | None = None,
-    ) -> RigidContactsParams:
+    ) -> Self:
         """Create a `RigidContactParams` instance"""
-        return RigidContactsParams(
+
+        return cls(
             mu=mu or cls.__dataclass_fields__["mu"].default,
             K=K or cls.__dataclass_fields__["K"].default,
             D=D or cls.__dataclass_fields__["D"].default,
         )
 
-    def valid(self) -> bool:
+    def valid(self) -> jtp.BoolLike:
+
         return bool(
             jnp.all(self.mu >= 0.0)
             and jnp.all(self.K >= 0.0)
@@ -77,18 +85,19 @@ class RigidContactsState(ContactsState):
     def __eq__(self, other: RigidContactsState) -> bool:
         return hash(self) == hash(other)
 
-    @staticmethod
-    def build(**kwargs) -> RigidContactsState:
+    @classmethod
+    def build(cls: type[Self]) -> Self:
         """Create a `RigidContactsState` instance"""
 
-        return RigidContactsState()
+        return cls()
 
-    @staticmethod
-    def zero(**kwargs) -> RigidContactsState:
+    @classmethod
+    def zero(cls: type[Self]) -> Self:
         """Build a zero `RigidContactsState` instance from a `JaxSimModel`."""
-        return RigidContactsState.build()
 
-    def valid(self, **kwargs) -> bool:
+        return cls.build()
+
+    def valid(self) -> jtp.BoolLike:
         return True
 
 
@@ -117,7 +126,8 @@ class RigidContacts(ContactModel):
             terrain_height: The height of the terrain at the collidable point position.
 
         Returns:
-            A tuple containing the activation state of the collidable points and the contact penetration depth h.
+            A tuple containing the activation state of the collidable points
+            and the contact penetration depth h.
         """
 
         # TODO: reduce code duplication with js.contact.in_contact
@@ -154,8 +164,8 @@ class RigidContacts(ContactModel):
 
         Args:
             inactive_collidable_points: The activation state of the collidable points.
-            M: The mass matrix of the system.
-            J_WC: The Jacobian matrix of the collidable points.
+            M: The mass matrix of the system (in mixed representation).
+            J_WC: The Jacobian matrix of the collidable points (in mixed representation).
             data: The `JaxSimModelData` instance.
         """
 
@@ -206,10 +216,12 @@ class RigidContacts(ContactModel):
 
         return BW_ν_post_impact
 
+    @jax.jit
     def compute_contact_forces(
         self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        *,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         link_forces: jtp.MatrixLike | None = None,