jaxsim 0.6.2.dev233__py3-none-any.whl → 0.6.2.dev238__py3-none-any.whl

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/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/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.dist-info → jaxsim-0.6.2.dev238.dist-info}/METADATA

@@ -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.dist-info → jaxsim-0.6.2.dev238.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=b8dYoVXqtHxHcF56iM2xgKk78lsvmGrfDlvdwaGasgs,3388
-jaxsim/_version.py,sha256=sPJq_dpGjKmFrDc3jN3zh6Jvj0t6QyFFNuDkHRcOzlE,528
+jaxsim/_version.py,sha256=eVqBRpQ8sW7XIxtY3zbtRorKBdbdS2mv3DLVp8TmLjI,528
 jaxsim/exceptions.py,sha256=MQ3LRMfVMX2-g3qYj7mUVNV9OLlIA48TANJegbcQyXI,2641
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=7msl8t5Jt09RNYfKdPJtpjLfWurldcycDappb045Eso,761
@@ -8,7 +8,7 @@ jaxsim/api/actuation_model.py,sha256=L8AzxIiEquWeG8UGGJaYr2Alt4dkkOROlbsCn9hUYik
 jaxsim/api/com.py,sha256=47a9SSaXY540RCkVnHodwLNUrodIfJIkguIYdSEQVwQ,13697
 jaxsim/api/common.py,sha256=yTaRXDYkXmISBOhZ93f9TssR0p4wq7qj7B6OsvYzRME,6942
 jaxsim/api/contact.py,sha256=dlKKDQUG-KQ5qQaYBv2NmZLDb1OnJdltZv8MWXkD_W0,20969
-jaxsim/api/data.py,sha256=aW6buyH3YszzEBqBXqBgHf6uaSilZcqU7PyP60qgLco,23336
+jaxsim/api/data.py,sha256=KBDN7zFrNDSkry3Pow1LdBZ6kmfwk7dR5jaLA6DNAVc,23468
 jaxsim/api/frame.py,sha256=4wg6GsyBQgYhSvc-ry_31JsaL66sZt3TtgwjB7NrHmk,14583
 jaxsim/api/integrators.py,sha256=DgOnzLepy45e-TM6Infk8qfPXn0r8GubCdJQZmNLP8w,5269
 jaxsim/api/joint.py,sha256=AnqlNWmBOay-gsoo0y4AbfFQ2OCJm-8T1E0IMhZeLoY,7457
@@ -54,7 +54,7 @@ jaxsim/rbda/rnea.py,sha256=lMU7xxdPqGGzk0QwteB-IYjL4auHOpd78C1YqAXlp9s,7588
 jaxsim/rbda/utils.py,sha256=6JwEDQqLMsBX7CUmPYEhdPEscXmGbWVYg6xEriPOgvE,5587
 jaxsim/rbda/contacts/__init__.py,sha256=resrBkTdOA-1YMdcdUH2RATEhAf_Ye6MQNtjG3ClMYQ,371
 jaxsim/rbda/contacts/common.py,sha256=qVm3Ghoytg1HAeykNrYw5-4rQJ4Mv7h0Pk75ETzGXyc,9045
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=zoMbc1edTsbUBf3hVCUoJb-xTCrtQT4RXKaILFup3KI,21152
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=RjeLF06Pp19qio447U9z5EdhdM6nyMh-ISQX_2-vdaE,21349
 jaxsim/rbda/contacts/rigid.py,sha256=ctJ4_4qSHaqGxCBMTYhvWWm8D4hJ-YdDiorJuURiNWw,17601
 jaxsim/rbda/contacts/soft.py,sha256=Ac9aWDdjAm55Mv9LLnEs3nj7hX_NvJMnicV35SbFLSY,15282
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
@@ -63,8 +63,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=XzmZeIibcaOzaxpprsGSxH3UrM66PAO456rFV91sNXg,11453
 jaxsim/utils/tracing.py,sha256=Btwxdfhb7fJLk3r5PlQkGYj60Y2KbFT1gANGIA697FU,530
 jaxsim/utils/wrappers.py,sha256=3IMwydqFgmSPqeuUQ3PRmdhDc1IoT6XC23jPC_LjWXs,4175
-jaxsim-0.6.2.dev233.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.6.2.dev233.dist-info/METADATA,sha256=9kHFCW__ENh9MIjH0wu2LxivnchhSt1vXL4Obu_BUD4,19650
-jaxsim-0.6.2.dev233.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
-jaxsim-0.6.2.dev233.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.6.2.dev233.dist-info/RECORD,,
+jaxsim-0.6.2.dev238.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.6.2.dev238.dist-info/METADATA,sha256=RqW1Edan0yEe43Bp8uljqRkFdaTqvSffYxCg1bGKEjY,19650
+jaxsim-0.6.2.dev238.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+jaxsim-0.6.2.dev238.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.6.2.dev238.dist-info/RECORD,,

{jaxsim-0.6.2.dev233.dist-info → jaxsim-0.6.2.dev238.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.0.2)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any