jaxsim 0.6.2.dev242__tar.gz → 0.6.2.dev259__tar.gz

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev242
+Version: 0.6.2.dev259
 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.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.2.dev242'
-__version_tuple__ = version_tuple = (0, 6, 2, 'dev242')
+__version__ = version = '0.6.2.dev259'
+__version_tuple__ = version_tuple = (0, 6, 2, 'dev259')

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim/api/actuation_model.py

@@ -92,5 +92,34 @@ def compute_resultant_torques(
     # ===============================
 
     τ_total = τ_references + τ_friction + τ_position_limit
-
+    τ_lim = tn_curve_fn(model=model, data=data)
+    τ_total = jnp.clip(τ_total, -τ_lim, τ_lim)
     return τ_total
+
+
+def tn_curve_fn(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    """
+    Compute the torque limits using the tn curve.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The torque limits.
+    """
+
+    τ_max = model.actuation_params.torque_max  # Max torque (Nm)
+    ω_th = model.actuation_params.omega_th  # Threshold speed (rad/s)
+    ω_max = model.actuation_params.omega_max  # Max speed for torque drop-off (rad/s)
+    abs_vel = jnp.abs(data.joint_velocities)
+    τ_lim = jnp.where(
+        abs_vel <= ω_th,
+        τ_max,
+        jnp.where(
+            abs_vel <= ω_max, τ_max * (1 - (abs_vel - ω_th) / (ω_max - ω_th)), 0.0
+        ),
+    )
+    return τ_lim

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim/api/data.py

@@ -431,13 +431,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         if contact_state is None:
             contact_state = self.contact_state
 
-        if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
-            contact_state = {
-                "tangential_deformation": jnp.zeros_like(
-                    contact_state["tangential_deformation"]
-                )
-            }
-
         # Normalize the quaternion to avoid numerical issues.
         base_quaternion_norm = jaxsim.math.safe_norm(
             base_quaternion, axis=-1, keepdims=True
@@ -525,6 +518,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             _joint_transforms=joint_transforms,
             _link_transforms=link_transforms,
             _link_velocities=link_velocities,
+            contact_state=contact_state,
             validate=validate,
         )
 

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim/api/integrators.py

@@ -156,9 +156,117 @@ def rk4_integration(
     return data_tf.replace(model=model)
 
 
+def rk4fast_integration(
+    model: js.model.JaxSimModel,
+    data: JaxSimModelData,
+    link_forces: jtp.Vector,
+    joint_torques: jtp.Vector,
+) -> JaxSimModelData:
+    """
+    Integrate the system state using the Runge-Kutta 4 fast method.
+
+    Note:
+        This method is a faster version of the RK4 method, but it may not be as accurate.
+        It computes the contact forces only once at the beginning of the integration step.
+    """
+
+    dt = model.time_step
+
+    if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
+
+        # Compute the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
+        # with the terrain.
+        W_f_L_terrain, contact_state_derivative = js.contact.link_contact_forces(
+            model=model,
+            data=data,
+            link_forces=link_forces,
+            joint_torques=joint_torques,
+        )
+
+    W_f_L_total = link_forces + W_f_L_terrain
+
+    # Update the contact state data. This is necessary only for the contact models
+    # that require propagation and integration of contact state.
+    contact_state = model.contact_model.update_contact_state(contact_state_derivative)
+
+    def f(x) -> dict[str, jtp.Matrix]:
+
+        with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
+
+            data_ti = data.replace(model=model, **x)
+
+            W_v̇_WB, s̈ = js.model.forward_dynamics_aba(
+                model=model,
+                data=data_ti,
+                joint_forces=joint_torques,
+                link_forces=W_f_L_total,
+            )
+
+            W_ṗ_B, W_Q̇_B, ṡ = js.ode.system_position_dynamics(
+                data=data,
+                baumgarte_quaternion_regularization=1.0,
+            )
+
+        return dict(
+            base_position=W_ṗ_B,
+            base_quaternion=W_Q̇_B,
+            joint_positions=ṡ,
+            base_linear_velocity=W_v̇_WB[0:3],
+            base_angular_velocity=W_v̇_WB[3:6],
+            joint_velocities=s̈,
+            # The contact state is not updated here, as it is assumed to be constant.
+            contact_state=data_ti.contact_state,
+        )
+
+    base_quaternion_norm = jaxsim.math.safe_norm(data._base_quaternion, axis=-1)
+    base_quaternion = data._base_quaternion / jnp.where(
+        base_quaternion_norm == 0, 1.0, base_quaternion_norm
+    )
+
+    x_t0 = dict(
+        base_position=data._base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=data._joint_positions,
+        base_linear_velocity=data._base_linear_velocity,
+        base_angular_velocity=data._base_angular_velocity,
+        joint_velocities=data._joint_velocities,
+        contact_state=contact_state,
+    )
+
+    euler_mid = lambda x, dxdt: x + (0.5 * dt) * dxdt
+    euler_fin = lambda x, dxdt: x + dt * dxdt
+
+    k1 = f(x_t0)
+    k2 = f(jax.tree.map(euler_mid, x_t0, k1))
+    k3 = f(jax.tree.map(euler_mid, x_t0, k2))
+    k4 = f(jax.tree.map(euler_fin, x_t0, k3))
+
+    # 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)
+
+    data_tf = dataclasses.replace(
+        data,
+        _base_position=x_tf["base_position"],
+        _base_quaternion=x_tf["base_quaternion"],
+        _joint_positions=x_tf["joint_positions"],
+        _base_linear_velocity=x_tf["base_linear_velocity"],
+        _base_angular_velocity=x_tf["base_angular_velocity"],
+        _joint_velocities=x_tf["joint_velocities"],
+        contact_state=x_tf["contact_state"],
+    )
+
+    return data_tf.replace(model=model)
+
+
 _INTEGRATORS_MAP: dict[
     js.model.IntegratorType, Callable[..., js.data.JaxSimModelData]
 ] = {
     js.model.IntegratorType.SemiImplicitEuler: semi_implicit_euler_integration,
     js.model.IntegratorType.RungeKutta4: rk4_integration,
+    js.model.IntegratorType.RungeKutta4Fast: rk4fast_integration,
 }

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim/api/model.py

@@ -29,6 +29,7 @@ class IntegratorType(enum.IntEnum):
 
     SemiImplicitEuler = enum.auto()
     RungeKutta4 = enum.auto()
+    RungeKutta4Fast = enum.auto()
 
 
 @jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
@@ -57,6 +58,10 @@ class JaxSimModel(JaxsimDataclass):
         default=None, repr=False
     )
 
+    actuation_params: Static[jaxsim.rbda.actuation.ActuationParams] = dataclasses.field(
+        default=None, repr=False
+    )
+
     kin_dyn_parameters: js.kin_dyn_parameters.KinDynParameters | None = (
         dataclasses.field(default=None, repr=False)
     )
@@ -121,6 +126,7 @@ class JaxSimModel(JaxsimDataclass):
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         contact_params: jaxsim.rbda.contacts.ContactsParams | None = None,
+        actuation_params: jaxsim.rbda.actuation.ActuationParams | None = None,
         integrator: IntegratorType | None = None,
         is_urdf: bool | None = None,
         considered_joints: Sequence[str] | None = None,
@@ -143,6 +149,7 @@ class JaxSimModel(JaxsimDataclass):
                 The contact model to consider.
                 If not specified, a soft contacts model is used.
             contact_params: The parameters of the contact model.
+            actuation_params: The parameters of the actuation model.
             integrator: The integrator to use for the simulation.
             is_urdf:
                 The optional flag to force the model description to be parsed as a URDF.
@@ -177,6 +184,7 @@ class JaxSimModel(JaxsimDataclass):
             time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
+            actuation_params=actuation_params,
             contact_params=contact_params,
             integrator=integrator,
             gravity=-gravity,
@@ -198,6 +206,7 @@ class JaxSimModel(JaxsimDataclass):
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         contact_params: jaxsim.rbda.contacts.ContactsParams | None = None,
+        actuation_params: jaxsim.rbda.actuation.ActuationParams | None = None,
         integrator: IntegratorType | None = None,
         gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
     ) -> JaxSimModel:
@@ -219,6 +228,7 @@ class JaxSimModel(JaxsimDataclass):
                 The contact model to consider.
                 If not specified, a relaxed-constraints rigid contacts model is used.
             contact_params: The parameters of the contact model.
+            actuation_params: The parameters of the actuation model.
             integrator: The integrator to use for the simulation.
             gravity: The gravity constant.
 
@@ -255,6 +265,9 @@ class JaxSimModel(JaxsimDataclass):
         if contact_params is None:
             contact_params = contact_model._parameters_class()
 
+        if actuation_params is None:
+            actuation_params = jaxsim.rbda.actuation.ActuationParams()
+
         # Consider the default integrator if not specified.
         integrator = (
             integrator
@@ -272,6 +285,7 @@ class JaxSimModel(JaxsimDataclass):
             terrain=terrain,
             contact_model=contact_model,
             contact_params=contact_params,
+            actuation_params=actuation_params,
             integrator=integrator,
             gravity=gravity,
             # The following is wrapped as hashless since it's a static argument, and we
@@ -474,6 +488,7 @@ def reduce(
         terrain=model.terrain,
         contact_model=model.contact_model,
         contact_params=model.contact_params,
+        actuation_params=model.actuation_params,
         gravity=model.gravity,
         integrator=model.integrator,
     )

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim/rbda/__init__.py

@@ -1,4 +1,4 @@
-from . import contacts
+from . import actuation, contacts
 from .aba import aba
 from .collidable_points import collidable_points_pos_vel
 from .crba import crba

jaxsim-0.6.2.dev259/src/jaxsim/rbda/actuation/__init__.py

@@ -0,0 +1 @@
+from .common import ActuationParams

jaxsim-0.6.2.dev259/src/jaxsim/rbda/actuation/common.py

@@ -0,0 +1,17 @@
+import dataclasses
+
+import jax_dataclasses
+
+import jaxsim.typing as jtp
+from jaxsim.utils import JaxsimDataclass
+
+
+@jax_dataclasses.pytree_dataclass
+class ActuationParams(JaxsimDataclass):
+    """
+    Parameters class for the actuation model.
+    """
+
+    torque_max: jtp.Float = dataclasses.field(default=3000.0)  # (Nm)
+    omega_th: jtp.Float = dataclasses.field(default=30.0)  # (rad/s)
+    omega_max: jtp.Float = dataclasses.field(default=100.0)  # (rad/s)

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev242
+Version: 0.6.2.dev259
 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.6.2.dev242 → jaxsim-0.6.2.dev259}/src/jaxsim.egg-info/SOURCES.txt

@@ -99,6 +99,8 @@ src/jaxsim/rbda/forward_kinematics.py
 src/jaxsim/rbda/jacobian.py
 src/jaxsim/rbda/rnea.py
 src/jaxsim/rbda/utils.py
+src/jaxsim/rbda/actuation/__init__.py
+src/jaxsim/rbda/actuation/common.py
 src/jaxsim/rbda/contacts/__init__.py
 src/jaxsim/rbda/contacts/common.py
 src/jaxsim/rbda/contacts/relaxed_rigid.py
@@ -112,6 +114,7 @@ src/jaxsim/utils/tracing.py
 src/jaxsim/utils/wrappers.py
 tests/__init__.py
 tests/conftest.py
+tests/test_actuation.py
 tests/test_api_com.py
 tests/test_api_contact.py
 tests/test_api_data.py

{jaxsim-0.6.2.dev242 → jaxsim-0.6.2.dev259}/tests/conftest.py

@@ -133,6 +133,7 @@ def velocity_representation(request) -> jaxsim.VelRepr:
     params=[
         pytest.param(IntegratorType.SemiImplicitEuler, id="semi_implicit_euler"),
         pytest.param(IntegratorType.RungeKutta4, id="runge_kutta_4"),
+        pytest.param(IntegratorType.RungeKutta4Fast, id="runge_kutta_4_fast"),
     ],
 )
 def integrator(request) -> str:

jaxsim-0.6.2.dev259/tests/test_actuation.py

@@ -0,0 +1,37 @@
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.rbda
+from jaxsim import VelRepr
+
+
+def test_tn_curve(jaxsim_model_single_pendulum: js.model.JaxSimModel):
+
+    model = jaxsim_model_single_pendulum
+    new_act_params = jaxsim.rbda.actuation.ActuationParams()
+    with new_act_params.editable(validate=False) as new_act_params:
+        new_act_params.torque_max = 10
+        new_act_params.omega_th = 1
+        new_act_params.omega_max = 2
+    with model.editable(validate=False) as model:
+        model.actuation_params = new_act_params
+    data = js.data.JaxSimModelData.build(
+        model=model,
+        velocity_representation=VelRepr.Inertial,
+    )
+    new_joint_velocities = 1.5 * jnp.ones(model.dofs())
+    joint_torques_0 = 30 * jnp.ones(model.dofs())
+    data_0 = data.replace(model=model, joint_velocities=new_joint_velocities)
+
+    τ_total = js.actuation_model.compute_resultant_torques(
+        model, data_0, joint_force_references=joint_torques_0
+    )
+    assert (τ_total <= joint_torques_0).all()
+    new_joint_velocities = 2.5 * jnp.ones(model.dofs())
+    joint_torques_0 = 30 * jnp.ones(model.dofs())
+    data_0 = data.replace(model=model, joint_velocities=new_joint_velocities)
+
+    τ_total = js.actuation_model.compute_resultant_torques(
+        model, data_0, joint_force_references=joint_torques_0
+    )
+    assert (τ_total == jnp.zeros(model.dofs())).all()