jaxsim 0.4.3.dev282__tar.gz → 0.4.3.dev295__tar.gz

{jaxsim-0.4.3.dev282 → jaxsim-0.4.3.dev295}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev282
+Version: 0.4.3.dev295
 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 → jaxsim-0.4.3.dev295}/src/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.dev295'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev295')

{jaxsim-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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-0.4.3.dev282 → jaxsim-0.4.3.dev295}/src/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]),
+            ]
         )