jaxsim 0.2.dev65__py3-none-any.whl → 0.2.dev77__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.dev65'
-__version_tuple__ = version_tuple = (0, 2, 'dev65')
+__version__ = version = '0.2.dev77'
+__version_tuple__ = version_tuple = (0, 2, 'dev77')

jaxsim/simulation/integrators.py

@@ -1,4 +1,5 @@
-from typing import Any, Callable, Dict, Tuple, Union
+import enum
+from typing import Any, Callable
 
 import jax
 import jax.numpy as jnp
@@ -19,30 +20,39 @@ State = jtp.PyTree
 StateDerivative = jtp.PyTree
 
 StateDerivativeCallable = Callable[
-    [State, Time], Tuple[StateDerivative, Dict[str, Any]]
+    [State, Time], tuple[StateDerivative, dict[str, Any]]
 ]
 
 
+class IntegratorType(enum.IntEnum):
+    RungeKutta4 = enum.auto()
+    EulerForward = enum.auto()
+    EulerSemiImplicit = enum.auto()
+    EulerSemiImplicitManifold = enum.auto()
+
+
 # =======================
 # Single-step integration
 # =======================
 
 
-def odeint_euler_one_step(
+def integrator_fixed_single_step(
     dx_dt: StateDerivativeCallable,
-    x0: State,
+    x0: State | ODEState,
     t0: Time,
     tf: Time,
+    integrator_type: IntegratorType,
     num_sub_steps: int = 1,
-) -> Tuple[State, Dict[str, Any]]:
+) -> tuple[State | ODEState, dict[str, Any]]:
     """
-    Forward Euler integrator.
+    Advance a state vector by integrating a sytem dynamics with a fixed-step integrator.
 
     Args:
         dx_dt: Callable that computes the state derivative.
         x0: Initial state.
         t0: Initial time.
         tf: Final time.
+        integrator_type: Integrator type.
         num_sub_steps: Number of sub-steps to break the integration into.
 
     Returns:
@@ -55,10 +65,14 @@ def odeint_euler_one_step(
     sub_step_dt = dt / num_sub_steps
 
     # Initialize the carry
-    Carry = Tuple[State, Time]
+    Carry = tuple[State | ODEState, Time]
     carry_init: Carry = (x0, t0)
 
-    def body_fun(carry: Carry, xs: None) -> Tuple[Carry, None]:
+    def forward_euler_body_fun(carry: Carry, xs: None) -> tuple[Carry, None]:
+        """
+        Forward Euler integrator.
+        """
+
         # Unpack the carry
         x_t0, t0 = carry
 
@@ -78,48 +92,11 @@ def odeint_euler_one_step(
 
         return carry, None
 
-    # Integrate over the given horizon
-    (x_tf, _), _ = jax.lax.scan(
-        f=body_fun, init=carry_init, xs=None, length=num_sub_steps
-    )
+    def rk4_body_fun(carry: Carry, xs: None) -> tuple[Carry, None]:
+        """
+        Runge-Kutta 4 integrator.
+        """
 
-    # Compute the aux dictionary at t0
-    _, aux_t0 = dx_dt(x0, t0)
-
-    return x_tf, aux_t0
-
-
-def odeint_rk4_one_step(
-    dx_dt: StateDerivativeCallable,
-    x0: State,
-    t0: Time,
-    tf: Time,
-    num_sub_steps: int = 1,
-) -> Tuple[State, Dict[str, Any]]:
-    """
-    Runge-Kutta 4 integrator.
-
-    Args:
-        dx_dt: Callable that computes the state derivative.
-        x0: Initial state.
-        t0: Initial time.
-        tf: Final time.
-        num_sub_steps: Number of sub-steps to break the integration into.
-
-    Returns:
-        The final state and a dictionary including auxiliary data at t0.
-    """
-
-    # Compute the sub-step size.
-    # We break dt in configurable sub-steps.
-    dt = tf - t0
-    sub_step_dt = dt / num_sub_steps
-
-    # Initialize the carry
-    Carry = Tuple[State, Time]
-    carry_init: Carry = (x0, t0)
-
-    def body_fun(carry: Carry, xs: None) -> Tuple[Carry, None]:
         # Unpack the carry
         x_t0, t0 = carry
 
@@ -148,49 +125,11 @@ def odeint_rk4_one_step(
 
         return carry, None
 
-    # Integrate over the given horizon
-    (x_tf, _), _ = jax.lax.scan(
-        f=body_fun, init=carry_init, xs=None, length=num_sub_steps
-    )
-
-    # Compute the aux dictionary at t0
-    _, aux_t0 = dx_dt(x0, t0)
-
-    return x_tf, aux_t0
-
-
-def odeint_euler_semi_implicit_one_step(
-    dx_dt: StateDerivativeCallable,
-    x0: ODEState,
-    t0: Time,
-    tf: Time,
-    num_sub_steps: int = 1,
-) -> Tuple[ODEState, Dict[str, Any]]:
-    """
-    Semi-implicit Euler integrator.
-
-    Args:
-        dx_dt: Callable that computes the state derivative.
-        x0: Initial state as ODEState object.
-        t0: Initial time.
-        tf: Final time.
-        num_sub_steps: Number of sub-steps to break the integration into.
-
-    Returns:
-        A tuple having as first element the final state as ODEState object,
-        and as second element a dictionary including auxiliary data at t0.
-    """
+    def semi_implicit_euler_body_fun(carry: Carry, xs: None) -> tuple[Carry, None]:
+        """
+        Semi-implicit Euler integrator.
+        """
 
-    # Compute the sub-step size.
-    # We break dt in configurable sub-steps.
-    dt = tf - t0
-    sub_step_dt = dt / num_sub_steps
-
-    # Initialize the carry
-    Carry = Tuple[ODEState, Time]
-    carry_init: Carry = (x0, t0)
-
-    def body_fun(carry: Carry, xs: None) -> Tuple[Carry, None]:
         # Unpack the carry
         x_t0, t0 = carry
 
@@ -218,6 +157,7 @@ def odeint_euler_semi_implicit_one_step(
         # 2. Compute the derivative of the generalized position
         # 3. Integrate the implicit velocities
         # 4. Integrate the remaining state
+        # 5. Outside the loop: integrate the quaternion on SO(3) manifold
 
         # ----------------------------------------------------------------
         # 1. Integrate the accelerations obtaining the implicit velocities
@@ -254,13 +194,27 @@ def odeint_euler_semi_implicit_one_step(
         BW_vl_WB = (BW_Xv_W @ W_v_WB)[0:3]
 
         # Compute the derivative of the generalized position
-        d_pos_tf = jnp.hstack([BW_vl_WB, W_Qd_B, vel_tf[6:]])
+        d_pos_tf = (
+            jnp.hstack([BW_vl_WB, vel_tf[6:]])
+            if integrator_type is IntegratorType.EulerSemiImplicitManifold
+            else jnp.hstack([BW_vl_WB, W_Qd_B, vel_tf[6:]])
+        )
 
         # ------------------------------------
         # 3. Integrate the implicit velocities
         # ------------------------------------
 
         pos_tf = pos_t0 + sub_step_dt * d_pos_tf
+        joint_positions = (
+            pos_tf[3:]
+            if integrator_type is IntegratorType.EulerSemiImplicitManifold
+            else pos_tf[7:]
+        )
+        base_quaternion = (
+            jnp.zeros_like(x_t0.base_quaternion)
+            if integrator_type is IntegratorType.EulerSemiImplicitManifold
+            else pos_tf[3:7]
+        )
 
         # ---------------------------------
         # 4. Integrate the remaining state
@@ -275,8 +229,8 @@ def odeint_euler_semi_implicit_one_step(
         x_tf = ODEState(
             physics_model=PhysicsModelState(
                 base_position=pos_tf[0:3],
-                base_quaternion=pos_tf[3:7],
-                joint_positions=pos_tf[7:],
+                base_quaternion=base_quaternion,
+                joint_positions=joint_positions,
                 base_linear_velocity=vel_tf[0:3],
                 base_angular_velocity=vel_tf[3:6],
                 joint_velocities=vel_tf[6:],
@@ -294,176 +248,43 @@ def odeint_euler_semi_implicit_one_step(
 
         return carry, None
 
-    # Integrate over the given horizon
-    (x_tf, _), _ = jax.lax.scan(
-        f=body_fun, init=carry_init, xs=None, length=num_sub_steps
-    )
-
-    # Compute the aux dictionary at t0
-    _, aux_t0 = dx_dt(x0, t0)
-
-    return x_tf, aux_t0
-
-
-def odeint_euler_semi_implicit_manifold_one_step(
-    dx_dt: StateDerivativeCallable,
-    x0: ODEState,
-    t0: Time,
-    tf: Time,
-    num_sub_steps: int = 1,
-) -> Tuple[ODEState, Dict[str, Any]]:
-    """
-    Semi-implicit Euler integrator with quaternion integration on SO(3).
-
-    Args:
-        dx_dt: Callable that computes the state derivative.
-        x0: Initial state as ODEState object.
-        t0: Initial time.
-        tf: Final time.
-        num_sub_steps: Number of sub-steps to break the integration into.
-
-    Returns:
-        A tuple having as first element the final state as ODEState object,
-        and as second element a dictionary including auxiliary data at t0.
-    """
-
-    # Compute the sub-step size.
-    # We break dt in configurable sub-steps.
-    dt = tf - t0
-    sub_step_dt = dt / num_sub_steps
-
-    # Integrate the quaternion on its manifold using the new angular velocity
-
-    # Initialize the carry
-    Carry = Tuple[ODEState, Time]
-    carry_init: Carry = (x0, t0)
-
-    def body_fun(carry: Carry, xs: None) -> Tuple[Carry, None]:
-        # Unpack the carry
-        x_t0, t0 = carry
-
-        # Compute the state derivative.
-        # We only keep the quantities related to the acceleration and discard those
-        # related to the velocity since we are going to use those implicitly integrated
-        # from the accelerations.
-        StateDerivative = ODEState
-        dxdt_t0: StateDerivative = dx_dt(x_t0, t0)[0]
-
-        # Extract the initial position ∈ ℝ⁷⁺ⁿ and initial velocity ∈ ℝ⁶⁺ⁿ.
-        # This integrator, contrarily to most of the other ones, is not generic.
-        # It expects to operate on an x object of class ODEState.
-        pos_t0 = x_t0.physics_model.position()
-        vel_t0 = x_t0.physics_model.velocity()
-
-        # Extract the velocity derivative
-        d_vel_dt = dxdt_t0.physics_model.velocity()
-
-        # =============================================
-        # Perform semi-implicit Euler integration [1-4]
-        # =============================================
+    _integrator_registry = {
+        IntegratorType.RungeKutta4: rk4_body_fun,
+        IntegratorType.EulerForward: forward_euler_body_fun,
+        IntegratorType.EulerSemiImplicit: semi_implicit_euler_body_fun,
+        IntegratorType.EulerSemiImplicitManifold: semi_implicit_euler_body_fun,
+    }
 
-        # 1. Integrate the accelerations obtaining the implicit velocities
-        # 2. Compute the derivative of the generalized position (w/o quaternion)
-        # 3. Integrate the implicit velocities (w/o quaternion)
-        # 4. Integrate the remaining state
-        # 5. Outside the loop: integrate the quaternion on SO(3) manifold
-
-        # ----------------------------------------------------------------
-        # 1. Integrate the accelerations obtaining the implicit velocities
-        # ----------------------------------------------------------------
-
-        vel_tf = vel_t0 + sub_step_dt * d_vel_dt
-
-        # ----------------------------------------------------------------------
-        # 2. Compute the derivative of the generalized position (w/o quaternion)
-        # ----------------------------------------------------------------------
-
-        # Compute the transform of the mixed base frame at t0
-        W_H_BW = jnp.vstack(
-            [
-                jnp.block([jnp.eye(3), jnp.vstack(x_t0.physics_model.base_position)]),
-                jnp.array([0, 0, 0, 1]),
-            ]
-        )
-
-        # The derivative W_ṗ_B of the base position is the linear component of the
-        # mixed velocity B[W]_v_WB. We need to compute it from the velocity in
-        # inertial-fixed representation W_vl_WB.
-        W_v_WB = vel_tf[0:6]
-        BW_Xv_W = se3.SE3.from_matrix(W_H_BW).inverse().adjoint()
-        BW_vl_WB = (BW_Xv_W @ W_v_WB)[0:3]
-
-        # Compute the derivative of the generalized position excluding the quaternion
-        pos_no_quat_t0 = jnp.hstack([pos_t0[0:3], pos_t0[7:]])
-        d_pos_no_quat_tf = jnp.hstack([BW_vl_WB, vel_tf[6:]])
-
-        # -----------------------------------------------------
-        # 3. Integrate the implicit velocities (w/o quaternion)
-        # -----------------------------------------------------
-
-        pos_no_quat_tf = pos_no_quat_t0 + sub_step_dt * d_pos_no_quat_tf
-
-        # ---------------------------------
-        # 4. Integrate the remaining state
-        # ---------------------------------
-
-        # Integrate the derivative of the tangential material deformation
-        m = x_t0.soft_contacts.tangential_deformation
-        ṁ = dxdt_t0.soft_contacts.tangential_deformation
-        tangential_deformation_tf = m + sub_step_dt * ṁ
-
-        # Pack the new state into an ODEState object.
-        # We store a zero quaternion as placeholder, it will be replaced later.
-        x_tf = ODEState(
-            physics_model=PhysicsModelState(
-                base_position=pos_no_quat_tf[0:3],
-                base_quaternion=jnp.zeros_like(x_t0.physics_model.base_quaternion),
-                joint_positions=pos_no_quat_tf[3:],
-                base_linear_velocity=vel_tf[0:3],
-                base_angular_velocity=vel_tf[3:6],
-                joint_velocities=vel_tf[6:],
-            ),
-            soft_contacts=SoftContactsState(
-                tangential_deformation=tangential_deformation_tf
-            ),
-        )
-
-        # Update the time
-        tf = t0 + sub_step_dt
-
-        # Pack the carry
-        carry = (x_tf, tf)
-
-        return carry, None
+    # Get the body function for the selected integrator
+    body_fun = _integrator_registry[integrator_type]
 
     # Integrate over the given horizon
-    (x_no_quat_tf, _), _ = jax.lax.scan(
+    (x_tf, _), _ = jax.lax.scan(
         f=body_fun, init=carry_init, xs=None, length=num_sub_steps
     )
 
-    # ---------------------------------------------
-    # 5. Integrate the quaternion on SO(3) manifold
-    # ---------------------------------------------
-
-    # Indices to convert quaternions between serializations
-    to_xyzw = jnp.array([1, 2, 3, 0])
-    to_wxyz = jnp.array([3, 0, 1, 2])
+    if integrator_type is IntegratorType.EulerSemiImplicitManifold:
+        # 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 and the implicitly integrated angular velocity
-    W_ω_WB_tf = x_no_quat_tf.physics_model.base_angular_velocity
-    W_Q_B_t0 = so3.SO3.from_quaternion_xyzw(x0.physics_model.base_quaternion[to_xyzw])
+        # Get the initial quaternion and the implicitly integrated angular velocity
+        W_ω_WB_tf = x_tf.physics_model.base_angular_velocity
+        W_Q_B_t0 = so3.SO3.from_quaternion_xyzw(
+            x0.physics_model.base_quaternion[to_xyzw]
+        )
 
-    # Integrate the quaternion on its manifold using the implicit angular velocity,
-    # transformed in body-fixed representation since jaxlie uses this convention
-    B_R_W = W_Q_B_t0.inverse().as_matrix()
-    W_Q_B_tf = W_Q_B_t0 @ so3.SO3.exp(tangent=dt * B_R_W @ W_ω_WB_tf)
+        # Integrate the quaternion on its manifold using the implicit angular velocity,
+        # transformed in body-fixed representation since jaxlie uses this convention
+        B_R_W = W_Q_B_t0.inverse().as_matrix()
+        W_Q_B_tf = W_Q_B_t0 @ so3.SO3.exp(tangent=dt * B_R_W @ W_ω_WB_tf)
 
-    # Store the quaternion in the final state
-    x_tf = x_no_quat_tf.replace(
-        physics_model=x_no_quat_tf.physics_model.replace(
-            base_quaternion=W_Q_B_tf.as_quaternion_xyzw()[to_wxyz]
+        # Store the quaternion in the final state
+        x_tf = x_tf.replace(
+            physics_model=x_tf.physics_model.replace(
+                base_quaternion=W_Q_B_tf.as_quaternion_xyzw()[to_wxyz]
+            )
         )
-    )
 
     # Compute the aux dictionary at t0
     _, aux_t0 = dx_dt(x0, t0)
@@ -477,10 +298,10 @@ def odeint_euler_semi_implicit_manifold_one_step(
 
 
 def integrate_single_step_over_horizon(
-    integrator_single_step: Callable[[Time, Time, State], Tuple[State, Dict[str, Any]]],
+    integrator_single_step: Callable[[Time, Time, State], tuple[State, dict[str, Any]]],
     t: TimeHorizon,
     x0: State,
-) -> Tuple[State, Dict[str, Any]]:
+) -> tuple[State, dict[str, Any]]:
     """
     Integrate a single-step integrator over a given horizon.
 
@@ -496,7 +317,7 @@ def integrate_single_step_over_horizon(
     # Initialize the carry
     carry_init = (x0, t)
 
-    def body_fun(carry: Tuple, idx: int) -> Tuple[Tuple, jtp.PyTree]:
+    def body_fun(carry: tuple, idx: int) -> tuple[tuple, jtp.PyTree]:
         # Unpack the carry
         x_t0, horizon = carry
 
@@ -526,96 +347,17 @@ def integrate_single_step_over_horizon(
 # ===================================================================
 
 
-def odeint_euler(
-    func,
-    y0: State,
-    t: TimeHorizon,
-    *args,
-    num_sub_steps: int = 1,
-    return_aux: bool = False
-) -> Union[State, Tuple[State, Dict[str, Any]]]:
-    """
-    Integrate a system of ODEs using the Euler method.
-
-    Args:
-        func: A function that computes the time-derivative of the state.
-        y0: The initial state.
-        t: The vector of time instants of the integration horizon.
-        *args: Additional arguments to be passed to the function func.
-        num_sub_steps: The number of sub-steps to be performed within each integration step.
-        return_aux: Whether to return the auxiliary data produced by the integrator.
-
-    Returns:
-        The state of the system at the end of the integration horizon, and optionally
-        the auxiliary data produced by the integrator.
-    """
-
-    # Close func over additional inputs and parameters
-    dx_dt_closure_aux = lambda x, ts: func(x, ts, *args)
-
-    # Close one-step integration over its arguments
-    integrator_single_step = lambda t0, tf, x0: odeint_euler_one_step(
-        dx_dt=dx_dt_closure_aux, x0=x0, t0=t0, tf=tf, num_sub_steps=num_sub_steps
-    )
-
-    # Integrate the state and compute optional auxiliary data over the horizon
-    out, aux = integrate_single_step_over_horizon(
-        integrator_single_step=integrator_single_step, t=t, x0=y0
-    )
-
-    return (out, aux) if return_aux else out
-
-
-def odeint_euler_semi_implicit(
-    func,
-    y0: ODEState,
-    t: TimeHorizon,
-    *args,
-    num_sub_steps: int = 1,
-    return_aux: bool = False
-) -> Union[ODEState, Tuple[ODEState, Dict[str, Any]]]:
-    """
-    Integrate a system of ODEs using the Semi-Implicit Euler method.
-
-    Args:
-        func: A function that computes the time-derivative of the state.
-        y0: The initial state as ODEState object.
-        t: The vector of time instants of the integration horizon.
-        *args: Additional arguments to be passed to the function func.
-        num_sub_steps: The number of sub-steps to be performed within each integration step.
-        return_aux: Whether to return the auxiliary data produced by the integrator.
-
-    Returns:
-        The state of the system at the end of the integration horizon as ODEState object,
-        and optionally the auxiliary data produced by the integrator.
-    """
-
-    # Close func over additional inputs and parameters
-    dx_dt_closure_aux = lambda x, ts: func(x, ts, *args)
-
-    # Close one-step integration over its arguments
-    integrator_single_step = lambda t0, tf, x0: odeint_euler_semi_implicit_one_step(
-        dx_dt=dx_dt_closure_aux, x0=x0, t0=t0, tf=tf, num_sub_steps=num_sub_steps
-    )
-
-    # Integrate the state and compute optional auxiliary data over the horizon
-    out, aux = integrate_single_step_over_horizon(
-        integrator_single_step=integrator_single_step, t=t, x0=y0
-    )
-
-    return (out, aux) if return_aux else out
-
-
-def odeint_rk4(
+def odeint(
     func,
     y0: State,
     t: TimeHorizon,
     *args,
     num_sub_steps: int = 1,
-    return_aux: bool = False
-) -> Union[State, Tuple[State, Dict[str, Any]]]:
+    return_aux: bool = False,
+    integrator_type: IntegratorType = None,
+):
     """
-    Integrate a system of ODEs using the Runge-Kutta 4 method.
+    Integrate a system of ODEs with a fixed-step integrator.
 
     Args:
         func: A function that computes the time-derivative of the state.
@@ -634,8 +376,13 @@ def odeint_rk4(
     dx_dt_closure = lambda x, ts: func(x, ts, *args)
 
     # Close one-step integration over its arguments
-    integrator_single_step = lambda t0, tf, x0: odeint_rk4_one_step(
-        dx_dt=dx_dt_closure, x0=x0, t0=t0, tf=tf, num_sub_steps=num_sub_steps
+    integrator_single_step = lambda t0, tf, x0: integrator_fixed_single_step(
+        dx_dt=dx_dt_closure,
+        x0=x0,
+        t0=t0,
+        tf=tf,
+        num_sub_steps=num_sub_steps,
+        integrator_type=integrator_type,
     )
 
     # Integrate the state and compute optional auxiliary data over the horizon

jaxsim/simulation/ode_integration.py

@@ -10,21 +10,7 @@ from jaxsim.physics.algos.soft_contacts import SoftContactsParams
 from jaxsim.physics.algos.terrain import FlatTerrain, Terrain
 from jaxsim.physics.model.physics_model import PhysicsModel
 from jaxsim.simulation import integrators, ode
-
-
-class IntegratorType(enum.IntEnum):
-    RungeKutta4 = enum.auto()
-    EulerForward = enum.auto()
-    EulerSemiImplicit = enum.auto()
-    EulerSemiImplicitManifold = enum.auto()
-
-
-_integrator_registry = {
-    IntegratorType.RungeKutta4: integrators.odeint_rk4,
-    IntegratorType.EulerForward: integrators.odeint_euler,
-    IntegratorType.EulerSemiImplicit: integrators.odeint_euler_semi_implicit,
-    IntegratorType.EulerSemiImplicitManifold: integrators.odeint_euler_semi_implicit_manifold_one_step,
-}
+from jaxsim.simulation.integrators import IntegratorType
 
 
 @jax.jit
@@ -62,12 +48,13 @@ def ode_integration_fixed_step(
     )
 
     # Integrate over the horizon
-    out = _integrator_registry[integrator_type](
+    out = integrators.odeint(
         func=dx_dt_closure,
         y0=x0,
         t=t,
         num_sub_steps=num_sub_steps,
         return_aux=return_aux,
+        integrator_type=integrator_type,
     )
 
     # Return output pytree and, optionally, the aux dict

{jaxsim-0.2.dev65.dist-info → jaxsim-0.2.dev77.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.dev65
+Version: 0.2.dev77
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo

{jaxsim-0.2.dev65.dist-info → jaxsim-0.2.dev77.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=LJhCG4rsmCrTKTocwRIvllPQeYTxDn-VFn6NjPngn4s,1877
-jaxsim/_version.py,sha256=dLd0JOuKbONNejYzPFlPHkUUqmQSlB4Sz4jslPyREVs,421
+jaxsim/_version.py,sha256=ea7HN3P3BTsbrEcL4jflFM4NRajTns0RZu9M7F_bdGw,421
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=ErTscpEljFyrhPCisZnLEUt6FWLAuEAh-72Teb8Nz98,626
 jaxsim/high_level/__init__.py,sha256=aWYBCsYmEO76Qt4GEi91Hye_ifGFLvc_bpy9OQplz2o,69
@@ -49,10 +49,10 @@ jaxsim/physics/model/ground_contact.py,sha256=mva-yDzYHREmgUu8jGJmIAsf66_SF6ZISm
 jaxsim/physics/model/physics_model.py,sha256=kVTIaJQrxALzyWjWrDLnwDOcxmzaPGSpUOS8BCq-g6M,13249
 jaxsim/physics/model/physics_model_state.py,sha256=LTC-uqUCP1-7-mLHMa6aY4xfBYWuHIextxDH0EEqEmE,5729
 jaxsim/simulation/__init__.py,sha256=WOWkzq7rMGa4xWvjNqTYtD0Nl4yLQtULGW1xU7hD9m0,182
-jaxsim/simulation/integrators.py,sha256=w9hhQULUFyyyw3p5sFn1YPax1TmLJBbzfQw1iEAvq60,21884
+jaxsim/simulation/integrators.py,sha256=WIlL7xi4UocSlWg4Qms8-6puqRYnK5A4r7TJUNPg5g0,13022
 jaxsim/simulation/ode.py,sha256=ntq_iQPIw3SHj64CZWD2mHAKmt05ZgRpw2UwyTxHDOQ,10380
 jaxsim/simulation/ode_data.py,sha256=spzHU5LnOL6mJPuuhho-J61koT-bcTRonqMMkiPo3M4,1750
-jaxsim/simulation/ode_integration.py,sha256=DdA09pQeP53AwYssWON4v9xcMU3LxlVdmW5_76e7Kv4,2373
+jaxsim/simulation/ode_integration.py,sha256=VDprQYoHEE_iI7ia1Mm3RyYl-LRvHU8dJEvRoGA4TFA,1947
 jaxsim/simulation/simulator.py,sha256=qCI5QG0WKkBC5GNqauSvI7rSlGD7CLttTzCgLED7iJM,18123
 jaxsim/simulation/simulator_callbacks.py,sha256=QWdY7dilmjrxeieWCB6RQ-cWpwLuUOK8fYWXpnnBcyU,2217
 jaxsim/simulation/utils.py,sha256=YdNA1mYGBAE7xVA-Dw7_OoBEuh0J8RS2X0RPQZf4c5E,329
@@ -62,8 +62,8 @@ jaxsim/utils/jaxsim_dataclass.py,sha256=FbjfEoCoYC_F-M3wUggXiEhQ7MMS-V_ciYQca-uS
 jaxsim/utils/oop.py,sha256=LQhBXkSOD0zgYNJLO7Bl0FPRg-LvtvPzxyQa1WFP0rM,22616
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/vmappable.py,sha256=NqGL9nGFRI5OorCfnjXsjR_yXigzDxL0lW1YhQ_nMTY,3655
-jaxsim-0.2.dev65.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
-jaxsim-0.2.dev65.dist-info/METADATA,sha256=iojkcf3T3vzR1nKobiTzSlBdsbw_qLvMXX7Yo25A3rI,7486
-jaxsim-0.2.dev65.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-jaxsim-0.2.dev65.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.dev65.dist-info/RECORD,,
+jaxsim-0.2.dev77.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
+jaxsim-0.2.dev77.dist-info/METADATA,sha256=zOGN2gEgqKAs7iRkr9rjGcSsZWUS90wnoekgcZCgaKs,7486
+jaxsim-0.2.dev77.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+jaxsim-0.2.dev77.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.dev77.dist-info/RECORD,,