jaxsim 0.2.dev191__py3-none-any.whl → 0.2.dev364__py3-none-any.whl

jaxsim/api/references.py

@@ -7,10 +7,11 @@ import jax.numpy as jnp
 import jax_dataclasses
 
 import jaxsim.api as js
-import jaxsim.physics.model.physics_model_state
 import jaxsim.typing as jtp
-from jaxsim import VelRepr
-from jaxsim.simulation.ode_data import ODEInput
+from jaxsim.utils.tracing import not_tracing
+
+from .common import VelRepr
+from .ode_data import ODEInput
 
 try:
     from typing import Self
@@ -95,7 +96,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # Create a zero references object.
         references = JaxSimModelReferences(
-            input=ODEInput.zero(physics_model=model.physics_model),
+            input=ODEInput.zero(model=model),
             velocity_representation=velocity_representation,
         )
 
@@ -132,7 +133,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         valid = True
 
         if model is not None:
-            valid = valid and self.input.valid(physics_model=model.physics_model)
+            valid = valid and self.input.valid(model=model)
 
         return valid
 
@@ -188,7 +189,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # If we have the model, we can extract the link names, if not provided.
         link_names = link_names if link_names is not None else model.link_names()
-        link_idxs = jaxsim.api.link.names_to_idxs(link_names=link_names, model=model)
+        link_idxs = js.link.names_to_idxs(link_names=link_names, model=model)
 
         # In inertial-fixed representation, we already have the link forces.
         if self.velocity_representation is VelRepr.Inertial:
@@ -198,7 +199,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             msg = "Missing model data to use a representation different from {}"
             raise ValueError(msg.format(VelRepr.Inertial.name))
 
-        if not data.valid(model=model):
+        if not_tracing(self.input.physics_model.f_ext) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
         # Helper function to convert a single 6D force to the active representation.
@@ -252,7 +253,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
             return self.input.physics_model.tau
 
-        if not self.valid(model=model):
+        if not_tracing(self.input.physics_model.tau) and not self.valid(model=model):
             msg = "The actuation object is not compatible with the provided model"
             raise ValueError(msg)
 
@@ -303,7 +304,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         if model is None:
             return replace(forces=forces)
 
-        if not self.valid(model=model):
+        if not_tracing(forces) and not self.valid(model=model):
             msg = "The references object is not compatible with the provided model"
             raise ValueError(msg)
 
@@ -379,7 +380,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # If we have the model, we can extract the link names if not provided.
         link_names = link_names if link_names is not None else model.link_names()
-        link_idxs = jaxsim.api.link.names_to_idxs(link_names=link_names, model=model)
+        link_idxs = js.link.names_to_idxs(link_names=link_names, model=model)
 
         # Compute the bias depending on whether we either set or add the link forces.
         W_f0_L = (
@@ -401,7 +402,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             msg = "Missing model data to use a representation different from {}"
             raise ValueError(msg.format(VelRepr.Inertial.name))
 
-        if not data.valid(model=model):
+        if not_tracing(forces) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
         # Helper function to convert a single 6D force to the inertial representation.

jaxsim/integrators/__init__.py

@@ -1,2 +1,2 @@
-from . import fixed_step
-from .common import Integrator, Time, TimeStep
+from . import fixed_step, variable_step
+from .common import Integrator, SystemDynamics, Time, TimeStep

jaxsim/integrators/common.py

@@ -5,9 +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.typing as jtp
+from jaxsim.math import Quaternion
 from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
 
 try:
@@ -46,6 +49,9 @@ class SystemDynamics(Protocol[State, StateDerivative]):
 @jax_dataclasses.pytree_dataclass
 class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
+    AfterInitKey: ClassVar[str] = "after_init"
+    InitializingKey: ClassVar[str] = "initializing"
+
     AuxDictDynamicsKey: ClassVar[str] = "aux_dict_dynamics"
 
     dynamics: Static[SystemDynamics[State, StateDerivative]] = dataclasses.field(
@@ -58,7 +64,10 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
     @classmethod
     def build(
-        cls: Type[Self], *, dynamics: SystemDynamics[State, StateDerivative], **kwargs
+        cls: Type[Self],
+        *,
+        dynamics: SystemDynamics[State, StateDerivative],
+        **kwargs,
     ) -> Self:
         """
         Build the integrator object.
@@ -102,9 +111,9 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         with integrator.mutable_context(mutability=Mutability.MUTABLE):
             xf = integrator(x0, t0, dt, **kwargs)
 
-        assert Integrator.AuxDictDynamicsKey in integrator.params
-
-        return xf, integrator.params
+        return xf, integrator.params | {
+            Integrator.AfterInitKey: jnp.array(False).astype(bool)
+        }
 
     @abc.abstractmethod
     def __call__(self, x0: State, t0: Time, dt: TimeStep, **kwargs) -> NextState:
@@ -116,7 +125,7 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         t0: Time,
         dt: TimeStep,
         *,
-        key: jax.Array | None = None,
+        include_dynamics_aux_dict: bool = False,
         **kwargs,
     ) -> dict[str, Any]:
         """
@@ -126,7 +135,6 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
             x0: The initial state of the system.
             t0: The initial time of the system.
             dt: The time step of the integration.
-            key: An optional random key to initialize the integrator.
 
         Returns:
             The auxiliary dictionary of the integrator.
@@ -141,17 +149,43 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
             the first step will be wrong.
         """
 
-        _, aux_dict_dynamics = self.dynamics(x0, t0)
-
         with self.editable(validate=False) as integrator:
+
+            # Initialize the integrator parameters.
+            # For initialization purpose, the integrators can check if the
+            # `Integrator.InitializingKey` is present in their parameters.
+            # The AfterInitKey is used in the first step after initialization.
+            integrator.params = {
+                Integrator.InitializingKey: jnp.array(True),
+                Integrator.AfterInitKey: jnp.array(False),
+            }
+
+            # Run a dummy call of the integrator.
+            # It is used only to get the params so that we know the structure
+            # of the corresponding pytree.
             _ = integrator(x0, t0, dt, **kwargs)
-            aux_dict_step = integrator.params
 
-        if Integrator.AuxDictDynamicsKey in aux_dict_dynamics:
-            msg = "You cannot create a key '{}' in the __call__ method."
-            raise KeyError(msg.format(Integrator.AuxDictDynamicsKey))
+        # Remove the injected key.
+        _ = integrator.params.pop(Integrator.InitializingKey)
 
-        return {Integrator.AuxDictDynamicsKey: aux_dict_dynamics} | aux_dict_step
+        # Make sure that all leafs of the dictionary are JAX arrays.
+        # Also, since these are dummy parameters, set them all to zero.
+        params_after_init = jax.tree_util.tree_map(
+            lambda l: jnp.zeros_like(l), integrator.params
+        )
+
+        # Mark the next step as first step after initialization.
+        params_after_init = params_after_init | {
+            Integrator.AfterInitKey: jnp.array(True)
+        }
+
+        # Store the zero parameters in the integrator.
+        # When the integrator is stepped, this is used to check if the passed
+        # parameters are valid.
+        with self.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
+            self.params = params_after_init
+
+        return params_after_init
 
 
 @jax_dataclasses.pytree_dataclass
@@ -209,20 +243,10 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             The integrator object.
         """
 
-        # Adjust the shape of the tableau coefficients.
-        c = jnp.atleast_1d(cls.c.squeeze())
-        b = jnp.atleast_2d(jnp.vstack(cls.b.squeeze()))
-        A = jnp.atleast_2d(cls.A.squeeze())
-
         # Check validity of the Butcher tableau.
-        if not ExplicitRungeKutta.butcher_tableau_is_valid(A=A, b=b, c=c):
+        if not ExplicitRungeKutta.butcher_tableau_is_valid(A=cls.A, b=cls.b, c=cls.c):
             raise ValueError("The Butcher tableau of this class is not valid.")
 
-        # Store the adjusted shapes of the tableau coefficients.
-        cls.c = c
-        cls.b = b
-        cls.A = A
-
         # Check that b.T has enough rows based on the configured index of the solution.
         if cls.row_index_of_solution >= cls.b.T.shape[0]:
             msg = "The index of the solution ({}-th row of `b.T`) is out of range ({})."
@@ -506,3 +530,65 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         # possibly intermediate kᵢ derivative).
         # Note that if multiple rows match (it should not), we return the first match.
         return True, int(jnp.where(rows_of_A_with_fsal == True)[0].tolist()[0])
+
+
+class ExplicitRungeKuttaSO3Mixin:
+    """
+    Mixin class to apply over explicit RK integrators defined on
+    `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_t0 = x0.physics_model.base_quaternion
+        W_ω_WB_t0 = x0.physics_model.base_angular_velocity
+
+        return xf.replace(
+            physics_model=xf.physics_model.replace(
+                base_quaternion=Quaternion.integration(
+                    quaternion=W_Q_B_t0,
+                    dt=dt,
+                    omega=W_ω_WB_t0,
+                    omega_in_body_fixed=False,
+                ),
+            )
+        )
+
+    @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])
+        to_wxyz = jnp.array([3, 0, 1, 2])
+
+        # Get the initial quaternion.
+        W_Q_B_t0 = jaxlie.SO3.from_quaternion_xyzw(
+            xyzw=x0.physics_model.base_quaternion[to_xyzw]
+        )
+
+        # 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 quaternion on SO(3).
+        # Note that we left-multiply with the exponential map since the angular
+        # velocity is expressed in the inertial frame.
+        W_Q_B_tf = jaxlie.SO3.exp(tangent=dt * W_ω_WB_tf) @ W_Q_B_t0
+
+        # Replace the quaternion in the final state.
+        return xf.replace(
+            physics_model=xf.physics_model.replace(
+                base_quaternion=W_Q_B_tf.as_quaternion_xyzw()[to_wxyz]
+            ),
+            validate=True,
+        )

jaxsim/integrators/fixed_step.py

@@ -3,14 +3,12 @@ from typing import ClassVar, Generic
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 
-from jaxsim.simulation.ode_data import ODEState
+import jaxsim.api as js
 
-from .common import ExplicitRungeKutta, PyTreeType, Time, TimeStep
-
-ODEStateDerivative = ODEState
+from .common import ExplicitRungeKutta, ExplicitRungeKuttaSO3Mixin, PyTreeType
 
+ODEStateDerivative = js.ode_data.ODEState
 
 # =====================================================
 # Explicit Runge-Kutta integrators operating on PyTrees
@@ -20,37 +18,23 @@ ODEStateDerivative = ODEState
 @jax_dataclasses.pytree_dataclass
 class ForwardEuler(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
-    A: ClassVar[jax.typing.ArrayLike] = jnp.array(
-        [
-            [0],
-        ]
-    ).astype(float)
+    A: ClassVar[jax.typing.ArrayLike] = jnp.atleast_2d(0).astype(float)
 
-    b: ClassVar[jax.typing.ArrayLike] = (
-        jnp.array(
-            [
-                [1],
-            ]
-        )
-        .astype(float)
-        .transpose()
-    )
+    b: ClassVar[jax.typing.ArrayLike] = jnp.atleast_2d(1).astype(float).transpose()
 
-    c: ClassVar[jax.typing.ArrayLike] = jnp.array(
-        [0],
-    ).astype(float)
+    c: ClassVar[jax.typing.ArrayLike] = jnp.atleast_1d(0).astype(float)
 
     row_index_of_solution: ClassVar[int] = 0
     order_of_bT_rows: ClassVar[tuple[int, ...]] = (1,)
 
 
 @jax_dataclasses.pytree_dataclass
-class Heun(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+class Heun2(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
     A: ClassVar[jax.typing.ArrayLike] = jnp.array(
         [
             [0, 0],
-            [1 / 2, 0],
+            [1, 0],
         ]
     ).astype(float)
 
@@ -103,56 +87,16 @@ class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 # ===============================================================================
 
 
-class ExplicitRungeKuttaSO3Mixin:
-    """
-    Mixin class to apply over explicit RK integrators defined on
-    `PyTreeType = ODEState` to integrate the quaternion on SO(3).
-    """
-
-    @classmethod
-    def post_process_state(
-        cls, x0: ODEState, t0: Time, xf: ODEState, dt: TimeStep
-    ) -> ODEState:
-
-        # Indices to convert quaternions between serializations.
-        to_xyzw = jnp.array([1, 2, 3, 0])
-        to_wxyz = jnp.array([3, 0, 1, 2])
-
-        # Get the initial quaternion.
-        W_Q_B_t0 = jaxlie.SO3.from_quaternion_xyzw(
-            xyzw=x0.physics_model.base_quaternion[to_xyzw]
-        )
-
-        # 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 quaternion on SO(3).
-        # Note that we left-multiply with the exponential map since the angular
-        # velocity is expressed in the inertial frame.
-        W_Q_B_tf = jaxlie.SO3.exp(tangent=dt * W_ω_WB_tf) @ W_Q_B_t0
-
-        # Replace the quaternion in the final state.
-        return xf.replace(
-            physics_model=xf.physics_model.replace(
-                base_quaternion=W_Q_B_tf.as_quaternion_xyzw()[to_wxyz]
-            ),
-            validate=True,
-        )
-
-
 @jax_dataclasses.pytree_dataclass
-class ForwardEulerSO3(ExplicitRungeKuttaSO3Mixin, Heun[ODEState]):
+class ForwardEulerSO3(ExplicitRungeKuttaSO3Mixin, ForwardEuler[js.ode_data.ODEState]):
     pass
 
 
 @jax_dataclasses.pytree_dataclass
-class HeunSO3(ExplicitRungeKuttaSO3Mixin, Heun[ODEState]):
+class Heun2SO3(ExplicitRungeKuttaSO3Mixin, Heun2[js.ode_data.ODEState]):
     pass
 
 
 @jax_dataclasses.pytree_dataclass
-class RungeKutta4SO3(ExplicitRungeKuttaSO3Mixin, RungeKutta4[ODEState]):
+class RungeKutta4SO3(ExplicitRungeKuttaSO3Mixin, RungeKutta4[js.ode_data.ODEState]):
     pass