jaxsim 0.4.1.dev9__tar.gz → 0.4.1.dev15__tar.gz

{jaxsim-0.4.1.dev9 → jaxsim-0.4.1.dev15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.1.dev9
+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.dev9 → jaxsim-0.4.1.dev15}/src/jaxsim/_version.py

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

{jaxsim-0.4.1.dev9 → jaxsim-0.4.1.dev15}/src/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-0.4.1.dev9 → jaxsim-0.4.1.dev15}/src/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.dev9 → jaxsim-0.4.1.dev15}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.1.dev9
+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.dev9 → jaxsim-0.4.1.dev15}/tests/test_automatic_differentiation.py

@@ -93,7 +93,7 @@ def test_ad_aba(
     aba = lambda W_p_B, W_Q_B, s, W_v_WB, ṡ, τ, W_f_L, g: jaxsim.rbda.aba(
         model=model,
         base_position=W_p_B,
-        base_quaternion=W_Q_B,
+        base_quaternion=W_Q_B / jnp.linalg.norm(W_Q_B),
         joint_positions=s,
         base_linear_velocity=W_v_WB[0:3],
         base_angular_velocity=W_v_WB[3:6],
@@ -150,7 +150,7 @@ def test_ad_rnea(
     rnea = lambda W_p_B, W_Q_B, s, W_v_WB, ṡ, W_v̇_WB, s̈, W_f_L, g: jaxsim.rbda.rnea(
         model=model,
         base_position=W_p_B,
-        base_quaternion=W_Q_B,
+        base_quaternion=W_Q_B / jnp.linalg.norm(W_Q_B),
         joint_positions=s,
         base_linear_velocity=W_v_WB[0:3],
         base_angular_velocity=W_v_WB[3:6],
@@ -229,7 +229,7 @@ def test_ad_fk(
     fk = lambda W_p_B, W_Q_B, s: jaxsim.rbda.forward_kinematics_model(
         model=model,
         base_position=W_p_B,
-        base_quaternion=W_Q_B,
+        base_quaternion=W_Q_B / jnp.linalg.norm(W_Q_B),
         joint_positions=s,
     )
 

{jaxsim-0.4.1.dev9 → jaxsim-0.4.1.dev15}/tests/test_simulations.py

@@ -6,7 +6,7 @@ import jaxsim.api as js
 import jaxsim.integrators
 import jaxsim.rbda
 from jaxsim import VelRepr
-from jaxsim.rbda.contacts.soft import SoftContactsParams
+from jaxsim.utils import Mutability
 
 
 def test_box_with_external_forces(
@@ -102,23 +102,19 @@ def test_box_with_zero_gravity(
 
     model = jaxsim_model_box
 
+    # Move the terrain (almost) infinitely far away from the box.
+    with model.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
+        model.terrain = jaxsim.terrain.FlatTerrain.build(height=-1e9)
+
     # Split the PRNG key.
-    _, subkey, subkey2 = jax.random.split(prng_key, num=3)
+    _, subkey = jax.random.split(prng_key, num=2)
 
     # Build the data of the model.
     data0 = js.data.JaxSimModelData.build(
         model=model,
-        base_position=jax.random.uniform(subkey2, shape=(3,)),
+        base_position=jax.random.uniform(subkey, shape=(3,)),
         velocity_representation=velocity_representation,
         standard_gravity=0.0,
-        contacts_params=SoftContactsParams.build(K=0.0, D=0.0, mu=0.0),
-    )
-
-    # Generate a random linear force.
-    L_f = (
-        jax.random.uniform(subkey, shape=(model.number_of_links(), 6))
-        .at[:, 3:]
-        .set(jnp.zeros(3))
     )
 
     # Initialize a references object that simplifies handling external forces.
@@ -129,13 +125,29 @@ def test_box_with_zero_gravity(
     )
 
     # Apply a link forces to the base link.
-    references = references.apply_link_forces(
-        forces=jnp.atleast_2d(L_f),
-        link_names=model.link_names(),
-        model=model,
-        data=data0,
-        additive=False,
-    )
+    with references.switch_velocity_representation(jaxsim.VelRepr.Mixed):
+
+        # Generate a random linear force.
+        # We enforce them to be the same for all velocity representations so that
+        # we can compare their outcomes.
+        LW_f = 10.0 * (
+            jax.random.uniform(jax.random.key(0), shape=(model.number_of_links(), 6))
+            .at[:, 3:]
+            .set(jnp.zeros(3))
+        )
+
+        # Note that the context manager does not switch back the newly created
+        # `references` (that is not the yielded object) to the original representation.
+        # In the simulation loop below, we need to make sure that we switch both `data`
+        # and `references` to the same representation before extracting the information
+        # passed to the step function.
+        references = references.apply_link_forces(
+            forces=jnp.atleast_2d(LW_f),
+            link_names=model.link_names(),
+            model=model,
+            data=data0,
+            additive=False,
+        )
 
     # Create the integrator.
     integrator = jaxsim.integrators.fixed_step.RungeKutta4SO3.build(
@@ -145,8 +157,7 @@ def test_box_with_zero_gravity(
     )
 
     # Initialize the integrator.
-    tf = 1.0
-    dt = 0.010
+    tf, dt = 1.0, 0.010
     T = jnp.arange(start=0, stop=tf * 1e9, step=dt * 1e9, dtype=int)
     integrator_state = integrator.init(x0=data0.state, t0=0.0, dt=dt)
 
@@ -156,19 +167,28 @@ def test_box_with_zero_gravity(
     # ... and step the simulation.
     for t_ns in T:
 
-        data, integrator_state = js.model.step(
-            model=model,
-            data=data,
-            dt=dt,
-            integrator=integrator,
-            integrator_state=integrator_state,
-            link_forces=references.link_forces(model=model, data=data),
-        )
+        assert data.time() == t_ns / 1e9
+
+        with (
+            data.switch_velocity_representation(velocity_representation),
+            references.switch_velocity_representation(velocity_representation),
+        ):
+
+            data, integrator_state = js.model.step(
+                model=model,
+                data=data,
+                dt=dt,
+                integrator=integrator,
+                integrator_state=integrator_state,
+                link_forces=references.link_forces(model=model, data=data),
+            )
+
+    # Check the final simulation time.
+    assert data.time() == T[-1] / 1e9 + dt
 
     # Check that the box moved as expected.
-    assert data.time() == t_ns / 1e9 + dt
     assert data.base_position() == pytest.approx(
         data0.base_position()
-        + 0.5 * L_f[:, :3].squeeze() / js.model.total_mass(model=model) * tf**2,
+        + 0.5 * LW_f[:, :3].squeeze() / js.model.total_mass(model=model) * tf**2,
         abs=1e-3,
     )