jaxsim 0.6.2.dev233__tar.gz → 0.6.2.dev238__tar.gz

{jaxsim-0.6.2.dev233 → jaxsim-0.6.2.dev238}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev233
+Version: 0.6.2.dev238
 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.6.2.dev233 → jaxsim-0.6.2.dev238}/src/jaxsim/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.2.dev233'
-__version_tuple__ = version_tuple = (0, 6, 2, 'dev233')
+__version__ = version = '0.6.2.dev238'
+__version_tuple__ = version_tuple = (0, 6, 2, 'dev238')

{jaxsim-0.6.2.dev233 → jaxsim-0.6.2.dev238}/src/jaxsim/api/data.py

@@ -60,7 +60,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     _link_velocities: jtp.Matrix = dataclasses.field(repr=False, default=None)
 
     # Extended state for soft and rigid contact models.
-    contact_state: dict[str, jtp.Array] = dataclasses.field(default=None)
+    contact_state: dict[str, jtp.Array] = dataclasses.field(default_factory=dict)
 
     @staticmethod
     def build(
@@ -174,7 +174,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         contact_state = contact_state or {}
 
         if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
-            contact_state.setdefault(
+            contact_state["tangential_deformation"] = contact_state.get(
                 "tangential_deformation",
                 jnp.zeros_like(model.kin_dyn_parameters.contact_parameters.point),
             )
@@ -420,11 +420,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         Replace the attributes of the `JaxSimModelData` object.
         """
 
-        # Extract the batch size.
-        batch_size = (
-            self._base_transform.shape[0] if self._base_transform.ndim > 2 else 1
-        )
-
         if joint_positions is None:
             joint_positions = self.joint_positions
         if joint_velocities is None:
@@ -437,10 +432,11 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             contact_state = self.contact_state
 
         if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
-            contact_state.setdefault(
-                "tangential_deformation",
-                jnp.zeros_like(model.kin_dyn_parameters.contact_parameters.point),
-            )
+            contact_state = {
+                "tangential_deformation": jnp.zeros_like(
+                    contact_state["tangential_deformation"]
+                )
+            }
 
         # Normalize the quaternion to avoid numerical issues.
         base_quaternion_norm = jaxsim.math.safe_norm(
@@ -450,20 +446,18 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_quaternion_norm == 0, 1.0, base_quaternion_norm
         )
 
-        joint_positions = jnp.atleast_1d(joint_positions.squeeze()).astype(float)
-        joint_velocities = jnp.atleast_1d(joint_velocities.squeeze()).astype(float)
-        base_quaternion = jnp.atleast_1d(base_quaternion.squeeze()).astype(float)
-        base_position = jnp.atleast_1d(base_position.squeeze()).astype(float)
+        joint_positions = jnp.atleast_2d(joint_positions.squeeze()).astype(float)
+        joint_velocities = jnp.atleast_2d(joint_velocities.squeeze()).astype(float)
+        base_quaternion = jnp.atleast_2d(base_quaternion.squeeze()).astype(float)
+        base_position = jnp.atleast_2d(base_position.squeeze()).astype(float)
 
         base_transform = jaxsim.math.Transform.from_quaternion_and_translation(
             translation=base_position, quaternion=base_quaternion
-        )
+        ).reshape((-1, 4, 4))
 
         joint_transforms = jax.vmap(model.kin_dyn_parameters.joint_transforms)(
-            joint_positions=jnp.broadcast_to(
-                joint_positions, (batch_size, model.dofs())
-            ),
-            base_transform=jnp.broadcast_to(base_transform, (batch_size, 4, 4)),
+            joint_positions=joint_positions,
+            base_transform=base_transform,
         )
 
         if base_linear_velocity is None and base_angular_velocity is None:
@@ -494,27 +488,31 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             jaxsim.rbda.forward_kinematics_model, in_axes=(None,)
         )(
             model,
-            base_position=jnp.broadcast_to(base_position, (batch_size, 3)),
-            base_quaternion=jnp.broadcast_to(base_quaternion, (batch_size, 4)),
-            joint_positions=jnp.broadcast_to(
-                joint_positions, (batch_size, model.dofs())
-            ),
-            joint_velocities=jnp.broadcast_to(
-                joint_velocities, (batch_size, model.dofs())
-            ),
-            base_linear_velocity_inertial=jnp.broadcast_to(
-                base_linear_velocity_inertial, (batch_size, 3)
-            ),
-            base_angular_velocity_inertial=jnp.broadcast_to(
-                base_angular_velocity_inertial, (batch_size, 3)
+            base_position=base_position,
+            base_quaternion=base_quaternion,
+            joint_positions=joint_positions,
+            joint_velocities=joint_velocities,
+            base_linear_velocity_inertial=jnp.atleast_2d(base_linear_velocity_inertial),
+            base_angular_velocity_inertial=jnp.atleast_2d(
+                base_angular_velocity_inertial
             ),
         )
 
         # Adjust the output shapes.
-        if batch_size == 1:
-            link_transforms = link_transforms.reshape(self._link_transforms.shape)
-            link_velocities = link_velocities.reshape(self._link_velocities.shape)
-            joint_transforms = joint_transforms.reshape(self._joint_transforms.shape)
+        joint_positions = joint_positions.reshape(self._joint_positions.shape)
+        joint_velocities = joint_velocities.reshape(self._joint_velocities.shape)
+        base_quaternion = base_quaternion.reshape(self._base_quaternion.shape)
+        base_linear_velocity_inertial = base_linear_velocity_inertial.reshape(
+            self._base_linear_velocity.shape
+        )
+        base_angular_velocity_inertial = base_angular_velocity_inertial.reshape(
+            self._base_angular_velocity.shape
+        )
+        base_position = base_position.reshape(self._base_position.shape)
+        base_transform = base_transform.reshape(self._base_transform.shape)
+        joint_transforms = joint_transforms.reshape(self._joint_transforms.shape)
+        link_transforms = link_transforms.reshape(self._link_transforms.shape)
+        link_velocities = link_velocities.reshape(self._link_velocities.shape)
 
         return super().replace(
             _joint_positions=joint_positions,

{jaxsim-0.6.2.dev233 → jaxsim-0.6.2.dev238}/src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -477,8 +477,7 @@ class RelaxedRigidContacts(common.ContactModel):
         tol = solver_options.pop("tol")
         maxiter = solver_options.pop("maxiter")
 
-        # Compute the 3D linear force in C[W] frame.
-        solution, _ = run_optimization(
+        solve_fn = lambda *_: run_optimization(
             init_params=init_params,
             fun=objective,
             opt=optax.lbfgs(**solver_options),
@@ -486,6 +485,15 @@ class RelaxedRigidContacts(common.ContactModel):
             maxiter=maxiter,
         )
 
+        # Compute the 3D linear force in C[W] frame.
+        solution, _ = jax.lax.custom_linear_solve(
+            lambda x: A @ x,
+            -b,
+            solve=solve_fn,
+            symmetric=True,
+            has_aux=True,
+        )
+
         # Reshape the optimized solution to be a matrix of 3D contact forces.
         CW_fl_C = solution.reshape(-1, 3)
 

{jaxsim-0.6.2.dev233 → jaxsim-0.6.2.dev238}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev233
+Version: 0.6.2.dev238
 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.6.2.dev233 → jaxsim-0.6.2.dev238}/tests/test_automatic_differentiation.py

@@ -3,8 +3,6 @@ import os
 import jax
 import jax.numpy as jnp
 import numpy as np
-import optax
-import pytest
 from jax.test_util import check_grads
 
 import jaxsim.api as js
@@ -350,11 +348,6 @@ def test_ad_integration(
 
     model = jaxsim_models_types
 
-    # TODO: Remove when https://github.com/google-deepmind/optax/pull/1190 is included in a release.
-    # Skip if `optax` version is less or equal to "0.2.4" and the model is ergoCub.
-    if model.name() == "ergoCub" and optax.__version__ <= "0.2.4":
-        pytest.skip("Skipping ergoCub model with optax version <= 0.2.4.")
-
     _, subkey = jax.random.split(prng_key, num=2)
     data, references = get_random_data_and_references(
         model=model, velocity_representation=VelRepr.Inertial, key=subkey
@@ -416,13 +409,11 @@ def test_ad_integration(
         return xf_W_p_B, xf_W_Q_B, xf_s, xf_W_v_WB, xf_ṡ
 
     # Check derivatives against finite differences.
-    # We set forward mode only because the backward mode is not supported by the
-    # current implementation of `optax` optimizers in the relaxed rigid contact model.
     check_grads(
         f=step,
         args=(W_p_B, W_Q_B, s, W_v_WB, ṡ, τ, W_f_L),
         order=AD_ORDER,
-        modes=["fwd"],
+        modes=["fwd", "rev"],
         eps=ε,
     )