jaxsim 0.5.1.dev84__tar.gz → 0.5.1.dev91__tar.gz

{jaxsim-0.5.1.dev84 → jaxsim-0.5.1.dev91}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev84
+Version: 0.5.1.dev91
 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.5.1.dev84 → jaxsim-0.5.1.dev91}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev84'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev84')
+__version__ = version = '0.5.1.dev91'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev91')

{jaxsim-0.5.1.dev84 → jaxsim-0.5.1.dev91}/src/jaxsim/integrators/common.py

@@ -53,10 +53,6 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         repr=False, hash=False, compare=False, kw_only=True
     )
 
-    metadata: dict[str, Any] = dataclasses.field(
-        default_factory=dict, repr=False, hash=False, compare=False, kw_only=True
-    )
-
     @classmethod
     def build(
         cls: type[Self],
@@ -102,10 +98,7 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
         metadata = metadata if metadata is not None else {}
 
-        with self.editable(validate=False) as integrator:
-            integrator.metadata = metadata
-
-        with integrator.mutable_context(mutability=Mutability.MUTABLE):
+        with self.mutable_context(mutability=Mutability.MUTABLE) as integrator:
             xf, metadata_step = integrator(x0, t0, dt, **kwargs)
 
         return (
@@ -315,6 +308,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         b = self.b
         A = self.A
 
+        # Extract metadata from the kwargs.
+        metadata = kwargs.pop("metadata", {})
+
         # Close f over optional kwargs.
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
 
@@ -327,7 +323,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         # 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
+            return 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
@@ -381,7 +377,8 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.metadata["dxdt0"] = jax.tree.map(lambda l: l[self.index_of_fsal], K)
+            # Store the first derivative of the next step in the metadata.
+            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`.
@@ -394,7 +391,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             lambda xf: self.post_process_state(x0=x0, t0=t0, xf=xf, dt=dt)
         )(z)
 
-        return z_transformed, aux_dict
+        return z_transformed, aux_dict | {"metadata": metadata}
 
     @staticmethod
     def butcher_tableau_is_valid(

{jaxsim-0.5.1.dev84 → jaxsim-0.5.1.dev91}/src/jaxsim/integrators/variable_step.py

@@ -14,7 +14,6 @@ from jax_dataclasses import Static
 
 import jaxsim.utils.tracing
 from jaxsim import typing as jtp
-from jaxsim.utils import Mutability
 
 from .common import (
     ExplicitRungeKutta,
@@ -271,30 +270,27 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             # 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)}
+            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
+                x0,
+                jnp.array(t0, dtype=float),
+                jnp.array(dt, dtype=float),
+                **(kwargs | {"metadata": metadata}),
             )
 
         # Remove the injected key.
-        _ = integrator.metadata.pop(EmbeddedRungeKutta.InitializingKey)
+        _ = 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
+            lambda l: jnp.nan * jnp.zeros_like(l), 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__(
@@ -307,7 +303,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         #    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.
+        #    `metadata` argument, and it populates with the actual metadata.
         # 3. After the first step, this method operates on the actual metadata.
         #
         # In particular, we store the following information in the metadata:
@@ -318,8 +314,10 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         #   evaluate the dynamics at the final state of the previous step, that matches
         #   the initial state of the current step.
         #
+        metadata = kwargs.pop("metadata", {})
+
         integrator_init = jnp.array(
-            self.metadata.get(self.InitializingKey, False), dtype=bool
+            metadata.get(self.InitializingKey, False), dtype=bool
         )
 
         # Close f over optional kwargs.
@@ -335,24 +333,23 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
         # 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(),
+        metadata["dt0"], metadata["dxdt0"] = jax.lax.cond(
+            pred=("dt0" in metadata) & ~jnp.isnan(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]),
+                metadata.get("dxdt0", f(x0, t0)[0]),
             ),
             false_fun=lambda aux: estimate_step_size(
                 x0=x0, t0=t0, f=f, order=p, atol=self.atol, rtol=self.rtol
             ),
-            operand=self.metadata,
+            operand=metadata,
         )
 
         # Clip the estimated initial step size to the given bounds, if necessary.
-        self.metadata["dt0"] = jnp.clip(
-            self.metadata["dt0"],
-            jnp.minimum(self.dt_min, self.metadata["dt0"]),
-            jnp.minimum(self.dt_max, self.metadata["dt0"]),
+        metadata["dt0"] = jnp.clip(
+            metadata["dt0"],
+            jnp.minimum(self.dt_min, metadata["dt0"]),
+            jnp.minimum(self.dt_max, metadata["dt0"]),
         )
 
         # =========================================================
@@ -364,7 +361,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         carry0: Carry = (
             x0,
             jnp.array(t0).astype(float),
-            self.metadata,
+            metadata,
             jnp.array(0, dtype=int),
             jnp.array(False).astype(bool),
         )
@@ -392,9 +389,10 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             # 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.metadata = metadata
-                z, _ = integrator._compute_next_state(x0=x0, t0=t0, dt=Δt0, **kwargs)
-                metadata_next = integrator.metadata
+                z, aux_dict = integrator._compute_next_state(
+                    x0=x0, t0=t0, dt=Δt0, **kwargs
+                )
+                metadata_next = aux_dict["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)
@@ -481,10 +479,10 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                 metadata_next,
                 discarded_steps,
             ) = jax.lax.cond(
-                pred=discarded_steps
-                >= self.max_step_rejections | local_error
-                <= 1.0 | Δt_next
-                < self.dt_min | integrator_init,
+                pred=(discarded_steps >= self.max_step_rejections)
+                | (local_error <= 1.0)
+                | (Δt_next < self.dt_min)
+                | integrator_init,
                 true_fun=accept_step,
                 false_fun=reject_step,
             )
@@ -510,12 +508,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             init_val=carry0,
         )
 
-        # 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.metadata = metadata_tf
-
-        return xf, {}
+        return xf, {"metadata": metadata_tf}
 
     @property
     def order_of_solution(self) -> int:

{jaxsim-0.5.1.dev84 → jaxsim-0.5.1.dev91}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev84
+Version: 0.5.1.dev91
 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.5.1.dev84 → jaxsim-0.5.1.dev91}/tests/conftest.py

@@ -426,7 +426,7 @@ def jaxsim_model_single_pendulum() -> js.model.JaxSimModel:
 
     rod_model.model.resolve_frames()
 
-    urdf_string = rod.urdf.exporter.UrdfExporter.sdf_to_urdf_string(
+    urdf_string = rod.urdf.exporter.UrdfExporter(pretty=True).to_urdf_string(
         sdf=rod_model.models()[0]
     )
 

{jaxsim-0.5.1.dev84 → jaxsim-0.5.1.dev91}/tests/test_simulations.py

@@ -201,8 +201,19 @@ def run_simulation(
     return data
 
 
+@pytest.mark.parametrize(
+    "integrator",
+    [
+        jaxsim.integrators.fixed_step.ForwardEuler,
+        jaxsim.integrators.fixed_step.ForwardEulerSO3,
+        jaxsim.integrators.fixed_step.RungeKutta4,
+        jaxsim.integrators.fixed_step.RungeKutta4SO3,
+        jaxsim.integrators.variable_step.BogackiShampineSO3,
+    ],
+)
 def test_simulation_with_soft_contacts(
     jaxsim_model_box: js.model.JaxSimModel,
+    integrator: jaxsim.integrators.Integrator,
 ):
 
     model = jaxsim_model_box
@@ -218,6 +229,11 @@ def test_simulation_with_soft_contacts(
         model.kin_dyn_parameters.contact_parameters.enabled = tuple(
             enabled_collidable_points_mask.tolist()
         )
+        model.integrator = integrator.build(
+            dynamics=js.ode.wrap_system_dynamics_for_integration(
+                system_dynamics=js.ode.system_dynamics
+            )
+        )
 
     assert np.sum(model.kin_dyn_parameters.contact_parameters.enabled) == 4