jaxsim 0.4.3.dev282__py3-none-any.whl → 0.4.3.dev298__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/contact.py

@@ -243,9 +243,7 @@ def in_contact(
         A boolean vector indicating whether the links are in contact with the terrain.
     """
 
-    link_names = link_names if link_names is not None else model.link_names()
-
-    if set(link_names).difference(model.link_names()):
+    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")
 
     W_p_Ci = collidable_point_positions(model=model, data=data)
@@ -256,13 +254,19 @@ def in_contact(
 
     below_terrain = W_p_Ci[:, 2] <= terrain_height
 
+    link_idxs = (
+        js.link.names_to_idxs(link_names=link_names, model=model)
+        if link_names is not None
+        else jnp.arange(model.number_of_links())
+    )
+
     links_in_contact = jax.vmap(
         lambda link_index: jnp.where(
             jnp.array(model.kin_dyn_parameters.contact_parameters.body) == link_index,
             below_terrain,
             jnp.zeros_like(below_terrain, dtype=bool),
         ).any()
-    )(js.link.names_to_idxs(link_names=link_names, model=model))
+    )(link_idxs)
 
     return links_in_contact
 

jaxsim/api/data.py

@@ -292,11 +292,13 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             msg = "The data object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
-
-        return self.state.physics_model.joint_positions[
+        joint_idxs = (
             js.joint.names_to_idxs(joint_names=joint_names, model=model)
-        ]
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
+
+        return self.state.physics_model.joint_positions[joint_idxs]
 
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def joint_velocities(
@@ -337,11 +339,13 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             msg = "The data object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
-
-        return self.state.physics_model.joint_velocities[
+        joint_idxs = (
             js.joint.names_to_idxs(joint_names=joint_names, model=model)
-        ]
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
+
+        return self.state.physics_model.joint_velocities[joint_idxs]
 
     @jax.jit
     def base_position(self) -> jtp.Vector:
@@ -374,10 +378,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         # we introduce a Baumgarte stabilization to let the quaternion converge to
         # a unit quaternion. In this case, it is not guaranteed that the quaternion
         # stored in the state is a unit quaternion.
-        W_Q_B = jax.lax.select(
-            pred=jnp.allclose(jnp.linalg.norm(W_Q_B), 1.0, atol=1e-6, rtol=0.0),
-            on_true=W_Q_B,
-            on_false=W_Q_B / jnp.linalg.norm(W_Q_B),
+        W_Q_B = jnp.where(
+            jnp.allclose(W_Q_B.dot(W_Q_B), 1.0), W_Q_B, W_Q_B / jnp.linalg.norm(W_Q_B)
         )
 
         return (W_Q_B if not dcm else jaxsim.math.Quaternion.to_dcm(W_Q_B)).astype(
@@ -502,12 +504,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             msg = "The data object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
+        joint_idxs = (
+            js.joint.names_to_idxs(joint_names=joint_names, model=model)
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
 
         return replace(
-            s=self.state.physics_model.joint_positions.at[
-                js.joint.names_to_idxs(joint_names=joint_names, model=model)
-            ].set(positions)
+            s=self.state.physics_model.joint_positions.at[joint_idxs].set(positions)
         )
 
     @functools.partial(jax.jit, static_argnames=["joint_names"])
@@ -548,12 +552,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             msg = "The data object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
+        joint_idxs = (
+            js.joint.names_to_idxs(joint_names=joint_names, model=model)
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
 
         return replace(
-            ṡ=self.state.physics_model.joint_velocities.at[
-                js.joint.names_to_idxs(joint_names=joint_names, model=model)
-            ].set(velocities)
+            ṡ=self.state.physics_model.joint_velocities.at[joint_idxs].set(velocities)
         )
 
     @jax.jit

jaxsim/api/joint.py

@@ -157,13 +157,19 @@ def position_limits(
         The position limits of the joints.
     """
 
-    joint_names = joint_names if joint_names is not None else model.joint_names()
+    joint_idxs = (
+        names_to_idxs(joint_names=joint_names, model=model)
+        if joint_names is not None
+        else jnp.arange(model.number_of_joints())
+    )
 
-    if len(joint_names) == 0:
+    if len(joint_idxs) == 0:
         return jnp.empty(0).astype(float), jnp.empty(0).astype(float)
 
-    joint_idxs = names_to_idxs(joint_names=joint_names, model=model)
-    return jax.vmap(lambda i: position_limit(model=model, joint_index=i))(joint_idxs)
+    s_min = model.kin_dyn_parameters.joint_parameters.position_limits_min[joint_idxs]
+    s_max = model.kin_dyn_parameters.joint_parameters.position_limits_max[joint_idxs]
+
+    return s_min.astype(float), s_max.astype(float)
 
 
 # ======================
@@ -203,7 +209,11 @@ def random_joint_positions(
     # Get the joint indices.
     # Note that it will trigger an exception if the given `joint_names` are not valid.
     joint_names = joint_names if joint_names is not None else model.joint_names()
-    joint_indices = names_to_idxs(model=model, joint_names=joint_names)
+    joint_indices = (
+        names_to_idxs(model=model, joint_names=joint_names)
+        if joint_names is not None
+        else jnp.arange(model.number_of_joints())
+    )
 
     from jaxsim.parsers.descriptions.joint import JointType
 

jaxsim/api/kin_dyn_parameters.py

@@ -398,9 +398,7 @@ class KynDynParameters(JaxsimDataclass):
         λ_H_pre = jnp.vstack(
             [
                 jnp.eye(4)[jnp.newaxis],
-                jax.vmap(
-                    lambda i: self.joint_model.parent_H_predecessor(joint_index=i)
-                )(jnp.arange(1, 1 + self.number_of_joints())),
+                self.joint_model.λ_H_pre[1 : 1 + self.number_of_joints()],
             ]
         )
 

jaxsim/api/link.py

@@ -423,156 +423,9 @@ def jacobian_derivative(
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
-    # Compute the derivative of the doubly-left free-floating full jacobian.
-    B_J̇_full_WX_B, B_H_L = jaxsim.rbda.jacobian_derivative_full_doubly_left(
-        model=model,
-        joint_positions=data.joint_positions(),
-        joint_velocities=data.joint_velocities(),
-    )
-
-    # Compute the actual doubly-left free-floating jacobian derivative of the link
-    # by zeroing the columns not in the path π_B(L) using the boolean κ(i).
-    κb = model.kin_dyn_parameters.support_body_array_bool[link_index]
-    B_J̇_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J̇_full_WX_B
-
-    # =====================================================
-    # Compute quantities to adjust the input representation
-    # =====================================================
-
-    In = jnp.eye(model.dofs())
-    On = jnp.zeros(shape=(model.dofs(), model.dofs()))
-
-    match data.velocity_representation:
-
-        case VelRepr.Inertial:
-
-            W_H_B = data.base_transform()
-            B_X_W = jaxsim.math.Adjoint.from_transform(transform=W_H_B, inverse=True)
-
-            with data.switch_velocity_representation(VelRepr.Inertial):
-                W_v_WB = data.base_velocity()
-                B_Ẋ_W = -B_X_W @ jaxsim.math.Cross.vx(W_v_WB)
-
-            # Compute the operator to change the representation of ν, and its
-            # time derivative.
-            T = jax.scipy.linalg.block_diag(B_X_W, In)
-            Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_W, On)
-
-        case VelRepr.Body:
-
-            B_X_B = jaxsim.math.Adjoint.from_rotation_and_translation(
-                translation=jnp.zeros(3), rotation=jnp.eye(3)
-            )
-
-            B_Ẋ_B = jnp.zeros(shape=(6, 6))
-
-            # Compute the operator to change the representation of ν, and its
-            # time derivative.
-            T = jax.scipy.linalg.block_diag(B_X_B, In)
-            Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_B, On)
-
-        case VelRepr.Mixed:
-
-            BW_H_B = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
-            B_X_BW = jaxsim.math.Adjoint.from_transform(transform=BW_H_B, inverse=True)
-
-            with data.switch_velocity_representation(VelRepr.Mixed):
-                BW_v_WB = data.base_velocity()
-                BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
-
-            BW_v_BW_B = BW_v_WB - BW_v_W_BW
-            B_Ẋ_BW = -B_X_BW @ jaxsim.math.Cross.vx(BW_v_BW_B)
-
-            # Compute the operator to change the representation of ν, and its
-            # time derivative.
-            T = jax.scipy.linalg.block_diag(B_X_BW, In)
-            Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_BW, On)
-
-        case _:
-            raise ValueError(data.velocity_representation)
-
-    # ======================================================
-    # Compute quantities to adjust the output representation
-    # ======================================================
-
-    match output_vel_repr:
-
-        case VelRepr.Inertial:
-
-            W_H_B = data.base_transform()
-            O_X_B = W_X_B = jaxsim.math.Adjoint.from_transform(transform=W_H_B)
-
-            with data.switch_velocity_representation(VelRepr.Body):
-                B_v_WB = data.base_velocity()
-
-            O_Ẋ_B = W_Ẋ_B = W_X_B @ jaxsim.math.Cross.vx(B_v_WB)  # noqa: F841
-
-        case VelRepr.Body:
-
-            O_X_B = L_X_B = jaxsim.math.Adjoint.from_transform(
-                transform=B_H_L[link_index, :, :], inverse=True
-            )
-
-            B_X_L = jaxsim.math.Adjoint.inverse(adjoint=L_X_B)
-
-            with data.switch_velocity_representation(VelRepr.Body):
-                B_v_WB = data.base_velocity()
-                L_v_WL = js.link.velocity(model=model, data=data, link_index=link_index)
-
-            O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
-                B_X_L @ L_v_WL - B_v_WB
-            )
-
-        case VelRepr.Mixed:
-
-            W_H_B = data.base_transform()
-            W_H_L = W_H_B @ B_H_L[link_index, :, :]
-            LW_H_L = W_H_L.at[0:3, 3].set(jnp.zeros(3))
-            LW_H_B = LW_H_L @ jaxsim.math.Transform.inverse(B_H_L[link_index, :, :])
-
-            O_X_B = LW_X_B = jaxsim.math.Adjoint.from_transform(transform=LW_H_B)
-
-            B_X_LW = jaxsim.math.Adjoint.inverse(adjoint=LW_X_B)
-
-            with data.switch_velocity_representation(VelRepr.Body):
-                B_v_WB = data.base_velocity()
-
-            with data.switch_velocity_representation(VelRepr.Mixed):
-                LW_v_WL = js.link.velocity(
-                    model=model, data=data, link_index=link_index
-                )
-                LW_v_W_LW = LW_v_WL.at[3:6].set(jnp.zeros(3))
-
-            LW_v_LW_L = LW_v_WL - LW_v_W_LW
-            LW_v_B_LW = LW_v_WL - LW_X_B @ B_v_WB - LW_v_LW_L
-
-            O_Ẋ_B = LW_Ẋ_B = -LW_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
-                B_X_LW @ LW_v_B_LW
-            )
-        case _:
-            raise ValueError(output_vel_repr)
-
-    # =============================================================
-    # Express the Jacobian derivative in the target representations
-    # =============================================================
-
-    # The derivative of the equation to change the input and output representations
-    # of the Jacobian derivative needs the computation of the plain link Jacobian.
-    # Compute here the full Jacobian of the model...
-    B_J_full_WL_B, _ = jaxsim.rbda.jacobian_full_doubly_left(
-        model=model,
-        joint_positions=data.joint_positions(),
-    )
-
-    # ... and extract the link Jacobian using the boolean support body array.
-    B_J_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J_full_WL_B
-
-    # Sum all the components that form the Jacobian derivative in the target
-    # input/output velocity representations.
-    O_J̇_WL_I = jnp.zeros(shape=(6, 6 + model.dofs()))
-    O_J̇_WL_I += O_Ẋ_B @ B_J_WL_B @ T
-    O_J̇_WL_I += O_X_B @ B_J̇_WL_B @ T
-    O_J̇_WL_I += O_X_B @ B_J_WL_B @ Ṫ
+    O_J̇_WL_I = js.model.generalized_free_floating_jacobian_derivative(
+        model=model, data=data, output_vel_repr=output_vel_repr
+    )[link_index]
 
     return O_J̇_WL_I
 

jaxsim/api/model.py

@@ -700,14 +700,173 @@ def generalized_free_floating_jacobian_derivative(
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
-    O_J̇_WL_I = jax.vmap(
-        lambda model, data, link_idxs, output_vel_repr: js.link.jacobian_derivative(
-            model, data, link_index=link_idxs, output_vel_repr=output_vel_repr
-        ),
-        in_axes=(None, None, 0, None),
-    )(model, data, jnp.arange(model.number_of_links()), output_vel_repr)
+    # Compute the derivative of the doubly-left free-floating full jacobian.
+    B_J̇_full_WX_B, B_H_L = jaxsim.rbda.jacobian_derivative_full_doubly_left(
+        model=model,
+        joint_positions=data.joint_positions(),
+        joint_velocities=data.joint_velocities(),
+    )
+
+    # The derivative of the equation to change the input and output representations
+    # of the Jacobian derivative needs the computation of the plain link Jacobian.
+    B_J_full_WL_B, _ = jaxsim.rbda.jacobian_full_doubly_left(
+        model=model,
+        joint_positions=data.joint_positions(),
+    )
+
+    # Compute the actual doubly-left free-floating jacobian derivative of the link
+    # by zeroing the columns not in the path π_B(L) using the boolean κ(i).
+    κb = model.kin_dyn_parameters.support_body_array_bool
+
+    # Compute the base transform.
+    W_H_B = data.base_transform()
+
+    @functools.partial(jax.vmap, in_axes=(0, None, None, 0))
+    def _compute_row(
+        B_H_L: jtp.Matrix,
+        B_J_full_WL_B: jtp.Matrix,
+        W_H_B: jtp.Matrix,
+        κb: jtp.Matrix,
+    ) -> jtp.Matrix:
+
+        # =====================================================
+        # Compute quantities to adjust the input representation
+        # =====================================================
+
+        In = jnp.eye(model.dofs())
+        On = jnp.zeros(shape=(model.dofs(), model.dofs()))
+
+        # Extract the link quantities using the boolean support body array.
+        B_J̇_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J̇_full_WX_B
+        B_J_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J_full_WL_B
+
+        match data.velocity_representation:
 
-    return O_J̇_WL_I
+            case VelRepr.Inertial:
+
+                B_X_W = jaxsim.math.Adjoint.from_transform(
+                    transform=W_H_B, inverse=True
+                )
+
+                W_v_WB = data.base_velocity()
+                B_Ẋ_W = -B_X_W @ jaxsim.math.Cross.vx(W_v_WB)
+
+                # Compute the operator to change the representation of ν, and its
+                # time derivative.
+                T = jax.scipy.linalg.block_diag(B_X_W, In)
+                Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_W, On)
+
+            case VelRepr.Body:
+
+                B_X_B = jaxsim.math.Adjoint.from_rotation_and_translation(
+                    translation=jnp.zeros(3), rotation=jnp.eye(3)
+                )
+
+                B_Ẋ_B = jnp.zeros(shape=(6, 6))
+
+                # Compute the operator to change the representation of ν, and its
+                # time derivative.
+                T = jax.scipy.linalg.block_diag(B_X_B, In)
+                Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_B, On)
+
+            case VelRepr.Mixed:
+
+                BW_H_B = W_H_B.at[0:3, 3].set(jnp.zeros(3))
+                B_X_BW = jaxsim.math.Adjoint.from_transform(
+                    transform=BW_H_B, inverse=True
+                )
+
+                BW_v_WB = data.base_velocity()
+                BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
+
+                BW_v_BW_B = BW_v_WB - BW_v_W_BW
+                B_Ẋ_BW = -B_X_BW @ jaxsim.math.Cross.vx(BW_v_BW_B)
+
+                # Compute the operator to change the representation of ν, and its
+                # time derivative.
+                T = jax.scipy.linalg.block_diag(B_X_BW, In)
+                Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_BW, On)
+
+            case _:
+                raise ValueError(data.velocity_representation)
+
+        # ======================================================
+        # Compute quantities to adjust the output representation
+        # ======================================================
+
+        match output_vel_repr:
+
+            case VelRepr.Inertial:
+
+                O_X_B = W_X_B = jaxsim.math.Adjoint.from_transform(transform=W_H_B)
+
+                with data.switch_velocity_representation(VelRepr.Body):
+                    B_v_WB = data.base_velocity()
+
+                O_Ẋ_B = W_Ẋ_B = W_X_B @ jaxsim.math.Cross.vx(B_v_WB)  # noqa: F841
+
+            case VelRepr.Body:
+
+                O_X_B = L_X_B = jaxsim.math.Adjoint.from_transform(
+                    transform=B_H_L, inverse=True
+                )
+
+                B_X_L = jaxsim.math.Adjoint.inverse(adjoint=L_X_B)
+
+                with data.switch_velocity_representation(VelRepr.Body):
+                    B_v_WB = data.base_velocity()
+                    L_v_WL = L_X_B @ B_J_WL_B @ data.generalized_velocity()
+
+                O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
+                    B_X_L @ L_v_WL - B_v_WB
+                )
+
+            case VelRepr.Mixed:
+
+                W_H_L = W_H_B @ B_H_L
+                LW_H_L = W_H_L.at[0:3, 3].set(jnp.zeros(3))
+                LW_H_B = LW_H_L @ jaxsim.math.Transform.inverse(B_H_L)
+
+                O_X_B = LW_X_B = jaxsim.math.Adjoint.from_transform(transform=LW_H_B)
+
+                B_X_LW = jaxsim.math.Adjoint.inverse(adjoint=LW_X_B)
+
+                with data.switch_velocity_representation(VelRepr.Body):
+                    B_v_WB = data.base_velocity()
+
+                with data.switch_velocity_representation(VelRepr.Mixed):
+                    BW_H_B = W_H_B.at[0:3, 3].set(jnp.zeros(3))
+                    B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
+                    LW_v_WL = LW_X_B @ (
+                        B_J_WL_B
+                        @ jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
+                        @ data.generalized_velocity()
+                    )
+                    LW_v_W_LW = LW_v_WL.at[3:6].set(jnp.zeros(3))
+
+                LW_v_LW_L = LW_v_WL - LW_v_W_LW
+                LW_v_B_LW = LW_v_WL - LW_X_B @ B_v_WB - LW_v_LW_L
+
+                O_Ẋ_B = LW_Ẋ_B = -LW_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
+                    B_X_LW @ LW_v_B_LW
+                )
+            case _:
+                raise ValueError(output_vel_repr)
+
+        # =============================================================
+        # Express the Jacobian derivative in the target representations
+        # =============================================================
+
+        # Sum all the components that form the Jacobian derivative in the target
+        # input/output velocity representations.
+        O_J̇_WL_I = jnp.zeros(shape=(6, 6 + model.dofs()))
+        O_J̇_WL_I += O_Ẋ_B @ B_J_WL_B @ T
+        O_J̇_WL_I += O_X_B @ B_J̇_WL_B @ T
+        O_J̇_WL_I += O_X_B @ B_J_WL_B @ Ṫ
+
+        return O_J̇_WL_I
+
+    return _compute_row(B_H_L, B_J_full_WL_B, W_H_B, κb)
 
 
 @functools.partial(jax.jit, static_argnames=["prefer_aba"])
@@ -1064,11 +1223,7 @@ def free_floating_coriolis_matrix(
         L_J_WL_B = generalized_free_floating_jacobian(model=model, data=data)
 
         # Doubly-left free-floating Jacobian derivative.
-        L_J̇_WL_B = jax.vmap(
-            lambda link_index: js.link.jacobian_derivative(
-                model=model, data=data, link_index=link_index
-            )
-        )(jnp.arange(model.number_of_links()))
+        L_J̇_WL_B = generalized_free_floating_jacobian_derivative(model=model, data=data)
 
     L_M_L = link_spatial_inertia_matrices(model=model)
 

jaxsim/api/references.py

@@ -193,8 +193,11 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             return self.input.physics_model.f_ext
 
         # If we have the model, we can extract the link names, if not provided.
-        link_names = link_names if link_names is not None else model.link_names()
-        link_idxs = js.link.names_to_idxs(link_names=link_names, model=model)
+        link_idxs = (
+            js.link.names_to_idxs(link_names=link_names, model=model)
+            if link_names is not None
+            else jnp.arange(model.number_of_links())
+        )
 
         # In inertial-fixed representation, we already have the link forces.
         if self.velocity_representation is VelRepr.Inertial:
@@ -267,8 +270,11 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             msg = "The actuation object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
-        joint_idxs = js.joint.names_to_idxs(joint_names=joint_names, model=model)
+        joint_idxs = (
+            js.joint.names_to_idxs(joint_names=joint_names, model=model)
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
 
         return jnp.atleast_1d(
             self.input.physics_model.tau[joint_idxs].squeeze()
@@ -318,8 +324,11 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             msg = "The references object is not compatible with the provided model"
             raise ValueError(msg)
 
-        joint_names = joint_names if joint_names is not None else model.joint_names()
-        joint_idxs = js.joint.names_to_idxs(joint_names=joint_names, model=model)
+        joint_idxs = (
+            js.joint.names_to_idxs(joint_names=joint_names, model=model)
+            if joint_names is not None
+            else jnp.arange(model.number_of_joints())
+        )
 
         return replace(forces=self.input.physics_model.tau.at[joint_idxs].set(forces))
 
@@ -388,18 +397,16 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
             return replace(forces=W_f0_L + W_f_L)
 
-        # If we have the model, we can extract the link names if not provided.
-        link_names = link_names if link_names is not None else model.link_names()
-
-        # Make sure that the link names are a tuple if they are provided by the user.
-        link_names = (link_names,) if isinstance(link_names, str) else link_names
-
-        if len(link_names) != f_L.shape[0]:
+        if link_names is not None and len(link_names) != f_L.shape[0]:
             msg = "The number of link names ({}) must match the number of forces ({})"
             raise ValueError(msg.format(len(link_names), f_L.shape[0]))
 
         # Extract the link indices.
-        link_idxs = js.link.names_to_idxs(link_names=link_names, model=model)
+        link_idxs = (
+            js.link.names_to_idxs(link_names=link_names, model=model)
+            if link_names is not None
+            else jnp.arange(model.number_of_links())
+        )
 
         # Compute the bias depending on whether we either set or add the link forces.
         W_f0_L = (
@@ -480,22 +487,21 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         f_F = jnp.atleast_2d(forces).astype(float)
 
-        # If we have the model, we can extract the frame names if not provided.
-        frame_names = frame_names if frame_names is not None else model.frame_names()
-
-        # Make sure that the frame names are a tuple if they are provided by the user.
-        frame_names = (frame_names,) if isinstance(frame_names, str) else frame_names
-
         if len(frame_names) != f_F.shape[0]:
             msg = "The number of frame names ({}) must match the number of forces ({})"
             raise ValueError(msg.format(len(frame_names), f_F.shape[0]))
 
         # Extract the frame indices.
-        frame_idxs = js.frame.names_to_idxs(frame_names=frame_names, model=model)
-        parent_link_idxs = jax.vmap(js.frame.idx_of_parent_link, in_axes=(None,))(
-            model, frame_index=frame_idxs
+        frame_idxs = (
+            js.frame.names_to_idxs(frame_names=frame_names, model=model)
+            if frame_names is not None
+            else jnp.arange(len(model.frame_names()))
         )
 
+        parent_link_idxs = jnp.array(model.kin_dyn_parameters.frame_parameters.body)[
+            frame_idxs - model.number_of_links()
+        ]
+
         exceptions.raise_value_error_if(
             condition=jnp.logical_not(data.valid(model=model)),
             msg="The provided data is not valid for the model",

jaxsim/integrators/common.py

@@ -319,9 +319,8 @@ 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(
-            lambda l: jnp.repeat(jnp.zeros_like(l)[jnp.newaxis, ...], c.size, axis=0),
-            x0,
+        carry0 = jax.tree_map(
+            lambda l: jnp.zeros((c.size, *l.shape), dtype=l.dtype), x0
         )
 
         # Closure on metadata to either evaluate the dynamics at the initial state

jaxsim/math/adjoint.py

@@ -137,11 +137,12 @@ class Adjoint:
             jtp.Matrix: The inverse adjoint matrix.
         """
         A_X_B = adjoint
-        A_H_B = Adjoint.to_transform(adjoint=A_X_B)
 
-        A_R_B = A_H_B[0:3, 0:3]
-        A_o_B = A_H_B[0:3, 3]
+        A_R_B = A_X_B[0:3, 0:3]
 
-        return Adjoint.from_rotation_and_translation(
-            rotation=A_R_B, translation=A_o_B, inverse=True
+        return jnp.vstack(
+            [
+                jnp.block([A_R_B.T, -A_R_B.T @ A_X_B[0:3, 3:6] @ A_R_B.T]),
+                jnp.block([jnp.zeros(shape=(3, 3)), A_R_B.T]),
+            ]
         )

jaxsim/parsers/rod/parser.py

@@ -113,7 +113,7 @@ def extract_model_data(
     links = [
         descriptions.LinkDescription(
             name=l.name,
-            mass=jnp.float32(l.inertial.mass),
+            mass=float(l.inertial.mass),
             inertia=utils.from_sdf_inertial(inertial=l.inertial),
             pose=l.pose.transform() if l.pose is not None else np.eye(4),
         )

jaxsim/rbda/contacts/rigid.py

@@ -437,29 +437,25 @@ class RigidContacts(ContactModel):
     def _compute_ineq_constraint_matrix(
         inactive_collidable_points: jtp.Vector, mu: jtp.FloatLike
     ) -> jtp.Matrix:
-
-        def compute_G_single_point(mu: float, c: float) -> jtp.Matrix:
-            """
-            Compute the inequality constraint matrix for a single collidable point
-            Rows 0-3: enforce the friction pyramid constraint,
-            Row 4: last one is for the non negativity of the vertical force
-            Row 5: contact complementarity condition
-            """
-            G_single_point = jnp.array(
-                [
-                    [1, 0, -mu],
-                    [0, 1, -mu],
-                    [-1, 0, -mu],
-                    [0, -1, -mu],
-                    [0, 0, -1],
-                    [0, 0, c],
-                ]
-            )
-            return G_single_point
-
-        G = jax.vmap(compute_G_single_point, in_axes=(None, 0))(
-            mu, inactive_collidable_points
+        """
+        Compute the inequality constraint matrix for a single collidable point
+        Rows 0-3: enforce the friction pyramid constraint,
+        Row 4: last one is for the non negativity of the vertical force
+        Row 5: contact complementarity condition
+        """
+        G_single_point = jnp.array(
+            [
+                [1, 0, -mu],
+                [0, 1, -mu],
+                [-1, 0, -mu],
+                [0, -1, -mu],
+                [0, 0, -1],
+                [0, 0, 0],
+            ]
         )
+        G = jnp.tile(G_single_point, (len(inactive_collidable_points), 1, 1))
+        G = G.at[:, 5, 2].set(inactive_collidable_points)
+
         G = jax.scipy.linalg.block_diag(*G)
         return G
 
@@ -498,28 +494,8 @@ class RigidContacts(ContactModel):
         D: jtp.FloatLike,
     ) -> jtp.Array:
 
-        def baumgarte_stabilization_of_single_point(
-            inactive: jtp.BoolLike,
-            δ: jtp.FloatLike,
-            δ_dot: jtp.FloatLike,
-            n: jtp.ArrayLike,
-            k_baumgarte: jtp.FloatLike,
-            d_baumgarte: jtp.FloatLike,
-        ) -> jtp.Array:
-
-            baumgarte_term = jax.lax.cond(
-                inactive,
-                lambda δ, δ_dot, n, K, D: jnp.zeros(3),
-                # This is equivalent to: K*(pT - p)⋅n̂ + D*(0 - v)⋅n̂,
-                # where pT is the point on the terrain surface vertical to p.
-                lambda δ, δ_dot, n, K, D: (K * δ + D * δ_dot) * n,
-                *(δ, δ_dot, n, k_baumgarte, d_baumgarte),
-            )
-
-            return baumgarte_term
-
-        baumgarte_term = jax.vmap(
-            baumgarte_stabilization_of_single_point, in_axes=(0, 0, 0, 0, None, None)
-        )(inactive_collidable_points, δ, δ_dot, n, K, D)
-
-        return baumgarte_term
+        return jnp.where(
+            inactive_collidable_points[:, jnp.newaxis],
+            jnp.zeros_like(n),
+            (K * δ + D * δ_dot)[:, jnp.newaxis] * n,
+        )

jaxsim/rbda/contacts/visco_elastic.py

@@ -407,10 +407,12 @@ class ViscoElasticContacts(common.ContactModel):
         # Compute the contact point dynamics
         # ==================================
 
-        p_t0 = js.contact.collidable_point_positions(model, data)[indices, :]
-        v_t0 = js.contact.collidable_point_velocities(model, data)[indices, :]
+        p_t0, v_t0 = js.contact.collidable_point_kinematics(model, data)
         m_t0 = data.state.extended["tangential_deformation"][indices, :]
 
+        p_t0 = p_t0[indices, :]
+        v_t0 = v_t0[indices, :]
+
         # Compute the linearized contact dynamics.
         # Note that it linearizes the (non-linear) contact model at (p, v, m)[t0].
         A, b, A_sc, b_sc = ViscoElasticContacts._contact_points_dynamics(

{jaxsim-0.4.3.dev282.dist-info → jaxsim-0.4.3.dev298.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev282
+Version: 0.4.3.dev298
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev282.dist-info → jaxsim-0.4.3.dev298.dist-info}/RECORD

@@ -1,27 +1,27 @@
 jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
-jaxsim/_version.py,sha256=8Lzzuj01bYiR6tHmTBRLQW2sSRCWlBO6B-razI_HdjU,428
+jaxsim/_version.py,sha256=Pwqoke-57lWzujPjc0ocPWyiVHY0CgqE-ALFKL9bumE,428
 jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
 jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
-jaxsim/api/contact.py,sha256=iiQc8vdlJ60Vl-ki4ieP7TW6MJhJfWREYpcVyyEWT0M,24034
-jaxsim/api/data.py,sha256=gQX6hfEaw0ooJYvpr5f8UvEJwqhtflEK_NHWn9XgTZY,28935
+jaxsim/api/contact.py,sha256=EM2pCPpFgbk0sjWBWZFYweRHfDNPa2jfnOZP6qjg-XY,24108
+jaxsim/api/data.py,sha256=ThRpoBlbdwf1N3xs8SWrY5d8RbfdYRwFcmkdIPgtee4,29004
 jaxsim/api/frame.py,sha256=yPSgNygHkvWlln4wShNt7vZm_fFobVEm7phsklNNyH8,12922
-jaxsim/api/joint.py,sha256=iz7kHLLXMjlDOjKN0T6tcutWCnpnFAOA2IEE9GuxlQI,7304
-jaxsim/api/kin_dyn_parameters.py,sha256=kbDN5n9uj8CamVJXk1U5oYLbxyjaWDIeUG0V68DCEFs,29578
-jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
-jaxsim/api/model.py,sha256=dpQZDT0BodMfK1wmpG-STFh-rFsJStobQ1fhrWILK9o,73410
+jaxsim/api/joint.py,sha256=Vl9VJs66_es88zjBBXqPzKkA3oAktV-DNiTkXwxOSsI,7562
+jaxsim/api/kin_dyn_parameters.py,sha256=fRB7LPGTIwqunJGSq5YfMwSuy2yM_c5gZbsjN1K19Jw,29482
+jaxsim/api/link.py,sha256=au47jV7bNdH2itvZ4qngKZyLi0UTIKzqVjjKosnlMsU,12858
+jaxsim/api/model.py,sha256=ynzI2Us2XNVgcz_rasxVjE5575MC04YINsAxNP-B5sU,79472
 jaxsim/api/ode.py,sha256=_t18avoCJngQk6eMFTGpaeahbpchQP20qJnUOCPkz8s,15360
 jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
-jaxsim/api/references.py,sha256=fW77LitZ8DYgT6ZmUInJfm5luBV1mTcqcNRiC_i79og,20862
+jaxsim/api/references.py,sha256=crPzgUCDHkTKUuXQcj9ygFAI0pewIx-fTPEsU8fBiU4,20555
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=_FZs7E0EazERGA3K0tGC1baUrs8sBDzYTf2U2mFYh9s,18329
+jaxsim/integrators/common.py,sha256=8Y3PGVOydPBJrqYtH33rtng48T12ySWfXTk-ED_2wao,18297
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
 jaxsim/integrators/variable_step.py,sha256=hGYKG3Sq3QITgzIePmCVCrrirwagqsKnB3aYifAcKR4,22848
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
-jaxsim/math/adjoint.py,sha256=o1FCipkGwPtMbN2gFNIyUV8ADF3TX5fxElpTEXK0bIs,4377
+jaxsim/math/adjoint.py,sha256=V7r5VrTCKPLEL5gavNSx9U7xSsrb11a5e4gWqJ2MuRo,4375
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
 jaxsim/math/inertia.py,sha256=01hz6wMFreN2jBA0rVoBS1YMVh77KvwuzXSOpI3pxNk,1614
 jaxsim/math/joint_model.py,sha256=EzAveaG5B6ZnCFNUzN30KEQUVesd83lfWXJarYR-kUw,9989
@@ -43,7 +43,7 @@ jaxsim/parsers/descriptions/joint.py,sha256=VSb6C0FBBKMqwrHBKfc-Bbn4rl_J0RzUxMQl
 jaxsim/parsers/descriptions/link.py,sha256=Eh0W5qL7_Uw0GV-BkNKXhm9Q2dRTfIWCX5D-87zQkxA,3711
 jaxsim/parsers/descriptions/model.py,sha256=I2Vsbv8Josl4Le7b5rIvhqA2k9Bbv5JxMqwytayxds0,9833
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=P3QFPdI2NBxL4TALcPhL6Exij9-Q9W716Aps2cb8s2I,14041
+jaxsim/parsers/rod/parser.py,sha256=UAL6vkAab2y6PgIv3tAL8UodNbq-mmNBuBotxHJOSVU,14035
 jaxsim/parsers/rod/utils.py,sha256=5DsF3OeePZGidOJ5GiFSZx-51uIdnFvMW9EK6SgOW6Q,5698
 jaxsim/rbda/__init__.py,sha256=kmy4G9aMkrqPNGdLSaSV3k15dpF52vBEUQXDFDuKIxU,337
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
@@ -56,17 +56,17 @@ jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
 jaxsim/rbda/contacts/__init__.py,sha256=L5MM-2pv76YPGzxExdz2EErgGBATuAjYnNHlq5QOySs,503
 jaxsim/rbda/contacts/common.py,sha256=ELASrMnkFuUZBRDPKu-HO7ymUukK3r-VZw9TjqOSZgM,11174
 jaxsim/rbda/contacts/relaxed_rigid.py,sha256=6EWTkHjixSv-9pBl7aO_abOxHrneUQKF6uUyt5Vbbuc,20512
-jaxsim/rbda/contacts/rigid.py,sha256=I_FrJv4rK9Lzj4i0EMaZHxOGyrPT_XuQUg6ZUzvp4y8,17943
+jaxsim/rbda/contacts/rigid.py,sha256=q9hXH2FpqAAf4UVt_BPUQ5DrLgitEARIu4j3kKr3Wpg,17005
 jaxsim/rbda/contacts/soft.py,sha256=wqwnsg0abWn64pYG7je08fsBwdgyJdd2k7dAG0Sx0_k,16432
-jaxsim/rbda/contacts/visco_elastic.py,sha256=S1GsYZQa8vMS1EtGmS7_Ql07h2AhJEmO67iShkGWmIU,41443
+jaxsim/rbda/contacts/visco_elastic.py,sha256=vQkfMuqQ3Qu8nbDTPY4jWBZjV3U7qtoRK1Aya3O3oFA,41424
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=K91HEzPqTSyNrc_j1KfAAEF_5oDeuk_-jnnZGrcMEcY,5015
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev282.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev282.dist-info/METADATA,sha256=Llp_WIbsROSnrBpSMJUdFnAOOkLDobo5f87CO8xvOuA,17276
-jaxsim-0.4.3.dev282.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-jaxsim-0.4.3.dev282.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev282.dist-info/RECORD,,
+jaxsim-0.4.3.dev298.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev298.dist-info/METADATA,sha256=2QuLbV-zLRmC7baPpENbJ9CXPiOoYuEzZEmpovjyQtw,17276
+jaxsim-0.4.3.dev298.dist-info/WHEEL,sha256=a7TGlA-5DaHMRrarXjVbQagU3Man_dCnGIWMJr5kRWo,91
+jaxsim-0.4.3.dev298.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev298.dist-info/RECORD,,

{jaxsim-0.4.3.dev282.dist-info → jaxsim-0.4.3.dev298.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.3.0)
+Generator: setuptools (75.4.0)
 Root-Is-Purelib: true
 Tag: py3-none-any