jaxsim 0.4.3.dev312__tar.gz → 0.4.3.dev350__tar.gz

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev312
+Version: 0.4.3.dev350
 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.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev312'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev312')
+__version__ = version = '0.4.3.dev350'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev350')

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/api/contact.py

@@ -138,7 +138,7 @@ def collidable_point_dynamics(
     **kwargs,
 ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
     r"""
-    Compute the 6D force applied to each collidable point.
+    Compute the 6D force applied to each enabled collidable point.
 
     Args:
         model: The model to consider.
@@ -151,7 +151,7 @@ def collidable_point_dynamics(
         kwargs: Additional keyword arguments to pass to the active contact model.
 
     Returns:
-        The 6D force applied to each collidable point and additional data based
+        The 6D force applied to each eneabled collidable point and additional data based
         on the contact model configured:
         - Soft: the material deformation rate.
         - Rigid: no additional data.
@@ -199,15 +199,19 @@ def collidable_point_dynamics(
     )
 
     # Compute the transforms of the implicit frames `C[L] = (W_p_C, [L])`
-    # associated to each collidable point.
+    # associated to the enabled collidable point.
     # In inertial-fixed representation, the computation of these transforms
     # is not necessary and the conversion below becomes a no-op.
+
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
     W_H_C = (
         js.contact.transforms(model=model, data=data)
         if data.velocity_representation is not VelRepr.Inertial
-        else jnp.zeros(
-            shape=(len(model.kin_dyn_parameters.contact_parameters.body), 4, 4)
-        )
+        else jnp.zeros(shape=(len(indices_of_enabled_collidable_points), 4, 4))
     )
 
     # Convert the 6D forces to the active representation.
@@ -246,6 +250,15 @@ def in_contact(
     if link_names is not None and set(link_names).difference(model.link_names()):
         raise ValueError("One or more link names are not part of the model")
 
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
     W_p_Ci = collidable_point_positions(model=model, data=data)
 
     terrain_height = jax.vmap(lambda x, y: model.terrain.height(x=x, y=y))(
@@ -262,7 +275,7 @@ def in_contact(
 
     links_in_contact = jax.vmap(
         lambda link_index: jnp.where(
-            jnp.array(model.kin_dyn_parameters.contact_parameters.body) == link_index,
+            parent_link_idx_of_enabled_collidable_points == link_index,
             below_terrain,
             jnp.zeros_like(below_terrain, dtype=bool),
         ).any()
@@ -426,14 +439,14 @@ def estimate_good_contact_parameters(
 @jax.jit
 def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
     r"""
-    Return the pose of the collidable points.
+    Return the pose of the enabled collidable points.
 
     Args:
         model: The model to consider.
         data: The data of the considered model.
 
     Returns:
-        The stacked SE(3) matrices of all collidable points.
+        The stacked SE(3) matrices of all enabled collidable points.
 
     Note:
         Each collidable point is implicitly associated with a frame
@@ -442,16 +455,27 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
         rigidly attached to.
     """
 
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
     # Get the transforms of the parent link of all collidable points.
     W_H_L = js.model.forward_kinematics(model=model, data=data)[
-        jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
+        parent_link_idx_of_enabled_collidable_points
+    ]
+
+    L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
+        indices_of_enabled_collidable_points
     ]
 
     # Build the link-to-point transform from the displacement between the link frame L
     # and the implicit contact frame C.
-    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(
-        model.kin_dyn_parameters.contact_parameters.point
-    )
+    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(L_p_Ci)
 
     # Compose the work-to-link and link-to-point transforms.
     return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
@@ -465,7 +489,7 @@ def jacobian(
     output_vel_repr: VelRepr | None = None,
 ) -> jtp.Array:
     r"""
-    Return the free-floating Jacobian of the collidable points.
+    Return the free-floating Jacobian of the enabled collidable points.
 
     Args:
         model: The model to consider.
@@ -475,7 +499,7 @@ def jacobian(
 
     Returns:
         The stacked :math:`6 \times (6+n)` free-floating jacobians of the frames associated to the
-        collidable points.
+        enabled collidable points.
 
     Note:
         Each collidable point is implicitly associated with a frame
@@ -488,6 +512,15 @@ def jacobian(
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
     # Compute the Jacobians of all links.
     W_J_WL = js.model.generalized_free_floating_jacobian(
         model=model, data=data, output_vel_repr=VelRepr.Inertial
@@ -496,9 +529,7 @@ def jacobian(
     # Compute the contact Jacobian.
     # In inertial-fixed output representation, the Jacobian of the parent link is also
     # the Jacobian of the frame C implicitly associated with the collidable point.
-    W_J_WC = W_J_WL[
-        jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
-    ]
+    W_J_WC = W_J_WL[parent_link_idx_of_enabled_collidable_points]
 
     # Adjust the output representation.
     match output_vel_repr:
@@ -550,7 +581,7 @@ def jacobian_derivative(
     output_vel_repr: VelRepr | None = None,
 ) -> jtp.Matrix:
     r"""
-    Compute the derivative of the free-floating jacobian of the contact points.
+    Compute the derivative of the free-floating jacobian of the enabled collidable points.
 
     Args:
         model: The model to consider.
@@ -559,7 +590,7 @@ def jacobian_derivative(
             The output velocity representation of the free-floating jacobian derivative.
 
     Returns:
-        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the contact points.
+        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the enabled collidable points.
 
     Note:
         The input representation of the free-floating jacobian derivative is the active
@@ -570,10 +601,18 @@ def jacobian_derivative(
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
     # Get the index of the parent link and the position of the collidable point.
-    parent_link_idxs = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
-    L_p_Ci = jnp.array(model.kin_dyn_parameters.contact_parameters.point)
-    contact_idxs = jnp.arange(L_p_Ci.shape[0])
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
+        indices_of_enabled_collidable_points
+    ]
 
     # Get the transforms of all the parent links.
     W_H_Li = js.model.forward_kinematics(model=model, data=data)
@@ -646,7 +685,7 @@ def jacobian_derivative(
             output_vel_repr=VelRepr.Inertial,
         )
 
-    # Get the Jacobian of the collidable points in the mixed representation.
+    # Get the Jacobian of the enabled collidable points in the mixed representation.
     with data.switch_velocity_representation(VelRepr.Mixed):
         CW_J_WC_BW = jacobian(
             model=model,
@@ -656,13 +695,11 @@ def jacobian_derivative(
 
     def compute_O_J̇_WC_I(
         L_p_C: jtp.Vector,
-        contact_idx: jtp.Int,
+        parent_link_idx: jtp.Int,
         CW_J_WC_BW: jtp.Matrix,
         W_H_L: jtp.Matrix,
     ) -> jtp.Matrix:
 
-        parent_link_idx = parent_link_idxs[contact_idx]
-
         match output_vel_repr:
             case VelRepr.Inertial:
                 O_X_W = W_X_W = Adjoint.from_transform(  # noqa: F841
@@ -703,7 +740,7 @@ def jacobian_derivative(
         return O_J̇_WC_I
 
     O_J̇_WC = jax.vmap(compute_O_J̇_WC_I, in_axes=(0, 0, 0, None))(
-        L_p_Ci, contact_idxs, CW_J_WC_BW, W_H_Li
+        L_p_Ci, parent_link_idx_of_enabled_collidable_points, CW_J_WC_BW, W_H_Li
     )
 
     return O_J̇_WC

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/api/joint.py

@@ -53,9 +53,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str
     """
 
     exceptions.raise_value_error_if(
-        condition=jnp.array(
-            [joint_index < 0, joint_index >= model.number_of_joints()]
-        ).any(),
+        condition=joint_index < 0,
         msg="Invalid joint index '{idx}'",
         idx=joint_index,
     )
@@ -123,10 +121,7 @@ def position_limit(
     """
 
     if model.number_of_joints() == 0:
-        s_min = model.kin_dyn_parameters.joint_parameters.position_limits_min
-        s_max = model.kin_dyn_parameters.joint_parameters.position_limits_max
-
-        return jnp.atleast_1d(s_min).astype(float), jnp.atleast_1d(s_max).astype(float)
+        return jnp.empty(0).astype(float), jnp.empty(0).astype(float)
 
     exceptions.raise_value_error_if(
         condition=jnp.array(
@@ -136,8 +131,12 @@ def position_limit(
         idx=joint_index,
     )
 
-    s_min = model.kin_dyn_parameters.joint_parameters.position_limits_min[joint_index]
-    s_max = model.kin_dyn_parameters.joint_parameters.position_limits_max[joint_index]
+    s_min = jnp.atleast_1d(
+        model.kin_dyn_parameters.joint_parameters.position_limits_min
+    )[joint_index]
+    s_max = jnp.atleast_1d(
+        model.kin_dyn_parameters.joint_parameters.position_limits_max
+    )[joint_index]
 
     return s_min.astype(float), s_max.astype(float)
 

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/api/kin_dyn_parameters.py

@@ -438,7 +438,9 @@ class KynDynParameters(JaxsimDataclass):
     # Helpers to update parameters
     # ============================
 
-    def set_link_mass(self, link_index: int, mass: jtp.FloatLike) -> KynDynParameters:
+    def set_link_mass(
+        self, link_index: jtp.IntLike, mass: jtp.FloatLike
+    ) -> KynDynParameters:
         """
         Set the mass of a link.
 
@@ -457,7 +459,7 @@ class KynDynParameters(JaxsimDataclass):
         return self.replace(link_parameters=link_parameters)
 
     def set_link_inertia(
-        self, link_index: int, inertia: jtp.MatrixLike
+        self, link_index: jtp.IntLike, inertia: jtp.MatrixLike
     ) -> KynDynParameters:
         r"""
         Set the inertia tensor of a link.
@@ -593,10 +595,10 @@ class LinkParameters(JaxsimDataclass):
         """
 
         # Extract the link parameters from the 6D spatial inertia.
-        m, L_p_CoM, I = Inertia.to_params(M=M)
+        m, L_p_CoM, I_CoM = Inertia.to_params(M=M)
 
         # Extract only the necessary elements of the inertia tensor.
-        inertia_elements = I[jnp.triu_indices(3)]
+        inertia_elements = I_CoM[jnp.triu_indices(3)]
 
         return LinkParameters(
             index=jnp.array(index).squeeze().astype(int),
@@ -743,6 +745,9 @@ class ContactParameters(JaxsimDataclass):
         point:
             The translations between the link frame and the collidable point, expressed
             in the coordinates of the parent link frame.
+        enabled:
+            A tuple of booleans representing, for each collidable point, whether it is
+            enabled or not in contact models.
 
     Note:
         Contrarily to LinkParameters and JointParameters, this class is not meant

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/api/link.py

@@ -4,6 +4,7 @@ from collections.abc import Sequence
 import jax
 import jax.numpy as jnp
 import jax.scipy.linalg
+import numpy as np
 
 import jaxsim.api as js
 import jaxsim.rbda
@@ -54,9 +55,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
     """
 
     exceptions.raise_value_error_if(
-        condition=jnp.array(
-            [link_index < 0, link_index >= model.number_of_links()]
-        ).any(),
+        condition=link_index < 0,
         msg="Invalid link index '{idx}'",
         idx=link_index,
     )
@@ -98,7 +97,7 @@ def idxs_to_names(
         The names of the links.
     """
 
-    return tuple(idx_to_name(model=model, link_index=idx) for idx in link_indices)
+    return tuple(np.array(model.kin_dyn_parameters.link_names)[list(link_indices)])
 
 
 # =========

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/api/model.py

@@ -304,7 +304,7 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.model_name
 
-    def number_of_links(self) -> jtp.Int:
+    def number_of_links(self) -> int:
         """
         Return the number of links in the model.
 
@@ -317,7 +317,7 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.kin_dyn_parameters.number_of_links()
 
-    def number_of_joints(self) -> jtp.Int:
+    def number_of_joints(self) -> int:
         """
         Return the number of joints in the model.
 
@@ -419,7 +419,7 @@ class JaxSimModel(JaxsimDataclass):
 def reduce(
     model: JaxSimModel,
     considered_joints: tuple[str, ...],
-    locked_joint_positions: dict[str, jtp.Float] | None = None,
+    locked_joint_positions: dict[str, jtp.FloatLike] | None = None,
 ) -> JaxSimModel:
     """
     Reduce the model by lumping together the links connected by removed joints.
@@ -1038,12 +1038,7 @@ def forward_dynamics_aba(
     C_v̇_WB = to_active(
         W_v̇_WB=W_v̇_WB,
         W_H_C=W_H_C,
-        W_v_WB=jnp.hstack(
-            [
-                data.state.physics_model.base_linear_velocity,
-                data.state.physics_model.base_angular_velocity,
-            ]
-        ),
+        W_v_WB=W_v_WB,
         W_v_WC=W_v_WC,
     )
 
@@ -2274,20 +2269,23 @@ def step(
             # Raise runtime error for not supported case in which Rigid contacts and
             # Baumgarte stabilization are enabled and used with ForwardEuler integrator.
             jaxsim.exceptions.raise_runtime_error_if(
-                condition=jnp.logical_and(
-                    isinstance(
-                        integrator,
-                        jaxsim.integrators.fixed_step.ForwardEuler
-                        | jaxsim.integrators.fixed_step.ForwardEulerSO3,
-                    ),
-                    jnp.array(
-                        [data_tf.contacts_params.K, data_tf.contacts_params.D]
-                    ).any(),
-                ),
+                condition=isinstance(
+                    integrator,
+                    jaxsim.integrators.fixed_step.ForwardEuler
+                    | jaxsim.integrators.fixed_step.ForwardEulerSO3,
+                )
+                & ((data_tf.contacts_params.K > 0) | (data_tf.contacts_params.D > 0)),
                 msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
             )
 
-            W_p_C = js.contact.collidable_point_positions(model, data_tf)
+            # Extract the indices corresponding to the enabled collidable points.
+            indices_of_enabled_collidable_points = (
+                model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+            )
+
+            W_p_C = js.contact.collidable_point_positions(model, data_tf)[
+                indices_of_enabled_collidable_points
+            ]
 
             # Compute the penetration depth of the collidable points.
             δ, *_ = jax.vmap(
@@ -2296,8 +2294,9 @@ def step(
             )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
 
             with data_tf.switch_velocity_representation(VelRepr.Mixed):
-
-                J_WC = js.contact.jacobian(model, data_tf)
+                J_WC = js.contact.jacobian(model, data_tf)[
+                    indices_of_enabled_collidable_points
+                ]
                 M = js.model.free_floating_mass_matrix(model, data_tf)
 
                 # Compute the impact velocity.

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/api/references.py

@@ -503,7 +503,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         ]
 
         exceptions.raise_value_error_if(
-            condition=jnp.logical_not(data.valid(model=model)),
+            condition=~data.valid(model=model),
             msg="The provided data is not valid for the model",
         )
         W_H_Fi = jax.vmap(

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/integrators/common.py

@@ -319,7 +319,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
 
         # Initialize the carry of the for loop with the stacked kᵢ vectors.
-        carry0 = jax.tree_map(
+        carry0 = jax.tree.map(
             lambda l: jnp.zeros((c.size, *l.shape), dtype=l.dtype), x0
         )
 
@@ -507,7 +507,7 @@ class ExplicitRungeKuttaSO3Mixin:
 
         # We assume that the initial quaternion is already unary.
         exceptions.raise_runtime_error_if(
-            condition=jnp.logical_not(jnp.allclose(W_Q_B_t0.dot(W_Q_B_t0), 1.0)),
+            condition=~jnp.allclose(W_Q_B_t0.dot(W_Q_B_t0), 1.0),
             msg="The SO(3) integrator received a quaternion at t0 that is not unary.",
         )
 

{jaxsim-0.4.3.dev312 → jaxsim-0.4.3.dev350}/src/jaxsim/integrators/variable_step.py

@@ -152,7 +152,7 @@ def compute_pytree_scale(
     """
 
     # Consider a zero second pytree, if not given.
-    x2 = jax.tree.map(lambda l: jnp.zeros_like(l), x1) if x2 is None else x2
+    x2 = jax.tree.map(jnp.zeros_like, x1) if x2 is None else x2
 
     # Compute the scaling factors of the initial state and its derivative.
     compute_scale = lambda l1, l2: atol + jnp.maximum(jnp.abs(l1), jnp.abs(l2)) * rtol
@@ -199,9 +199,7 @@ def local_error_estimation(
 
     # Consider a zero estimated final state, if not given.
     xf_estimate = (
-        jax.tree.map(lambda l: jnp.zeros_like(l), xf)
-        if xf_estimate is None
-        else xf_estimate
+        jax.tree.map(jnp.zeros_like, xf) if xf_estimate is None else xf_estimate
     )
 
     # Estimate the error.
@@ -483,14 +481,10 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                 metadata_next,
                 discarded_steps,
             ) = jax.lax.cond(
-                pred=jnp.array(
-                    [
-                        discarded_steps >= self.max_step_rejections,
-                        local_error <= 1.0,
-                        Δt_next < self.dt_min,
-                        integrator_init,
-                    ]
-                ).any(),
+                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,
             )