jaxsim 0.6.2.dev253__py3-none-any.whl → 0.6.2.dev264__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.dev253'
-__version_tuple__ = version_tuple = (0, 6, 2, 'dev253')
+__version__ = version = '0.6.2.dev264'
+__version_tuple__ = version_tuple = (0, 6, 2, 'dev264')

jaxsim/api/data.py

@@ -431,13 +431,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         if contact_state is None:
             contact_state = self.contact_state
 
-        if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
-            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(
             base_quaternion, axis=-1, keepdims=True
@@ -525,6 +518,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             _joint_transforms=joint_transforms,
             _link_transforms=link_transforms,
             _link_velocities=link_velocities,
+            contact_state=contact_state,
             validate=validate,
         )
 

jaxsim/api/integrators.py

@@ -156,9 +156,117 @@ def rk4_integration(
     return data_tf.replace(model=model)
 
 
+def rk4fast_integration(
+    model: js.model.JaxSimModel,
+    data: JaxSimModelData,
+    link_forces: jtp.Vector,
+    joint_torques: jtp.Vector,
+) -> JaxSimModelData:
+    """
+    Integrate the system state using the Runge-Kutta 4 fast method.
+
+    Note:
+        This method is a faster version of the RK4 method, but it may not be as accurate.
+        It computes the contact forces only once at the beginning of the integration step.
+    """
+
+    dt = model.time_step
+
+    if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
+
+        # Compute the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
+        # with the terrain.
+        W_f_L_terrain, contact_state_derivative = js.contact.link_contact_forces(
+            model=model,
+            data=data,
+            link_forces=link_forces,
+            joint_torques=joint_torques,
+        )
+
+    W_f_L_total = link_forces + W_f_L_terrain
+
+    # Update the contact state data. This is necessary only for the contact models
+    # that require propagation and integration of contact state.
+    contact_state = model.contact_model.update_contact_state(contact_state_derivative)
+
+    def f(x) -> dict[str, jtp.Matrix]:
+
+        with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
+
+            data_ti = data.replace(model=model, **x)
+
+            W_v̇_WB, s̈ = js.model.forward_dynamics_aba(
+                model=model,
+                data=data_ti,
+                joint_forces=joint_torques,
+                link_forces=W_f_L_total,
+            )
+
+            W_ṗ_B, W_Q̇_B, ṡ = js.ode.system_position_dynamics(
+                data=data,
+                baumgarte_quaternion_regularization=1.0,
+            )
+
+        return dict(
+            base_position=W_ṗ_B,
+            base_quaternion=W_Q̇_B,
+            joint_positions=ṡ,
+            base_linear_velocity=W_v̇_WB[0:3],
+            base_angular_velocity=W_v̇_WB[3:6],
+            joint_velocities=s̈,
+            # The contact state is not updated here, as it is assumed to be constant.
+            contact_state=data_ti.contact_state,
+        )
+
+    base_quaternion_norm = jaxsim.math.safe_norm(data._base_quaternion, axis=-1)
+    base_quaternion = data._base_quaternion / jnp.where(
+        base_quaternion_norm == 0, 1.0, base_quaternion_norm
+    )
+
+    x_t0 = dict(
+        base_position=data._base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=data._joint_positions,
+        base_linear_velocity=data._base_linear_velocity,
+        base_angular_velocity=data._base_angular_velocity,
+        joint_velocities=data._joint_velocities,
+        contact_state=contact_state,
+    )
+
+    euler_mid = lambda x, dxdt: x + (0.5 * dt) * dxdt
+    euler_fin = lambda x, dxdt: x + dt * dxdt
+
+    k1 = f(x_t0)
+    k2 = f(jax.tree.map(euler_mid, x_t0, k1))
+    k3 = f(jax.tree.map(euler_mid, x_t0, k2))
+    k4 = f(jax.tree.map(euler_fin, x_t0, k3))
+
+    # Average the slopes and compute the RK4 state derivative.
+    average = lambda k1, k2, k3, k4: (k1 + 2 * k2 + 2 * k3 + k4) / 6
+
+    dxdt = jax.tree_util.tree_map(average, k1, k2, k3, k4)
+
+    # Integrate the dynamics
+    x_tf = jax.tree_util.tree_map(euler_fin, x_t0, dxdt)
+
+    data_tf = dataclasses.replace(
+        data,
+        _base_position=x_tf["base_position"],
+        _base_quaternion=x_tf["base_quaternion"],
+        _joint_positions=x_tf["joint_positions"],
+        _base_linear_velocity=x_tf["base_linear_velocity"],
+        _base_angular_velocity=x_tf["base_angular_velocity"],
+        _joint_velocities=x_tf["joint_velocities"],
+        contact_state=x_tf["contact_state"],
+    )
+
+    return data_tf.replace(model=model)
+
+
 _INTEGRATORS_MAP: dict[
     js.model.IntegratorType, Callable[..., js.data.JaxSimModelData]
 ] = {
     js.model.IntegratorType.SemiImplicitEuler: semi_implicit_euler_integration,
     js.model.IntegratorType.RungeKutta4: rk4_integration,
+    js.model.IntegratorType.RungeKutta4Fast: rk4fast_integration,
 }

jaxsim/api/model.py

@@ -29,6 +29,7 @@ class IntegratorType(enum.IntEnum):
 
     SemiImplicitEuler = enum.auto()
     RungeKutta4 = enum.auto()
+    RungeKutta4Fast = enum.auto()
 
 
 @jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)

jaxsim/math/transform.py

@@ -92,11 +92,4 @@ class Transform:
             The 4x4 inverse transformation matrix.
         """
 
-        A_H_B = jnp.reshape(transform, (-1, 4, 4))
-
-        return (
-            jaxlie.SE3.from_matrix(matrix=A_H_B)
-            .inverse()
-            .as_matrix()
-            .reshape(transform.shape[:-2] + (4, 4))
-        )
+        return jaxlie.SE3.from_matrix(matrix=transform).inverse().as_matrix()

jaxsim/rbda/aba.py

@@ -201,7 +201,7 @@ def aba(
 
         # Propagate them to the parent, handling the base link.
         def propagate(
-            MA_pA: tuple[jtp.Matrix, jtp.Matrix]
+            MA_pA: tuple[jtp.Matrix, jtp.Matrix],
         ) -> tuple[jtp.Matrix, jtp.Matrix]:
 
             MA, pA = MA_pA

jaxsim/rbda/contacts/rigid.py

@@ -357,7 +357,7 @@ class RigidContacts(ContactModel):
         #      G x ≤ h
         #
         # TODO: add possibility to notify if the QP problem did not converge.
-        solution, _, _, _, converged, _ = qpax.solve_qp(  # noqa: F841
+        solution, _, _, _, converged, _ = qpax.solve_qp(  # noqa: RUF059
             Q=Q, q=q, A=A, b=b, G=G, h=h_bounds, **self.solver_options
         )
 

{jaxsim-0.6.2.dev253.dist-info → jaxsim-0.6.2.dev264.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev253
+Version: 0.6.2.dev264
 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.dev253.dist-info → jaxsim-0.6.2.dev264.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=b8dYoVXqtHxHcF56iM2xgKk78lsvmGrfDlvdwaGasgs,3388
-jaxsim/_version.py,sha256=tso8wO5VrX4RPYa0ZLA_aUJ6ylncC96kilEqh4mUO5U,528
+jaxsim/_version.py,sha256=oLeHDsqsLO5VnFp4ZF6v3W6PavO6GfFRTiLMYRl3wOg,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,13 +8,13 @@ jaxsim/api/actuation_model.py,sha256=GgLi-yhEpsczwhMNYIlMvet8tirmmED6S7AumbSbk4U
 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=KBDN7zFrNDSkry3Pow1LdBZ6kmfwk7dR5jaLA6DNAVc,23468
+jaxsim/api/data.py,sha256=9pxug2gFIZPwqi9kNYXhEziA5IQBB9KNNwIfPfc_kAU,23249
 jaxsim/api/frame.py,sha256=4wg6GsyBQgYhSvc-ry_31JsaL66sZt3TtgwjB7NrHmk,14583
-jaxsim/api/integrators.py,sha256=DgOnzLepy45e-TM6Infk8qfPXn0r8GubCdJQZmNLP8w,5269
+jaxsim/api/integrators.py,sha256=sHdTWw2Z-Q7jggA8zRkA1KYYd4HNIozXPwNvFwt0YBs,9011
 jaxsim/api/joint.py,sha256=AnqlNWmBOay-gsoo0y4AbfFQ2OCJm-8T1E0IMhZeLoY,7457
 jaxsim/api/kin_dyn_parameters.py,sha256=v1ZDb-Eqn552Zb344sJYy9e5sQT--_SYGIdtNC2JVSg,29828
 jaxsim/api/link.py,sha256=bSZOYJDY9HJMgY25VzevTTsdFZTJc6yRRpslc5FhGHE,12782
-jaxsim/api/model.py,sha256=rZK8YigkM1pHxx0Aj9AJDpFma7I_pNgcs8lzPF8ukUk,69914
+jaxsim/api/model.py,sha256=Do5q7vSyml854wFyN21qJtNfJfFo7TlEfxx_FoTxOW0,69948
 jaxsim/api/ode.py,sha256=fp20_LK9lXw2DfNkQgrfJmtd_iBXDNzZkOn0u5Pm8Qw,6190
 jaxsim/api/references.py,sha256=-vd50y3v-jkXAsILS432etIKV6e2EUE2oOeLHuUrfuQ,20399
 jaxsim/math/__init__.py,sha256=dNozvtm8WsB7nxw4uK29yQQKPcDUEczr2zcHoZfMItc,383
@@ -25,7 +25,7 @@ jaxsim/math/joint_model.py,sha256=mV6F3SVHUxHcXMIw-hc4fG341R4rkK8MvRPn53HcEpg,69
 jaxsim/math/quaternion.py,sha256=MSaZywzJDxs2te1ZELeIcupKSFIA9q_pdXy7fDAEqM4,4539
 jaxsim/math/rotation.py,sha256=TEUtT3X2tFieNxdlccup1pfaTgCTtfX-hTNotd8-nNk,1892
 jaxsim/math/skew.py,sha256=z_9YN-NDHL3n4KXWNbzTSMkFDZ0SDpz4RUcwwYFOaao,1402
-jaxsim/math/transform.py,sha256=NOWBQgL_5hVr9aHA7j0Ta-Zq_EXAKguUAx78fVPZJE4,3342
+jaxsim/math/transform.py,sha256=d0_m_obmUOmnI8Bte0ktvibR9Hv9M9qpg8tVuLON2g0,3192
 jaxsim/math/utils.py,sha256=IiH01iN54BtLnULC04pDfYe8Av99p3FGdYp2jJInm30,1166
 jaxsim/mujoco/__init__.py,sha256=1kAWzYOS7nP29S5FGyWPqiAnPf4yPSoaPW-WBGBjVV0,214
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
@@ -45,7 +45,7 @@ jaxsim/parsers/rod/meshes.py,sha256=yAXefG73_zqbVKRUdlcz9yFmypjDIpiP9cO96PeAozE,
 jaxsim/parsers/rod/parser.py,sha256=vaOlvlvgQCEGbymFc21praHKEVq_LqqlKeRJNv7tZiE,14398
 jaxsim/parsers/rod/utils.py,sha256=wmD-wCF1lLO8pknX7A3a8CGt9wDlTS_xCqQulcZ_XlM,8242
 jaxsim/rbda/__init__.py,sha256=ksfupKZzeJyysxrbyMyEfszUdBH6LfCfkSz3KLfODhY,328
-jaxsim/rbda/aba.py,sha256=QiajR4Uv06u-BX5V2kWpmFmP_kIfTxxcycdJhIbid9A,9029
+jaxsim/rbda/aba.py,sha256=jxvss3XL8pBaT40bWG5pHcH9f1DDl3LoRqdnpFSlCWo,9030
 jaxsim/rbda/collidable_points.py,sha256=XyeV1I43GL22j03rkNVocaIPOGYirt3PiDHrFMndziQ,2070
 jaxsim/rbda/crba.py,sha256=DC9kBXMG1qXaoAdo8K7OCnVHT_YUaL_t6Li56sRf8ro,5093
 jaxsim/rbda/forward_kinematics.py,sha256=qem7Yp-B2oNVOsU3Q2CWV2tbfZKJOCAdDozFgaPB8tg,3838
@@ -57,7 +57,7 @@ jaxsim/rbda/actuation/common.py,sha256=aGFqO4VTgQLsTJyOtVuoa_otT_RbkckmG3rq7wjOy
 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=RjeLF06Pp19qio447U9z5EdhdM6nyMh-ISQX_2-vdaE,21349
-jaxsim/rbda/contacts/rigid.py,sha256=ctJ4_4qSHaqGxCBMTYhvWWm8D4hJ-YdDiorJuURiNWw,17601
+jaxsim/rbda/contacts/rigid.py,sha256=gHOpNQ_sk7nQcpGAE_nFOSq8fza2yFv_Mqqs-e118q8,17603
 jaxsim/rbda/contacts/soft.py,sha256=Ac9aWDdjAm55Mv9LLnEs3nj7hX_NvJMnicV35SbFLSY,15282
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=so8kqVsUlPXqOKQ_UaYW6HE4XS8nTZRdFdXQyhB2tG4,6725
@@ -65,8 +65,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.dev253.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.6.2.dev253.dist-info/METADATA,sha256=F_UAPp-S0HPegiLJOZ2yiiu8y1qQRmIIqzwYiiuiieM,19658
-jaxsim-0.6.2.dev253.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-jaxsim-0.6.2.dev253.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.6.2.dev253.dist-info/RECORD,,
+jaxsim-0.6.2.dev264.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.6.2.dev264.dist-info/METADATA,sha256=KJmAE7MbfhUFSsRrMFTrHJTD3GUdaUPXVMuEFWU-HiU,19658
+jaxsim-0.6.2.dev264.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+jaxsim-0.6.2.dev264.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.6.2.dev264.dist-info/RECORD,,