jaxsim 0.4.1.dev11__py3-none-any.whl → 0.4.1.dev15__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.1.dev11'
-__version_tuple__ = version_tuple = (0, 4, 1, 'dev11')
+__version__ = version = '0.4.1.dev15'
+__version_tuple__ = version_tuple = (0, 4, 1, 'dev15')

jaxsim/integrators/common.py

@@ -5,11 +5,12 @@ from typing import Any, ClassVar, Generic, Protocol, Type, TypeVar
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.api as js
+import jaxsim.math
 import jaxsim.typing as jtp
+from jaxsim import exceptions
 from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
 
 try:
@@ -539,48 +540,38 @@ class ExplicitRungeKuttaSO3Mixin:
     `PyTreeType = ODEState` to integrate the quaternion on SO(3).
     """
 
-    @classmethod
-    def integrate_rk_stage(
-        cls, x0: js.ode_data.ODEState, t0: Time, dt: TimeStep, k: js.ode_data.ODEState
-    ) -> js.ode_data.ODEState:
-
-        op = lambda x0_leaf, k_leaf: x0_leaf + dt * k_leaf
-        xf: js.ode_data.ODEState = jax.tree_util.tree_map(op, x0, k)
-
-        W_Q_B_tf = xf.physics_model.base_quaternion
-
-        return xf.replace(
-            physics_model=xf.physics_model.replace(
-                base_quaternion=W_Q_B_tf / jnp.linalg.norm(W_Q_B_tf)
-            )
-        )
-
     @classmethod
     def post_process_state(
         cls, x0: js.ode_data.ODEState, t0: Time, xf: js.ode_data.ODEState, dt: TimeStep
     ) -> js.ode_data.ODEState:
 
-        # Indices to convert quaternions between serializations.
-        to_xyzw = jnp.array([1, 2, 3, 0])
+        # Extract the initial base quaternion.
+        W_Q_B_t0 = x0.physics_model.base_quaternion
 
-        # Get the initial rotation.
-        W_R_B_t0 = jaxlie.SO3.from_quaternion_xyzw(
-            xyzw=x0.physics_model.base_quaternion[to_xyzw]
+        # We assume that the initial quaternion is already unary.
+        exceptions.raise_runtime_error_if(
+            condition=jnp.logical_not(jnp.allclose(W_Q_B_t0.dot(W_Q_B_t0), 1.0)),
+            msg="The SO(3) integrator received a quaternion at t0 that is not unary.",
         )
 
-        # Get the final angular velocity.
-        # This is already computed by averaging the kᵢ in RK-based schemes.
-        # Therefore, by using the ω at tf, we obtain a RK scheme operating
-        # on the SO(3) manifold.
-        W_ω_WB_tf = xf.physics_model.base_angular_velocity
-
-        # Integrate the orientation on SO(3).
-        # Note that we left-multiply with the exponential map since the angular
-        # velocity is expressed in the inertial frame.
-        W_R_B_tf = jaxlie.SO3.exp(tangent=dt * W_ω_WB_tf) @ W_R_B_t0
+        # Get the angular velocity ω to integrate the quaternion.
+        # This velocity ω[t0] is computed in the previous timestep by averaging the kᵢ
+        # corresponding to the active RK-based scheme. Therefore, by using the ω[t0],
+        # we obtain an explicit RK scheme operating on the SO(3) manifold.
+        # Note that the current integrator is not a semi-implicit scheme, therefore
+        # using the final ω[tf] would be not correct.
+        W_ω_WB_t0 = x0.physics_model.base_angular_velocity
+
+        # Integrate the quaternion on SO(3).
+        W_Q_B_tf = jaxsim.math.Quaternion.integration(
+            quaternion=W_Q_B_t0,
+            dt=dt,
+            omega=W_ω_WB_t0,
+            omega_in_body_fixed=False,
+        )
 
         # Replace the quaternion in the final state.
         return xf.replace(
-            physics_model=xf.physics_model.replace(base_quaternion=W_R_B_tf.wxyz),
+            physics_model=xf.physics_model.replace(base_quaternion=W_Q_B_tf),
             validate=True,
         )

jaxsim/rbda/utils.py

@@ -2,6 +2,7 @@ import jax.numpy as jnp
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim import exceptions
 from jaxsim.math import StandardGravity
 
 
@@ -131,6 +132,13 @@ def process_inputs(
     if W_Q_B.shape != (4,):
         raise ValueError(W_Q_B.shape, (4,))
 
+    # Check that the quaternion is unary since our RBDAs make this assumption in order
+    # to prevent introducing additional normalizations that would affect AD.
+    exceptions.raise_value_error_if(
+        condition=jnp.logical_not(jnp.allclose(W_Q_B.dot(W_Q_B), 1.0)),
+        msg="A RBDA received a quaternion that is not normalized.",
+    )
+
     # Pack the 6D base velocity and acceleration.
     W_v_WB = jnp.hstack([W_vl_WB, W_ω_WB])
     W_v̇_WB = jnp.hstack([W_v̇l_WB, W_ω̇_WB])

{jaxsim-0.4.1.dev11.dist-info → jaxsim-0.4.1.dev15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.1.dev11
+Version: 0.4.1.dev15
 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>

{jaxsim-0.4.1.dev11.dist-info → jaxsim-0.4.1.dev15.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
-jaxsim/_version.py,sha256=7ZFpPaz7M9u-cSnMJ8POHhkQ3_j7elMp-k9w58-1Rgw,426
+jaxsim/_version.py,sha256=GGa2wlqOOTBXMs_fQ4BwE7QCLSYjgvCVY3t6lhcCcz4,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
@@ -17,7 +17,7 @@ jaxsim/api/ode.py,sha256=NnLTBvpaT4kXnbjAghXIzLv9DTMJ8bele2iOlUQDv3Q,11028
 jaxsim/api/ode_data.py,sha256=9YZX-SK_KJtoIqG-zYWZsQInb2NA_LtxDn-jtLqm_3U,19759
 jaxsim/api/references.py,sha256=UA6kSQVBoq-bXSo99EOELf-_MD5MTy2zS0GtG3wQ410,16618
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=JXJECMkE-ZSmRt0-5koAw1vnmzkFewZ0aBrnP0bTpZY,20260
+jaxsim/integrators/common.py,sha256=iwFykYZxdchqJcmcx8MFWEVijS5Hx9wCNKLKAJdF4gE,20103
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
 jaxsim/integrators/variable_step.py,sha256=0FCmAZIFnhvQxVbAzNfZgCWN1yMRTGVdBm9UwwaXI1o,21280
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
@@ -51,7 +51,7 @@ jaxsim/rbda/crba.py,sha256=NhtZO48OUKKor7ddY7mB7h7a6idrmOyf0Vy4p7UCCgI,4724
 jaxsim/rbda/forward_kinematics.py,sha256=OEQYovnLKsWphUKhigmWa_384LwZW3Csp0MKufw4e1M,3415
 jaxsim/rbda/jacobian.py,sha256=I6mrlkk7Cpq3CE7k_tajOHCbT6vf2pW6vMS0TKNCnng,10725
 jaxsim/rbda/rnea.py,sha256=UrhcL93fp3pAKlGxOPS6X47L0ferH50bcSMzG55t4zY,7626
-jaxsim/rbda/utils.py,sha256=dPRWG8pV7rm3VivP09NE2ttX6wHCRNir4wuixGA_G2Y,5003
+jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
 jaxsim/rbda/contacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/rbda/contacts/common.py,sha256=iMKLP30Qft9eGTiHo2iY-UoACJjg1JphA9_pW8wRdjc,2410
 jaxsim/rbda/contacts/soft.py,sha256=3cDynim_tIgcbzRuqpHN82v4ELlxxK6lR-PG0haSK7Q,15660
@@ -61,8 +61,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=fLl1tY3DDb3lpIhG6BPqA5W34hM84oFzL-5cuz8k-68,11379
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=GOJQCJc5zwzoEGZB62wnWWGvUUQlXvDxz_A2Q-hFv7c,4027
-jaxsim-0.4.1.dev11.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.1.dev11.dist-info/METADATA,sha256=RQcmfKhS454m3fgEmDhJP00XtQFwnxWeXaBkTWiUKaQ,16779
-jaxsim-0.4.1.dev11.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-jaxsim-0.4.1.dev11.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.1.dev11.dist-info/RECORD,,
+jaxsim-0.4.1.dev15.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.1.dev15.dist-info/METADATA,sha256=IbvjxjjI0ys4OxT59PyJKmN5Y5xaTDKPT41Q3PTcmbY,16779
+jaxsim-0.4.1.dev15.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+jaxsim-0.4.1.dev15.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.1.dev15.dist-info/RECORD,,