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

jaxsim/simulation/integrators.py

@@ -1,393 +0,0 @@
-import enum
-from typing import Any, Callable
-
-import jax
-import jax.numpy as jnp
-from jax.tree_util import tree_map
-
-import jaxsim.typing as jtp
-from jaxsim.math.quaternion import Quaternion
-from jaxsim.physics.algos.soft_contacts import SoftContactsState
-from jaxsim.physics.model.physics_model_state import PhysicsModelState
-from jaxsim.simulation.ode_data import ODEState
-from jaxsim.sixd import se3, so3
-
-Time = jtp.FloatLike
-TimeStep = jtp.FloatLike
-TimeHorizon = jtp.VectorLike
-
-State = jtp.PyTree
-StateDerivative = jtp.PyTree
-
-StateDerivativeCallable = Callable[
-    [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 integrator_fixed_single_step(
-    dx_dt: StateDerivativeCallable,
-    x0: State | ODEState,
-    t0: Time,
-    tf: Time,
-    integrator_type: IntegratorType,
-    num_sub_steps: int = 1,
-) -> tuple[State | ODEState, dict[str, Any]]:
-    """
-    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:
-        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 | ODEState, Time]
-    carry_init: Carry = (x0, t0)
-
-    def forward_euler_body_fun(carry: Carry, xs: None) -> tuple[Carry, None]:
-        """
-        Forward Euler integrator.
-        """
-
-        # Unpack the carry
-        x_t0, t0 = carry
-
-        # Compute the state derivative
-        dxdt_t0, _ = dx_dt(x_t0, t0)
-
-        # Integrate the dynamics
-        x_tf = jax.tree_util.tree_map(
-            lambda x, dxdt: x + sub_step_dt * dxdt, x_t0, dxdt_t0
-        )
-
-        # Update the time
-        tf = t0 + sub_step_dt
-
-        # Pack the carry
-        carry = (x_tf, tf)
-
-        return carry, None
-
-    def rk4_body_fun(carry: Carry, xs: None) -> tuple[Carry, None]:
-        """
-        Runge-Kutta 4 integrator.
-        """
-
-        # Unpack the carry
-        x_t0, t0 = carry
-
-        # Helper to forward the state to compute k2 and k3 at midpoint and k4 at final
-        euler_mid = lambda x, dxdt: x + (0.5 * sub_step_dt) * dxdt
-        euler_fin = lambda x, dxdt: x + sub_step_dt * dxdt
-
-        # Compute the RK4 slopes
-        k1, _ = dx_dt(x_t0, t0)
-        k2, _ = dx_dt(tree_map(euler_mid, x_t0, k1), t0 + 0.5 * sub_step_dt)
-        k3, _ = dx_dt(tree_map(euler_mid, x_t0, k2), t0 + 0.5 * sub_step_dt)
-        k4, _ = dx_dt(tree_map(euler_fin, x_t0, k3), t0 + sub_step_dt)
-
-        # Average the slopes and compute the RK4 state derivative
-        average = lambda k1, k2, k3, k4: (k1 + 2 * k2 + 2 * k3 + k4) / 6
-        dxdt = jax.tree_util.tree_map(average, k1, k2, k3, k4)
-
-        # Integrate the dynamics
-        x_tf = jax.tree_util.tree_map(euler_fin, x_t0, dxdt)
-
-        # Update the time
-        tf = t0 + sub_step_dt
-
-        # Pack the carry
-        carry = (x_tf, tf)
-
-        return carry, None
-
-    def semi_implicit_euler_body_fun(carry: Carry, xs: None) -> tuple[Carry, None]:
-        """
-        Semi-implicit Euler integrator.
-        """
-
-        # 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]
-        # =============================================
-
-        # 1. Integrate the accelerations obtaining the implicit velocities
-        # 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
-        # ----------------------------------------------------------------
-
-        vel_tf = vel_t0 + sub_step_dt * d_vel_dt
-
-        # -----------------------------------------------------
-        # 2. Compute the derivative of the generalized position
-        # -----------------------------------------------------
-
-        # Extract the implicit angular velocity and the initial base quaternion
-        W_ω_WB = vel_tf[3:6]
-        W_Q_B = x_t0.physics_model.base_quaternion
-
-        # Compute the quaternion derivative and the base position derivative
-        W_Qd_B = Quaternion.derivative(
-            quaternion=W_Q_B, omega=W_ω_WB, omega_in_body_fixed=False
-        ).squeeze()
-
-        # 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
-        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
-        # ---------------------------------
-
-        # 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
-        x_tf = ODEState(
-            physics_model=PhysicsModelState(
-                base_position=pos_tf[0:3],
-                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:],
-            ),
-            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
-
-    _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,
-    }
-
-    # Get the body function for the selected integrator
-    body_fun = _integrator_registry[integrator_type]
-
-    # Integrate over the given horizon
-    (x_tf, _), _ = jax.lax.scan(
-        f=body_fun, init=carry_init, xs=None, length=num_sub_steps
-    )
-
-    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_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)
-
-        # 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)
-
-    return x_tf, aux_t0
-
-
-# ===============================
-# Adapter: single step -> horizon
-# ===============================
-
-
-def integrate_single_step_over_horizon(
-    integrator_single_step: Callable[[Time, Time, State], tuple[State, dict[str, Any]]],
-    t: TimeHorizon,
-    x0: State,
-) -> tuple[State, dict[str, Any]]:
-    """
-    Integrate a single-step integrator over a given horizon.
-
-    Args:
-        integrator_single_step: A single-step integrator.
-        t: The vector of time instants of the integration horizon.
-        x0: The initial state of the integration horizon.
-
-    Returns:
-        The final state and auxiliary data produced by the integrator.
-    """
-
-    # Initialize the carry
-    carry_init = (x0, t)
-
-    def body_fun(carry: tuple, idx: int) -> tuple[tuple, jtp.PyTree]:
-        # Unpack the carry
-        x_t0, horizon = carry
-
-        # Get the integration interval
-        t0 = horizon[idx]
-        tf = horizon[idx + 1]
-
-        # Perform a single-step integration of the ODE
-        x_tf, aux_t0 = integrator_single_step(t0, tf, x_t0)
-
-        # Prepare returned data
-        out = (x_t0, aux_t0)
-        carry = (x_tf, horizon)
-
-        return carry, out
-
-    # Integrate over the given horizon
-    _, (x_horizon, aux_horizon) = jax.lax.scan(
-        f=body_fun, init=carry_init, xs=jnp.arange(start=0, stop=len(t), dtype=int)
-    )
-
-    return x_horizon, aux_horizon
-
-
-# ===================================================================
-# Integration over horizon (same APIs of jax.experimental.ode.odeint)
-# ===================================================================
-
-
-def odeint(
-    func,
-    y0: State,
-    t: TimeHorizon,
-    *args,
-    num_sub_steps: int = 1,
-    return_aux: bool = False,
-    integrator_type: IntegratorType = None,
-):
-    """
-    Integrate a system of ODEs with a fixed-step integrator.
-
-    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 = lambda x, ts: func(x, ts, *args)
-
-    # Close one-step integration over its arguments
-    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
-    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

jaxsim/simulation/ode.py

@@ -1,290 +0,0 @@
-from typing import Any, Dict, Tuple
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import jaxsim.typing as jtp
-from jaxsim.physics import algos
-from jaxsim.physics.algos.soft_contacts import (
-    SoftContacts,
-    SoftContactsParams,
-    collidable_points_pos_vel,
-)
-from jaxsim.physics.algos.terrain import FlatTerrain, Terrain
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-from . import ode_data
-
-
-def compute_contact_forces(
-    physics_model: PhysicsModel,
-    ode_state: ode_data.ODEState,
-    soft_contacts_params: SoftContactsParams = SoftContactsParams(),
-    terrain: Terrain = FlatTerrain(),
-) -> Tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix]:
-    """
-    Compute the contact forces acting on the collidable points of the model.
-
-    Args:
-        physics_model: The physics model to consider.
-        ode_state: The state of the ODE corresponding to the physics model.
-        soft_contacts_params: The parameters of the soft contacts model.
-        terrain: The terrain model.
-
-    Returns:
-        A tuple containing:
-        - The contact forces expressed in the world frame acting on the model's links.
-        - The derivative of the tangential deformation of the terrain dynamics.
-        - The contact forces expressed in the world frame acting on the model's collidable points.
-    """
-
-    # Compute position and linear mixed velocity of all model's collidable points
-    # collidable_points_kinematics
-    pos_cp, vel_cp = collidable_points_pos_vel(
-        model=physics_model,
-        q=ode_state.physics_model.joint_positions,
-        qd=ode_state.physics_model.joint_velocities,
-        xfb=ode_state.physics_model.xfb(),
-    )
-
-    # Compute the forces acting on the collidable points due to contact with
-    # the compliant ground surface. Apply vmap to process all points together.
-    contact_forces_points, tangential_deformation_dot = jax.vmap(
-        SoftContacts(parameters=soft_contacts_params, terrain=terrain).contact_model
-    )(pos_cp.T, vel_cp.T, ode_state.soft_contacts.tangential_deformation.T)
-
-    contact_forces_points = contact_forces_points.T
-    tangential_deformation_dot = tangential_deformation_dot.T
-
-    # Initialize the contact forces, one per body
-    contact_forces_links = jnp.zeros_like(
-        ode_data.ODEInput.zero(physics_model).physics_model.f_ext
-    )
-
-    # Combine the contact forces of all collidable points belonging to the same body
-    for body_idx in set(physics_model.gc.body):
-        body_idx = int(body_idx)
-        contact_forces_links = contact_forces_links.at[body_idx, :].set(
-            jnp.sum(contact_forces_points[:, physics_model.gc.body == body_idx], axis=1)
-        )
-
-    return contact_forces_links, tangential_deformation_dot, contact_forces_points.T
-
-
-def dx_dt(
-    x: ode_data.ODEState,
-    t: jtp.Float | None,
-    physics_model: PhysicsModel,
-    soft_contacts_params: SoftContactsParams = SoftContactsParams(),
-    ode_input: ode_data.ODEInput | None = None,
-    terrain: Terrain = FlatTerrain(),
-) -> Tuple[ode_data.ODEState, Dict[str, Any]]:
-    """
-    Compute the state derivative of the ODE corresponding to the physics model.
-
-    Args:
-        x: The state of the ODE.
-        t: The current time.
-        physics_model: The physics model to consider.
-        soft_contacts_params: The parameters of the soft contacts model.
-        ode_input: The input of the ODE.
-        terrain: The terrain model.
-
-    Returns:
-        A tuple containing:
-        - The state derivative of the ODE.
-        - A dictionary containing auxiliary information.
-    """
-
-    if t is not None and isinstance(t, np.ndarray) and t.size != 1:
-        raise ValueError(t.size)
-
-    # Initialize arguments
-    ode_state = x
-    ode_input = (
-        ode_input
-        if ode_input is not None
-        else ode_data.ODEInput.zero(physics_model=physics_model)
-    )
-
-    # ======================
-    # Compute contact forces
-    # ======================
-
-    # Initialize the collidable points contact forces
-    contact_forces_points = None
-
-    # Initialize the contact forces, one per body
-    contact_forces_links = jnp.zeros_like(ode_input.physics_model.f_ext)
-
-    # Initialize the derivative of the tangential deformation
-    tangential_deformation_dot = jnp.zeros_like(
-        ode_state.soft_contacts.tangential_deformation
-    )
-
-    if physics_model.gc.body.size > 0:
-        (
-            contact_forces_links,
-            tangential_deformation_dot,
-            contact_forces_points,
-        ) = compute_contact_forces(
-            physics_model=physics_model,
-            soft_contacts_params=soft_contacts_params,
-            ode_state=ode_state,
-            terrain=terrain,
-        )
-
-    # =====================
-    # Joint position limits
-    # =====================
-
-    if physics_model.dofs() > 0:
-        # Get the joint position limits
-        s_min, s_max = jnp.array(
-            [j.position_limit for j in physics_model.description.joints_dict.values()]
-        ).T
-
-        # Get the spring/damper parameters of joint limits enforcement
-        k_damper = jnp.array(list(physics_model._joint_limit_damper.values()))
-
-        # Compute the joint torques that enforce joint limits
-        s = ode_state.physics_model.joint_positions
-        tau_min = jnp.where(s <= s_min, k_damper * (s_min - s), 0)
-        tau_max = jnp.where(s >= s_max, k_damper * (s_max - s), 0)
-        tau_limit = tau_max + tau_min
-
-    else:
-        tau_limit = jnp.zeros_like(ode_input.physics_model.tau)
-
-    # ==============
-    # Joint friction
-    # ==============
-
-    if physics_model.dofs() > 0:
-        # Static and viscous joint friction parameters
-        kc = jnp.array(list(physics_model._joint_friction_static.values()))
-        kv = jnp.array(list(physics_model._joint_friction_viscous.values()))
-
-        # Compute the joint friction torque
-        tau_friction = -(
-            jnp.diag(kc) @ jnp.sign(ode_state.physics_model.joint_positions)
-            + jnp.diag(kv) @ ode_state.physics_model.joint_velocities
-        )
-
-    else:
-        tau_friction = jnp.zeros_like(ode_input.physics_model.tau)
-
-    # ========================
-    # Compute forward dynamics
-    # ========================
-
-    # Compute the total forces applied to the bodies
-    total_forces = ode_input.physics_model.f_ext + contact_forces_links
-
-    # Compute the joint torques to actuate
-    tau = ode_input.physics_model.tau + tau_friction + tau_limit
-
-    # Compute forward dynamics with the ABA algorithm
-    W_a_WB, qdd = algos.aba.aba(
-        model=physics_model,
-        xfb=ode_state.physics_model.xfb(),
-        q=ode_state.physics_model.joint_positions,
-        qd=ode_state.physics_model.joint_velocities,
-        tau=tau,
-        f_ext=total_forces,
-    )
-
-    # =========================================
-    # Compute the state derivative of base link
-    # =========================================
-
-    if not physics_model.is_floating_base:
-        W_Qd_B = jnp.zeros(4)
-        BW_v_WB = jnp.zeros(3)
-
-    else:
-        from jaxsim.math.conv import Convert
-        from jaxsim.math.quaternion import Quaternion
-
-        W_Qd_B = Quaternion.derivative(
-            quaternion=ode_state.physics_model.base_quaternion,
-            omega=ode_state.physics_model.base_angular_velocity,
-            omega_in_body_fixed=False,
-        ).squeeze()
-
-        # Compute linear component of mixed velocity
-        BW_v_WB = Convert.velocities_threed(
-            v_6d=jnp.hstack(
-                [
-                    ode_state.physics_model.base_linear_velocity,
-                    ode_state.physics_model.base_angular_velocity,
-                ]
-            ),
-            p=ode_state.physics_model.base_position,
-        ).squeeze()
-
-    # Derivative of xfb (floating-base state)
-    xd_fb = jnp.hstack([W_Qd_B, BW_v_WB, W_a_WB.squeeze()]).squeeze()
-
-    # =====================================
-    # Build the full derivative of ODEState
-    # =====================================
-
-    def fix_one_dof(vector: jtp.Vector) -> jtp.Vector | None:
-        """Fix the shape of computed quantities for models with just 1 DoF."""
-
-        if vector is None:
-            return None
-
-        return jnp.array([vector]) if vector.shape == () else vector
-
-    # Fill the PhysicsModelState object included in the input ODEState to store the
-    # returned PhysicsModelState derivative
-    physics_model_state_derivative = ode_state.physics_model.replace(
-        joint_positions=fix_one_dof(ode_state.physics_model.joint_velocities.squeeze()),
-        joint_velocities=fix_one_dof(qdd.squeeze()),
-        base_quaternion=xd_fb.squeeze()[0:4],
-        base_position=xd_fb.squeeze()[4:7],
-        base_angular_velocity=xd_fb.squeeze()[10:13],
-        base_linear_velocity=xd_fb.squeeze()[7:10],
-    )
-
-    # Fill the SoftContactsState object included in the input ODEState to store the
-    # returned SoftContactsState derivative
-    soft_contacts_state_derivative = ode_state.soft_contacts.replace(
-        tangential_deformation=tangential_deformation_dot.squeeze(),
-    )
-
-    # We store the state derivative using the ODEState class so that the pytree
-    # structure remains consistent, allowing to use our generic pytree integrators
-    state_derivative = ode_data.ODEState(
-        physics_model=physics_model_state_derivative,
-        soft_contacts=soft_contacts_state_derivative,
-    )
-
-    # ===============================
-    # Build auxiliary data and return
-    # ===============================
-
-    # Real ODEInput containing the real joint forces that have been actuated and
-    # the total external forces (= original external forces + terrain + limits)
-    ode_input_real = ode_data.ODEInput(
-        physics_model=ode_data.PhysicsModelInput(tau=tau, f_ext=total_forces)
-    )
-
-    # Pack the inertial-fixed floating-base acceleration
-    W_nud_WB = jnp.hstack([W_a_WB.squeeze(), qdd.squeeze()])
-
-    # Build the auxiliary data
-    aux_dict = {
-        "model_acceleration": W_nud_WB,
-        "ode_input": ode_input,
-        "ode_input_real": ode_input_real,
-        "contact_forces_links": contact_forces_links,
-        "contact_forces_points": contact_forces_points,
-        "tangential_deformation_dot": tangential_deformation_dot,
-    }
-
-    # Return the state derivative as a generic PyTree, and the dict with auxiliary info
-    return state_derivative, aux_dict