jaxsim 0.4.3.dev245__py3-none-any.whl → 0.4.3.dev271__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.3.dev245'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev245')
+__version__ = version = '0.4.3.dev271'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev271')

jaxsim/api/model.py

@@ -54,6 +54,10 @@ class JaxSimModel(JaxsimDataclass):
         default=None, repr=False
     )
 
+    integrator: Static[jaxsim.integrators.Integrator | None] = dataclasses.field(
+        default=None, repr=False
+    )
+
     _description: Static[wrappers.HashlessObject[ModelDescription | None]] = (
         dataclasses.field(default=None, repr=False)
     )
@@ -93,12 +97,16 @@ class JaxSimModel(JaxsimDataclass):
     # Initialization and state
     # ========================
 
-    @staticmethod
+    @classmethod
     def build_from_model_description(
+        cls,
         model_description: str | pathlib.Path | rod.Model,
-        model_name: str | None = None,
         *,
+        model_name: str | None = None,
         time_step: jtp.FloatLike | None = None,
+        integrator: (
+            jaxsim.integrators.Integrator | type[jaxsim.integrators.Integrator] | None
+        ) = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         is_urdf: bool | None = None,
@@ -120,6 +128,10 @@ class JaxSimModel(JaxsimDataclass):
             contact_model:
                 The contact model to consider.
                 If not specified, a soft contacts model is used.
+            integrator:
+                The integrator to use. If not specified, a default one is used.
+                This argument can either be a pre-built integrator instance or one
+                of the integrator classes defined in JaxSim.
             is_urdf:
                 The optional flag to force the model description to be parsed as a URDF.
                 This is usually automatically inferred.
@@ -146,10 +158,11 @@ class JaxSimModel(JaxsimDataclass):
             )
 
         # Build the model.
-        model = JaxSimModel.build(
+        model = cls.build(
             model_description=intermediate_description,
             model_name=model_name,
             time_step=time_step,
+            integrator=integrator,
             terrain=terrain,
             contact_model=contact_model,
         )
@@ -160,12 +173,16 @@ class JaxSimModel(JaxsimDataclass):
 
         return model
 
-    @staticmethod
+    @classmethod
     def build(
+        cls,
         model_description: ModelDescription,
-        model_name: str | None = None,
         *,
+        model_name: str | None = None,
         time_step: jtp.FloatLike | None = None,
+        integrator: (
+            jaxsim.integrators.Integrator | type[jaxsim.integrators.Integrator] | None
+        ) = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
     ) -> JaxSimModel:
@@ -182,6 +199,11 @@ class JaxSimModel(JaxsimDataclass):
                 The default time step to consider for the simulation. It can be
                 manually overridden in the function that steps the simulation.
             terrain: The terrain to consider (the default is a flat infinite plane).
+                The optional name of the model overriding the physics model name.
+            integrator:
+                The integrator to use. If not specified, a default one is used.
+                This argument can either be a pre-built integrator instance or one
+                of the integrator classes defined in JaxSim.
             contact_model:
                 The contact model to consider.
                 If not specified, a soft contacts model is used.
@@ -195,23 +217,62 @@ class JaxSimModel(JaxsimDataclass):
 
         # Consider the default terrain (a flat infinite plane) if not specified.
         terrain = (
-            terrain or JaxSimModel.__dataclass_fields__["terrain"].default_factory()
+            terrain
+            if terrain is not None
+            else JaxSimModel.__dataclass_fields__["terrain"].default_factory()
         )
 
         # Consider the default time step if not specified.
         time_step = (
-            time_step or JaxSimModel.__dataclass_fields__["time_step"].default_factory()
+            time_step
+            if time_step is not None
+            else JaxSimModel.__dataclass_fields__["time_step"].default_factory()
         )
 
         # Create the default contact model.
         # It will be populated with an initial estimation of good parameters.
         # While these might not be the best, they are a good starting point.
-        contact_model = contact_model or jaxsim.rbda.contacts.SoftContacts.build(
-            terrain=terrain, parameters=None
+        contact_model = (
+            contact_model
+            if contact_model is not None
+            else jaxsim.rbda.contacts.SoftContacts.build(
+                terrain=terrain, parameters=None
+            )
         )
 
+        # Build the integrator if not provided.
+        match integrator:
+
+            # If None, build a default integrator.
+            case None:
+
+                integrator = jaxsim.integrators.fixed_step.Heun2SO3.build(
+                    dynamics=js.ode.wrap_system_dynamics_for_integration(
+                        system_dynamics=js.ode.system_dynamics
+                    )
+                )
+
+            # If it's a pre-built integrator (also a custom one from the user)
+            # just use it as is.
+            case _ if isinstance(integrator, jaxsim.integrators.Integrator):
+                pass
+
+            # If an integrator class is passed, assume that it is a JaxSim integrator
+            # and build it with the default system dynamics.
+            case _ if issubclass(integrator, jaxsim.integrators.Integrator):
+
+                integrator_cls = integrator
+                integrator = integrator_cls.build(
+                    dynamics=js.ode.wrap_system_dynamics_for_integration(
+                        system_dynamics=js.ode.system_dynamics
+                    )
+                )
+
+            case _:
+                raise ValueError(f"Invalid integrator: {integrator}")
+
         # Build the model.
-        model = JaxSimModel(
+        model = cls(
             model_name=model_name,
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
@@ -219,6 +280,7 @@ class JaxSimModel(JaxsimDataclass):
             time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
+            integrator=integrator,
             # The following is wrapped as hashless since it's a static argument, and we
             # don't want to trigger recompilation if it changes. All relevant parameters
             # needed to compute kinematics and dynamics quantities are stored in the
@@ -404,6 +466,7 @@ def reduce(
     reduced_model = JaxSimModel.build(
         model_description=reduced_intermediate_description,
         model_name=model.name(),
+        time_step=model.time_step,
         terrain=model.terrain,
         contact_model=model.contact_model,
     )
@@ -1912,10 +1975,10 @@ def step(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    integrator: jaxsim.integrators.Integrator,
     t0: jtp.FloatLike = 0.0,
     dt: jtp.FloatLike | None = None,
-    integrator_state: dict[str, Any] | None = None,
+    integrator: jaxsim.integrators.Integrator | None = None,
+    integrator_metadata: dict[str, Any] | None = None,
     link_forces: jtp.MatrixLike | None = None,
     joint_force_references: jtp.VectorLike | None = None,
     **kwargs,
@@ -1927,7 +1990,7 @@ def step(
         model: The model to consider.
         data: The data of the considered model.
         integrator: The integrator to use.
-        integrator_state: The state of the integrator.
+        integrator_metadata: The metadata of the integrator, if needed.
         t0: The initial time to consider. Only relevant for time-dependent dynamics.
         dt: The time step to consider. If not specified, it is read from the model.
         link_forces:
@@ -1937,8 +2000,9 @@ def step(
         kwargs: Additional kwargs to pass to the integrator.
 
     Returns:
-        A tuple containing the new data of the model
-        and the new state of the integrator.
+        A tuple containing the new data of the model and a dictionary of auxiliary
+        data computed during the step. If the integrator has metadata, the dictionary
+        will contain the new metadata stored in the `integrator_metadata` key.
 
     Note:
         In order to reduce the occurrences of frame conversions performed internally,
@@ -1953,8 +2017,9 @@ def step(
     integrator_kwargs = kwargs.pop("integrator_kwargs", {})
     integrator_kwargs = kwargs | integrator_kwargs
 
-    # Initialize the integrator state.
-    integrator_state_t0 = integrator_state if integrator_state is not None else dict()
+    # Extract the integrator and the optional metadata.
+    integrator_metadata_t0 = integrator_metadata
+    integrator = integrator if integrator is not None else model.integrator
 
     # Initialize the time-related variables.
     state_t0 = data.state
@@ -2010,11 +2075,11 @@ def step(
         τ_references = references.joint_force_references(model=model)
 
     # Step the dynamics forward.
-    state_tf, integrator_state_tf = integrator.step(
+    state_tf, integrator_metadata_tf = integrator.step(
         x0=state_t0,
         t0=t0,
         dt=dt,
-        params=integrator_state_t0,
+        metadata=integrator_metadata_t0,
         # Always inject the current (model, data) pair into the system dynamics
         # considered by the integrator, and include the input variables represented
         # by the pair (f_L, τ_references).
@@ -2091,13 +2156,17 @@ def step(
                     )
                 )
 
-            # Reset the generalized velocity.
-            data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
-            data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
+                # Reset the generalized velocity.
+                data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
+                data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
 
     # Restore the input velocity representation.
     data_tf = data_tf.replace(
         velocity_representation=data.velocity_representation, validate=False
     )
 
-    return data_tf, integrator_state_tf
+    return data_tf, {} | (
+        dict(integrator_metadata=integrator_metadata_tf)
+        if integrator_metadata is not None
+        else {}
+    )

jaxsim/api/ode.py

@@ -24,41 +24,45 @@ class SystemDynamicsFromModelAndData(Protocol):
 
 
 def wrap_system_dynamics_for_integration(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
     *,
     system_dynamics: SystemDynamicsFromModelAndData,
-    **kwargs,
+    **kwargs: dict[str, Any],
 ) -> jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]:
     """
-    Wrap generic system dynamics operating on `JaxSimModel` and `JaxSimModelData`
-    for integration with `jaxsim.integrators`.
+    Wrap the system dynamics considered by JaxSim integrators in a generic
+    `f(x, t, **u, **parameters)` function.
 
     Args:
-        model: The model to consider.
-        data: The data of the considered model.
         system_dynamics: The system dynamics to wrap.
         **kwargs: Additional kwargs to close over the system dynamics.
 
     Returns:
-        The system dynamics closed over the model, the data, and the additional kwargs.
+        The system dynamics closed over the additional kwargs to be used by
+        JaxSim integrators.
     """
 
-    # We allow to close `system_dynamics` over additional kwargs.
-    kwargs_closed = kwargs.copy()
-
-    # Create a local copy of model and data.
-    # The wrapped dynamics will hold a reference of this object.
-    model_closed = model.copy()
-    data_closed = data.copy().replace(
-        state=js.ode_data.ODEState.zero(model=model_closed, data=data)
-    )
-
+    # Close `system_dynamics` over additional kwargs.
+    # Similarly to what done in `jaxsim.api.model.step`, to be future-proof, we use the
+    # following logic to allow the caller to close over arguments having the same name
+    # of the ones used in the `wrap_system_dynamics_for_integration` function.
+    kwargs = kwargs.copy() if kwargs is not None else {}
+    colliding_system_dynamics_kwargs = kwargs.pop("system_dynamics_kwargs", {})
+    system_dynamics_kwargs = kwargs | colliding_system_dynamics_kwargs
+
+    # Remove `model` and `data` for backward compatibility.
+    # It's no longer necessary to close over them at this stage, as this is always
+    # done in `jaxsim.api.model.step`.
+    # We can remove the following lines in a few releases.
+    _ = system_dynamics_kwargs.pop("data", None)
+    _ = system_dynamics_kwargs.pop("model", None)
+
+    # Create the function with the signature expected by our generic integrators.
+    # Note that our system dynamics is time independent.
     def f(x: ODEState, t: Time, **kwargs_f) -> tuple[ODEState, dict[str, Any]]:
 
-        # Allow caller to override the closed data and model objects.
-        data_f = kwargs_f.pop("data", data_closed)
-        model_f = kwargs_f.pop("model", model_closed)
+        # Get the data and model objects from the kwargs.
+        data_f = kwargs_f.pop("data")
+        model_f = kwargs_f.pop("model")
 
         # Update the state and time stored inside data.
         with data_f.editable(validate=True) as data_rw:
@@ -69,7 +73,7 @@ def wrap_system_dynamics_for_integration(
         return system_dynamics(
             model=model_f,
             data=data_rw,
-            **(kwargs_closed | kwargs_f),
+            **(system_dynamics_kwargs | kwargs_f),
         )
 
     f: jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]

jaxsim/integrators/common.py

@@ -10,7 +10,7 @@ from jax_dataclasses import Static
 import jaxsim.api as js
 import jaxsim.math
 import jaxsim.typing as jtp
-from jaxsim import exceptions
+from jaxsim import exceptions, logging
 from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
 
 try:
@@ -49,16 +49,11 @@ 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(
         repr=False, hash=False, compare=False, kw_only=True
     )
 
-    params: dict[str, Any] = dataclasses.field(
+    metadata: dict[str, Any] = dataclasses.field(
         default_factory=dict, repr=False, hash=False, compare=False, kw_only=True
     )
 
@@ -88,9 +83,9 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         t0: Time,
         dt: TimeStep,
         *,
-        params: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
         **kwargs,
-    ) -> tuple[State, dict[str, Any]]:
+    ) -> tuple[NextState, dict[str, Any]]:
         """
         Perform a single integration step.
 
@@ -98,28 +93,30 @@ 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.
-            params: The auxiliary dictionary of the integrator.
+            metadata: The state auxiliary dictionary of the integrator.
             **kwargs: Additional keyword arguments.
 
         Returns:
             The final state of the system and the updated auxiliary dictionary.
         """
 
+        metadata = metadata if metadata is not None else {}
+
         with self.editable(validate=False) as integrator:
-            integrator.params = params
+            integrator.metadata = metadata
 
         with integrator.mutable_context(mutability=Mutability.MUTABLE):
-            xf, aux_dict = integrator(x0, t0, dt, **kwargs)
+            xf, metadata_step = integrator(x0, t0, dt, **kwargs)
 
         return (
             xf,
-            integrator.params
-            | {Integrator.AfterInitKey: jnp.array(False).astype(bool)}
-            | aux_dict,
+            metadata | metadata_step,
         )
 
     @abc.abstractmethod
-    def __call__(self, x0: State, t0: Time, dt: TimeStep, **kwargs) -> NextState:
+    def __call__(
+        self, x0: State, t0: Time, dt: TimeStep, **kwargs
+    ) -> tuple[NextState, dict[str, Any]]:
         pass
 
     def init(
@@ -131,62 +128,12 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         include_dynamics_aux_dict: bool = False,
         **kwargs,
     ) -> dict[str, Any]:
-        """
-        Initialize the integrator.
-
-        Args:
-            x0: The initial state of the system.
-            t0: The initial time of the system.
-            dt: The time step of the integration.
-
-        Returns:
-            The auxiliary dictionary of the integrator.
 
-        Note:
-            This method should have the same signature as the inherited `__call__`
-            method, including additional kwargs.
-
-        Note:
-            If the integrator supports FSAL, the pair `(x0, t0)` must match the real
-            initial state and time of the system, otherwise the initial derivative of
-            the first step will be wrong.
-        """
-
-        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)
-
-        # Remove the injected key.
-        _ = integrator.params.pop(Integrator.InitializingKey)
-
-        # 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.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
+        logging.warning(
+            "The 'init' method has been deprecated. There is no need to call it."
+        )
 
-        return params_after_init
+        return {}
 
 
 @jax_dataclasses.pytree_dataclass
@@ -377,8 +324,11 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             x0,
         )
 
-        # Apply FSAL property by passing ẋ0 = f(x0, t0) from the previous iteration.
-        get_ẋ0_and_aux_dict = lambda: self.params.get("dxdt0", f(x0, t0))
+        # Closure on metadata to either evaluate the dynamics at the initial state
+        # or to use the previous state derivative (only integrators supporting FSAL).
+        def get_ẋ0_and_aux_dict() -> tuple[StateDerivative, dict[str, Any]]:
+            ẋ0, aux_dict = f(x0, t0)
+            return self.metadata.get("dxdt0", ẋ0), aux_dict
 
         # We use a `jax.lax.scan` to compile the `f` function only once.
         # Otherwise, if we compute e.g. for RK4 sequentially, the jit-compiled code
@@ -405,8 +355,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
                 # Compute the next time for the kᵢ evaluation.
                 ti = t0 + c[i] * Δt
 
-                # This is kᵢ, aux_dict = f(xᵢ, tᵢ).
-                return f(xi, ti)
+                # Evaluate the dynamics.
+                ki, aux_dict = f(xi, ti)
+                return ki, aux_dict
 
             # This selector enables FSAL property in the first iteration (i=0).
             ki, aux_dict = jax.lax.cond(
@@ -431,7 +382,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.params["dxdt0"] = jax.tree.map(lambda l: l[self.index_of_fsal], K)
+            self.metadata["dxdt0"] = jax.tree.map(lambda l: l[self.index_of_fsal], K)
 
         # Compute the output state.
         # Note that z contains as many new states as the rows of `b.T`.
@@ -514,7 +465,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             raise ValueError("The Butcher tableau is not valid.")
 
         if not ExplicitRungeKutta.butcher_tableau_is_explicit(A=A):
-            return False
+            return False, None
 
         if index_of_solution >= b.T.shape[0]:
             msg = "The index of the solution (i-th row of `b.T`) is out of range."

jaxsim/integrators/variable_step.py

@@ -12,6 +12,7 @@ import jax.numpy as jnp
 import jax_dataclasses
 from jax_dataclasses import Static
 
+import jaxsim.utils.tracing
 from jaxsim import typing as jtp
 from jaxsim.utils import Mutability
 
@@ -219,6 +220,9 @@ def local_error_estimation(
 @jax_dataclasses.pytree_dataclass
 class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
+    AfterInitKey: ClassVar[str] = "after_init"
+    InitializingKey: ClassVar[str] = "initializing"
+
     # Define the row of the integration output corresponding to the solution estimate.
     # This is the row of b.T that produces the state used e.g. by embedded methods to
     # implement the adaptive timestep logic.
@@ -246,40 +250,79 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         self,
         x0: State,
         t0: Time,
-        dt: TimeStep | None = None,
-        *,
-        include_dynamics_aux_dict: bool = False,
+        dt: TimeStep,
         **kwargs,
     ) -> dict[str, Any]:
+        """
+        Initialize the integrator and get the metadata.
 
-        # In these type of integrators, it's not relevant picking a meaningful dt.
-        # We just need to execute __call__ once to initialize the dictionary of params.
-        return super().init(
-            x0=x0,
-            t0=t0,
-            dt=0.001,
-            include_dynamics_aux_dict=include_dynamics_aux_dict,
-            **kwargs,
+        Args:
+            x0: The initial state of the system.
+            t0: The initial time of the system.
+            dt: The time step of the integration.
+
+        Returns:
+            The metadata of the integrator to be passed to the first step.
+        """
+
+        if jaxsim.utils.tracing(var=jnp.zeros(0)):
+            raise RuntimeError("This method cannot be used within a JIT context")
+
+        with self.editable(validate=False) as integrator:
+
+            # Inject this key to signal that the integrator is initializing.
+            # This is used to allocate the arrays of the metadata dictionary,
+            # that are then filled with NaNs.
+            integrator.metadata = {EmbeddedRungeKutta.InitializingKey: jnp.array(True)}
+
+            # Run a dummy call of the integrator.
+            # It is used only to get the metadata so that we know the structure
+            # of the corresponding pytree.
+            _ = integrator(
+                x0, jnp.array(t0, dtype=float), jnp.array(dt, dtype=float), **kwargs
+            )
+
+        # Remove the injected key.
+        _ = integrator.metadata.pop(EmbeddedRungeKutta.InitializingKey)
+
+        # Make sure that all leafs of the dictionary are JAX arrays.
+        # Also, since these are dummy parameters, set them all to NaN.
+        metadata_after_init = jax.tree.map(
+            lambda l: jnp.nan * jnp.zeros_like(l), integrator.metadata
         )
 
+        # 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.metadata = metadata_after_init
+
+        return metadata_after_init
+
     def __call__(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
     ) -> tuple[NextState, dict[str, Any]]:
 
         # This method is called differently in three stages:
         #
-        # 1. During initialization, to allocate a dummy params dictionary.
-        # 2. During the first step, to compute the initial valid params dictionary.
-        # 3. After the first step, to compute the next state and the next valid params.
+        # 1. During initialization, to allocate a dummy metadata dictionary.
+        #    The metadata is a dictionary of float JAX arrays, that are initialized
+        #    with the right shape and filled with NaNs.
+        # 2. During the first step, this method operates on the Nan-filled
+        #    `self.metadata` attribute, and it populates with the actual metadata.
+        # 3. After the first step, this method operates on the actual metadata.
         #
-        # Stage 1 produces a zero-filled dummy dictionary.
-        # Stage 2 receives a dummy dictionary and produces valid parameters that can be
-        # fed to later steps.
-        # Stage 3 corresponds to any consecutive step after the first one. It can re-use
-        # data (like for FSAL) from previous steps.
+        # In particular, we store the following information in the metadata:
+        # - The first attempt of the step size, `dt0`. This is either estimated during
+        #   phase 2, or taken from the previous step during phase 3.
+        # - For integrators that support FSAL, the derivative at the initial state
+        #   computed during the previous step. This can be done because FSAL integrators
+        #   evaluate the dynamics at the final state of the previous step, that matches
+        #   the initial state of the current step.
         #
-        integrator_init = self.params.get(self.InitializingKey, jnp.array(False))
-        integrator_first_step = self.params.get(self.AfterInitKey, jnp.array(False))
+        integrator_init = jnp.array(
+            self.metadata.get(self.InitializingKey, False), dtype=bool
+        )
 
         # Close f over optional kwargs.
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
@@ -292,34 +335,26 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         p̂ = self.order_of_solution_estimate
         q = jnp.minimum(p, p̂)
 
-        # In Stage 1 and 2, estimate from scratch dt0 and dxdt0.
-        # In Stage 3, dt0 is taken from the previous step. If the integrator supports
-        # FSAL, dxdt0 is taken from the previous step. Otherwise, it is computed by
-        # evaluating the dynamics.
-        self.params["dt0"], self.params["dxdt0"], aux_dict = jax.lax.cond(
-            pred=jnp.logical_or("dt0" not in self.params, integrator_first_step),
-            true_fun=lambda params: (
-                *estimate_step_size(
-                    x0=x0, t0=t0, f=f, order=p, atol=self.atol, rtol=self.rtol
-                ),
-                self.params.get("dxdt0", f(x0, t0))[1],
+        # The value of dt0 is NaN (or, at least, it should be) only after initialization
+        # and before the first step.
+        self.metadata["dt0"], self.metadata["dxdt0"] = jax.lax.cond(
+            pred=("dt0" in self.metadata)
+            & ~jnp.isnan(self.metadata.get("dt0", 0.0)).any(),
+            true_fun=lambda metadata: (
+                metadata.get("dt0", jnp.array(0.0, dtype=float)),
+                self.metadata.get("dxdt0", f(x0, t0)[0]),
             ),
-            false_fun=lambda params: (
-                params.get("dt0", jnp.array(0).astype(float)),
-                *self.params.get("dxdt0", f(x0, t0)),
+            false_fun=lambda aux: estimate_step_size(
+                x0=x0, t0=t0, f=f, order=p, atol=self.atol, rtol=self.rtol
             ),
-            operand=self.params,
+            operand=self.metadata,
         )
 
-        # If the integrator does not support FSAL, it is useless to store dxdt0.
-        if not self.has_fsal:
-            _ = self.params.pop("dxdt0")
-
         # Clip the estimated initial step size to the given bounds, if necessary.
-        self.params["dt0"] = jnp.clip(
-            self.params["dt0"],
-            jnp.minimum(self.dt_min, self.params["dt0"]),
-            jnp.minimum(self.dt_max, self.params["dt0"]),
+        self.metadata["dt0"] = jnp.clip(
+            self.metadata["dt0"],
+            jnp.minimum(self.dt_min, self.metadata["dt0"]),
+            jnp.minimum(self.dt_max, self.metadata["dt0"]),
         )
 
         # =========================================================
@@ -331,7 +366,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         carry0: Carry = (
             x0,
             jnp.array(t0).astype(float),
-            self.params,
+            self.metadata,
             jnp.array(0, dtype=int),
             jnp.array(False).astype(bool),
         )
@@ -347,21 +382,21 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         def while_loop_body(carry: Carry) -> Carry:
 
             # Unpack the carry.
-            x0, t0, params, discarded_steps, _ = carry
+            x0, t0, metadata, discarded_steps, _ = carry
 
             # Take care of the final adaptive step.
             # We want the final Δt to let us reach tf exactly.
             # Then we can exit the while loop.
-            Δt0 = params["dt0"]
+            Δt0 = metadata["dt0"]
             Δt0 = jnp.where(t0 + Δt0 < tf, Δt0, tf - t0)
             break_loop = jnp.where(t0 + Δt0 < tf, False, True)
 
             # Run the underlying explicit RK integrator.
             # The output z contains multiple solutions (depending on the rows of b.T).
             with self.editable(validate=True) as integrator:
-                integrator.params = params
+                integrator.metadata = metadata
                 z, _ = integrator._compute_next_state(x0=x0, t0=t0, dt=Δt0, **kwargs)
-                params_next = integrator.params
+                metadata_next = integrator.metadata
 
             # Extract the high-order solution xf and the low-order estimate x̂f.
             xf = jax.tree.map(lambda l: l[self.row_index_of_solution], z)
@@ -394,11 +429,11 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             def accept_step():
                 # Use Δt_next in the next while loop.
                 # If it is the last one, and Δt0 was clipped, return the initial Δt0.
-                params_next_accepted = params_next | dict(
+                metadata_next_accepted = metadata_next | dict(
                     dt0=jnp.clip(
                         jax.lax.select(
                             pred=break_loop,
-                            on_true=params["dt0"],
+                            on_true=metadata["dt0"],
                             on_false=Δt_next,
                         ),
                         self.dt_min,
@@ -419,16 +454,16 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                     x0_next,
                     t0_next,
                     break_loop_next,
-                    params_next_accepted,
+                    metadata_next_accepted,
                     jnp.array(0, dtype=int),
                 )
 
             def reject_step():
-                # Get back the original params.
-                params_next_rejected = params
+                # Get back the original metadata.
+                metadata_next_rejected = metadata
 
                 # This time, with a reduced Δt.
-                params_next_rejected["dt0"] = jnp.clip(
+                metadata_next_rejected["dt0"] = jnp.clip(
                     Δt_next, self.dt_min, self.dt_max
                 )
 
@@ -436,7 +471,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                     x0,
                     t0,
                     False,
-                    params_next_rejected,
+                    metadata_next_rejected,
                     discarded_steps + 1,
                 )
 
@@ -445,7 +480,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                 x0_next,
                 t0_next,
                 break_loop,
-                params_next,
+                metadata_next,
                 discarded_steps,
             ) = jax.lax.cond(
                 pred=jnp.array(
@@ -463,7 +498,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             return (
                 x0_next,
                 t0_next,
-                params_next,
+                metadata_next,
                 discarded_steps,
                 break_loop,
             )
@@ -472,7 +507,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         (
             xf,
             tf,
-            params_tf,
+            metadata_tf,
             _,
             _,
         ) = jax.lax.while_loop(
@@ -484,9 +519,9 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         # Store the parameters.
         # They will be returned to the caller in a functional way in the step method.
         with self.mutable_context(mutability=Mutability.MUTABLE):
-            self.params = params_tf
+            self.metadata = metadata_tf
 
-        return xf, aux_dict
+        return xf, {}
 
     @property
     def order_of_solution(self) -> int:

{jaxsim-0.4.3.dev245.dist-info → jaxsim-0.4.3.dev271.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev245
+Version: 0.4.3.dev271
 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.3.dev245.dist-info → jaxsim-0.4.3.dev271.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
-jaxsim/_version.py,sha256=M-BxL76M-utEC1M8k4OEPThRT-EPO5mWYFUyfr6637A,428
+jaxsim/_version.py,sha256=WDllAgWs3R6C27z-j1tgLt8CxLvz_Ys39X00J18KcPI,428
 jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
@@ -12,14 +12,14 @@ jaxsim/api/frame.py,sha256=yPSgNygHkvWlln4wShNt7vZm_fFobVEm7phsklNNyH8,12922
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
 jaxsim/api/kin_dyn_parameters.py,sha256=kbDN5n9uj8CamVJXk1U5oYLbxyjaWDIeUG0V68DCEFs,29578
 jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
-jaxsim/api/model.py,sha256=FSK2KvH2oNoGIt67gxQD7ScbClqXCn1Oy5JLtc4QHYg,70585
-jaxsim/api/ode.py,sha256=2KvGT3WW1eWEme4fH-5LlOXdLF6JUP5Z-IGY93ashUc,13815
+jaxsim/api/model.py,sha256=dpQZDT0BodMfK1wmpG-STFh-rFsJStobQ1fhrWILK9o,73410
+jaxsim/api/ode.py,sha256=jFE4yk5lHSNk_SynbgA4tHcPdWq17cB-qUUW8KhcknQ,14289
 jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
 jaxsim/api/references.py,sha256=fW77LitZ8DYgT6ZmUInJfm5luBV1mTcqcNRiC_i79og,20862
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=78MBs89GxsL0wU2yAexjvBZt3HEtfZoGVIN9f0a8yTc,20305
+jaxsim/integrators/common.py,sha256=_FZs7E0EazERGA3K0tGC1baUrs8sBDzYTf2U2mFYh9s,18329
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
-jaxsim/integrators/variable_step.py,sha256=1VoSU3GeFcGEuP2dgZQ83sTkI5Xe-IThqKlRoVtwGSE,21270
+jaxsim/integrators/variable_step.py,sha256=hGYKG3Sq3QITgzIePmCVCrrirwagqsKnB3aYifAcKR4,22848
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
 jaxsim/math/adjoint.py,sha256=o1FCipkGwPtMbN2gFNIyUV8ADF3TX5fxElpTEXK0bIs,4377
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
@@ -65,8 +65,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev245.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev245.dist-info/METADATA,sha256=KkzN5EWZkz0Oj6awfhPKRHWrXTZNNTl5vVlfNK64ZY8,17276
-jaxsim-0.4.3.dev245.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-jaxsim-0.4.3.dev245.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev245.dist-info/RECORD,,
+jaxsim-0.4.3.dev271.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev271.dist-info/METADATA,sha256=3B9vB5QnQXwr70wS3SjMkXOMgqPauA_bSIOinzvb7xU,17276
+jaxsim-0.4.3.dev271.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+jaxsim-0.4.3.dev271.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev271.dist-info/RECORD,,

{jaxsim-0.4.3.dev245.dist-info → jaxsim-0.4.3.dev271.dist-info}/WHEEL

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