jaxsim 0.6.2.dev182__py3-none-any.whl → 0.6.2.dev225__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/__init__.py

@@ -4,7 +4,6 @@ from . import (
     actuation_model,
     com,
     contact,
-    contact_model,
     frame,
     integrators,
     joint,

jaxsim/api/com.py

@@ -301,9 +301,7 @@ def bias_acceleration(
             C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL  # noqa: F841
             C_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
 
-            L_H_C = L_H_W = jax.vmap(  # noqa: F841
-                lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
-            )(W_H_L)
+            L_H_C = L_H_W = jax.vmap(jaxsim.math.Transform.inverse)(W_H_L)  # noqa: F841
 
             L_v_LC = L_v_LW = jax.vmap(  # noqa: F841
                 lambda i: -js.link.velocity(

jaxsim/api/common.py

@@ -121,14 +121,8 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             The 6D quantity in the other representation.
         """
 
-        W_array = array.squeeze()
-        W_H_O = transform.squeeze()
-
-        if W_array.size != 6:
-            raise ValueError(W_array.size, 6)
-
-        if W_H_O.shape != (4, 4):
-            raise ValueError(W_H_O.shape, (4, 4))
+        W_array = array
+        W_H_O = transform
 
         match other_representation:
 
@@ -139,25 +133,24 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
 
                 if not is_force:
                     O_Xv_W = Adjoint.from_transform(transform=W_H_O, inverse=True)
-                    O_array = O_Xv_W @ W_array
+                    O_array = jnp.einsum("...ij,...j->...i", O_Xv_W, W_array)
 
                 else:
-                    O_Xf_W = Adjoint.from_transform(transform=W_H_O).T
-                    O_array = O_Xf_W @ W_array
+                    O_Xf_W = Adjoint.from_transform(transform=W_H_O).swapaxes(-1, -2)
+                    O_array = jnp.einsum("...ij,...j->...i", O_Xf_W, W_array)
 
                 return O_array
 
             case VelRepr.Mixed:
-                W_p_O = W_H_O[0:3, 3]
-                W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
+                W_H_OW = W_H_O.at[..., 0:3, 0:3].set(jnp.eye(3))
 
                 if not is_force:
                     OW_Xv_W = Adjoint.from_transform(transform=W_H_OW, inverse=True)
-                    OW_array = OW_Xv_W @ W_array
+                    OW_array = jnp.einsum("...ij,...j->...i", OW_Xv_W, W_array)
 
                 else:
-                    OW_Xf_W = Adjoint.from_transform(transform=W_H_OW).T
-                    OW_array = OW_Xf_W @ W_array
+                    OW_Xf_W = Adjoint.from_transform(transform=W_H_OW).swapaxes(-1, -2)
+                    OW_array = jnp.einsum("...ij,...j->...i", OW_Xf_W, W_array)
 
                 return OW_array
 
@@ -188,45 +181,40 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             The 6D quantity in the inertial-fixed representation.
         """
 
-        W_array = array.squeeze()
-        W_H_O = transform.squeeze()
-
-        if W_array.size != 6:
-            raise ValueError(W_array.size, 6)
-
-        if W_H_O.shape != (4, 4):
-            raise ValueError(W_H_O.shape, (4, 4))
+        O_array = array
+        W_H_O = transform
 
         match other_representation:
             case VelRepr.Inertial:
-                W_array = array
-                return W_array
+                return O_array
 
             case VelRepr.Body:
-                O_array = array
 
                 if not is_force:
-                    W_Xv_O: jtp.Array = Adjoint.from_transform(W_H_O)
-                    W_array = W_Xv_O @ O_array
+                    W_Xv_O = Adjoint.from_transform(W_H_O)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xv_O, O_array)
 
                 else:
-                    W_Xf_O = Adjoint.from_transform(transform=W_H_O, inverse=True).T
-                    W_array = W_Xf_O @ O_array
+                    W_Xf_O = Adjoint.from_transform(
+                        transform=W_H_O, inverse=True
+                    ).swapaxes(-1, -2)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xf_O, O_array)
 
                 return W_array
 
             case VelRepr.Mixed:
-                BW_array = array
-                W_p_O = W_H_O[0:3, 3]
-                W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
+
+                W_H_OW = W_H_O.at[..., 0:3, 0:3].set(jnp.eye(3))
 
                 if not is_force:
-                    W_Xv_BW: jtp.Array = Adjoint.from_transform(W_H_OW)
-                    W_array = W_Xv_BW @ BW_array
+                    W_Xv_BW = Adjoint.from_transform(W_H_OW)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xv_BW, O_array)
 
                 else:
-                    W_Xf_BW = Adjoint.from_transform(transform=W_H_OW, inverse=True).T
-                    W_array = W_Xf_BW @ BW_array
+                    W_Xf_BW = Adjoint.from_transform(
+                        transform=W_H_OW, inverse=True
+                    ).swapaxes(-1, -2)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xf_BW, O_array)
 
                 return W_array
 

jaxsim/api/contact.py

@@ -11,7 +11,7 @@ import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim import logging
 from jaxsim.math import Adjoint, Cross, Transform
-from jaxsim.rbda import contacts
+from jaxsim.rbda.contacts import SoftContacts
 
 from .common import VelRepr
 
@@ -37,14 +37,11 @@ def collidable_point_kinematics(
         the linear component of the mixed 6D frame velocity.
     """
 
-    # Switch to inertial-fixed since the RBDAs expect velocities in this representation.
-    with data.switch_velocity_representation(VelRepr.Inertial):
-
-        W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
-            model=model,
-            link_transforms=data._link_transforms,
-            link_velocities=data._link_velocities,
-        )
+    W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
+        model=model,
+        link_transforms=data._link_transforms,
+        link_velocities=data._link_velocities,
+    )
 
     return W_p_Ci, W_ṗ_Ci
 
@@ -164,18 +161,23 @@ def estimate_good_soft_contacts_parameters(
 def estimate_good_contact_parameters(
     model: js.model.JaxSimModel,
     *,
+    standard_gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
     static_friction_coefficient: jtp.FloatLike = 0.5,
-    **kwargs,
+    number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
+    damping_ratio: jtp.FloatLike = 1.0,
+    max_penetration: jtp.FloatLike | None = None,
 ) -> jaxsim.rbda.contacts.ContactParamsTypes:
     """
     Estimate good contact parameters.
 
     Args:
         model: The model to consider.
+        standard_gravity: The standard gravity acceleration.
         static_friction_coefficient: The static friction coefficient.
-        kwargs:
-            Additional model-specific parameters passed to the builder method of
-            the parameters class.
+        number_of_active_collidable_points_steady_state:
+            The number of active collidable points in steady state.
+        damping_ratio: The damping ratio.
+        max_penetration: The maximum penetration allowed.
 
     Returns:
         The estimated good contacts parameters.
@@ -190,20 +192,41 @@ def estimate_good_contact_parameters(
         specific application.
     """
 
-    match model.contact_model:
+    def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
+        """
+        Displacement between the CoM and the lowest collidable point using zero
+        joint positions.
+        """
+
+        zero_data = js.data.JaxSimModelData.build(
+            model=model,
+        )
+
+        W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
 
-        case contacts.RelaxedRigidContacts():
-            assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
+        if model.floating_base():
+            W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
+            return 2 * (W_pz_CoM - W_pz_C.min())
 
-            parameters = contacts.RelaxedRigidContactsParams.build(
-                mu=static_friction_coefficient,
-                **kwargs,
-            )
+        return 2 * W_pz_CoM
 
-        case _:
-            raise ValueError(f"Invalid contact model: {model.contact_model}")
+    max_δ = (
+        max_penetration
+        if max_penetration is not None
+        # Consider as default a 0.5% of the model height.
+        else 0.005 * estimate_model_height(model=model)
+    )
 
-    return parameters
+    nc = number_of_active_collidable_points_steady_state
+
+    return model.contact_model._parameters_class().build_default_from_jaxsim_model(
+        model=model,
+        standard_gravity=standard_gravity,
+        static_friction_coefficient=static_friction_coefficient,
+        max_penetration=max_δ,
+        number_of_active_collidable_points_steady_state=nc,
+        damping_ratio=damping_ratio,
+    )
 
 
 @jax.jit
@@ -244,7 +267,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
 
     # 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))(L_p_Ci)
+    L_H_C = jax.vmap(jnp.eye(4).at[0:3, 3].set)(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)
@@ -504,3 +527,96 @@ def jacobian_derivative(
     )
 
     return O_J̇_WC
+
+
+@jax.jit
+@js.common.named_scope
+def link_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_torques: jtp.VectorLike | None = None,
+) -> tuple[jtp.Matrix, dict[str, jtp.Matrix]]:
+    """
+    Compute the 6D contact forces of all links of the model in inertial representation.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_forces:
+            The 6D external forces to apply to the links expressed in inertial representation
+        joint_torques:
+            The joint torques acting on the joints.
+
+    Returns:
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
+        expressed in inertial representation.
+    """
+
+    # Compute the contact forces for each collidable point with the active contact model.
+    W_f_C, aux_dict = model.contact_model.compute_contact_forces(
+        model=model,
+        data=data,
+        **(
+            dict(link_forces=link_forces, joint_force_references=joint_torques)
+            if not isinstance(model.contact_model, SoftContacts)
+            else {}
+        ),
+    )
+
+    # Compute the 6D forces applied to the links equivalent to the forces applied
+    # to the frames associated to the collidable points.
+    W_f_L = link_forces_from_contact_forces(model=model, contact_forces=W_f_C)
+
+    return W_f_L, aux_dict
+
+
+@staticmethod
+def link_forces_from_contact_forces(
+    model: js.model.JaxSimModel,
+    *,
+    contact_forces: jtp.MatrixLike,
+) -> jtp.Matrix:
+    """
+    Compute the link forces from the contact forces.
+
+    Args:
+        model: The robot model considered by the contact model.
+        contact_forces: The contact forces computed by the contact model.
+
+    Returns:
+        The 6D contact forces applied to the links and expressed in the frame of
+        the velocity representation of data.
+    """
+
+    # Get the object storing the contact parameters of the model.
+    contact_parameters = model.kin_dyn_parameters.contact_parameters
+
+    # Extract the indices corresponding to the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    # Convert the contact forces to a JAX array.
+    W_f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
+
+    # Construct the vector defining the parent link index of each collidable point.
+    # We use this vector to sum the 6D forces of all collidable points rigidly
+    # attached to the same link.
+    parent_link_index_of_collidable_points = jnp.array(
+        contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    # Create the mask that associate each collidable point to their parent link.
+    # We use this mask to sum the collidable points to the right link.
+    mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+        model.number_of_links()
+    )
+
+    # Sum the forces of all collidable points rigidly attached to a body.
+    # Since the contact forces W_f_C are expressed in the world frame,
+    # we don't need any coordinate transformation.
+    W_f_L = mask.T @ W_f_C
+
+    return W_f_L

jaxsim/api/data.py

@@ -5,9 +5,9 @@ import functools
 from collections.abc import Sequence
 
 try:
-    from typing import override
+    from typing import Self, override
 except ImportError:
-    from typing_extensions import override
+    from typing_extensions import override, Self
 
 import jax
 import jax.numpy as jnp
@@ -22,11 +22,6 @@ import jaxsim.typing as jtp
 from . import common
 from .common import VelRepr
 
-try:
-    from typing import Self
-except ImportError:
-    from typing_extensions import Self
-
 
 @jax_dataclasses.pytree_dataclass
 class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
@@ -64,6 +59,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     _link_transforms: jtp.Matrix = dataclasses.field(repr=False, default=None)
     _link_velocities: jtp.Matrix = dataclasses.field(repr=False, default=None)
 
+    # Extended state for soft and rigid contact models.
+    contact_state: dict[str, jtp.Array] = dataclasses.field(default=None)
+
     @staticmethod
     def build(
         model: js.model.JaxSimModel,
@@ -73,6 +71,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         base_linear_velocity: jtp.VectorLike | None = None,
         base_angular_velocity: jtp.VectorLike | None = None,
         joint_velocities: jtp.VectorLike | None = None,
+        contact_state: dict[str, jtp.Array] | None = None,
         velocity_representation: VelRepr = VelRepr.Mixed,
     ) -> JaxSimModelData:
         """
@@ -89,6 +88,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 The base angular velocity in the selected representation.
             joint_velocities: The joint velocities.
             velocity_representation: The velocity representation to use. It defaults to mixed if not provided.
+            contact_state: The optional contact state.
 
         Returns:
             A `JaxSimModelData` initialized with the given state.
@@ -171,6 +171,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             )
         )
 
+        contact_state = contact_state or {}
+
+        if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
+            contact_state.setdefault(
+                "tangential_deformation",
+                jnp.zeros_like(model.kin_dyn_parameters.contact_parameters.point),
+            )
+
         model_data = JaxSimModelData(
             velocity_representation=velocity_representation,
             _base_quaternion=base_quaternion,
@@ -183,6 +191,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             _joint_transforms=joint_transforms,
             _link_transforms=link_transforms,
             _link_velocities=link_velocities_inertial,
+            contact_state=contact_state,
         )
 
         if not model_data.valid(model=model):
@@ -265,14 +274,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         """
 
         # Extract the base quaternion.
-        W_Q_B = self.base_quaternion.squeeze()
+        W_Q_B = self.base_quaternion
 
         # Always normalize the quaternion to avoid numerical issues.
         # If the active scheme does not integrate the quaternion on its manifold,
         # 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.
-        norm = jaxsim.math.safe_norm(W_Q_B)
+        norm = jaxsim.math.safe_norm(W_Q_B, axis=-1, keepdims=True)
         W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
         return W_Q_B
 
@@ -285,11 +294,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             The base 6D velocity in the active representation.
         """
 
-        W_v_WB = jnp.hstack(
-            [
-                self._base_linear_velocity,
-                self._base_angular_velocity,
-            ]
+        W_v_WB = jnp.concatenate(
+            [self._base_linear_velocity, self._base_angular_velocity], axis=-1
         )
 
         W_H_B = self._base_transform
@@ -350,11 +356,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     @js.common.named_scope
     @jax.jit
-    def reset_base_quaternion(self, base_quaternion: jtp.VectorLike) -> Self:
+    def reset_base_quaternion(
+        self, model: js.model.JaxSimModel, base_quaternion: jtp.VectorLike
+    ) -> Self:
         """
         Reset the base quaternion.
 
         Args:
+            model: The JaxSim model to use.
             base_quaternion: The base orientation as a quaternion.
 
         Returns:
@@ -363,18 +372,21 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         W_Q_B = jnp.array(base_quaternion, dtype=float)
 
-        norm = jaxsim.math.safe_norm(W_Q_B)
+        norm = jaxsim.math.safe_norm(W_Q_B, axis=-1)
         W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
 
-        return self.replace(validate=True, base_quaternion=W_Q_B)
+        return self.replace(model=model, base_quaternion=W_Q_B)
 
     @js.common.named_scope
     @jax.jit
-    def reset_base_pose(self, base_pose: jtp.MatrixLike) -> Self:
+    def reset_base_pose(
+        self, model: js.model.JaxSimModel, base_pose: jtp.MatrixLike
+    ) -> Self:
         """
         Reset the base pose.
 
         Args:
+            model: The JaxSim model to use.
             base_pose: The base pose as an SE(3) matrix.
 
         Returns:
@@ -385,6 +397,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         W_p_B = base_pose[0:3, 3]
         W_Q_B = jaxsim.math.Quaternion.from_dcm(dcm=base_pose[0:3, 0:3])
         return self.replace(
+            model=model,
             base_position=W_p_B,
             base_quaternion=W_Q_B,
         )
@@ -399,11 +412,19 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         base_linear_velocity: jtp.Vector | None = None,
         base_angular_velocity: jtp.Vector | None = None,
         base_position: jtp.Vector | None = None,
+        *,
+        contact_state: dict[str, jtp.Array] | None = None,
         validate: bool = False,
     ) -> Self:
         """
         Replace the attributes of the `JaxSimModelData` object.
         """
+
+        # Extract the batch size.
+        batch_size = (
+            self._base_transform.shape[0] if self._base_transform.ndim > 2 else 1
+        )
+
         if joint_positions is None:
             joint_positions = self.joint_positions
         if joint_velocities is None:
@@ -412,6 +433,22 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_quaternion = self.base_quaternion
         if base_position is None:
             base_position = self.base_position
+        if contact_state is None:
+            contact_state = self.contact_state
+
+        if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
+            contact_state.setdefault(
+                "tangential_deformation",
+                jnp.zeros_like(model.kin_dyn_parameters.contact_parameters.point),
+            )
+
+        # Normalize the quaternion to avoid numerical issues.
+        base_quaternion_norm = jaxsim.math.safe_norm(
+            base_quaternion, axis=-1, keepdims=True
+        )
+        base_quaternion = base_quaternion / jnp.where(
+            base_quaternion_norm == 0, 1.0, base_quaternion_norm
+        )
 
         joint_positions = jnp.atleast_1d(joint_positions.squeeze()).astype(float)
         joint_velocities = jnp.atleast_1d(joint_velocities.squeeze()).astype(float)
@@ -421,44 +458,70 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         base_transform = jaxsim.math.Transform.from_quaternion_and_translation(
             translation=base_position, quaternion=base_quaternion
         )
-        joint_transforms = model.kin_dyn_parameters.joint_transforms(
-            joint_positions=joint_positions, base_transform=base_transform
+
+        joint_transforms = jax.vmap(model.kin_dyn_parameters.joint_transforms)(
+            joint_positions=jnp.broadcast_to(
+                joint_positions, (batch_size, model.dofs())
+            ),
+            base_transform=jnp.broadcast_to(base_transform, (batch_size, 4, 4)),
         )
 
         if base_linear_velocity is None and base_angular_velocity is None:
-            base_linear_velocity = self._base_linear_velocity
-            base_angular_velocity = self._base_angular_velocity
+            base_linear_velocity_inertial = self._base_linear_velocity
+            base_angular_velocity_inertial = self._base_angular_velocity
         else:
             if base_linear_velocity is None:
                 base_linear_velocity = self.base_velocity[:3]
             if base_angular_velocity is None:
                 base_angular_velocity = self.base_velocity[3:]
+
             base_linear_velocity = jnp.atleast_1d(base_linear_velocity.squeeze())
             base_angular_velocity = jnp.atleast_1d(base_angular_velocity.squeeze())
+
             W_v_WB = JaxSimModelData.other_representation_to_inertial(
                 array=jnp.hstack([base_linear_velocity, base_angular_velocity]),
                 other_representation=self.velocity_representation,
                 transform=base_transform,
                 is_force=False,
             ).astype(float)
-            base_linear_velocity, base_angular_velocity = W_v_WB[:3], W_v_WB[3:]
 
-        link_transforms, link_velocities = jaxsim.rbda.forward_kinematics_model(
-            model=model,
-            base_position=base_position,
-            base_quaternion=base_quaternion,
-            joint_positions=joint_positions,
-            joint_velocities=joint_velocities,
-            base_linear_velocity_inertial=base_linear_velocity,
-            base_angular_velocity_inertial=base_angular_velocity,
+            base_linear_velocity_inertial, base_angular_velocity_inertial = (
+                W_v_WB[..., :3],
+                W_v_WB[..., 3:],
+            )
+
+        link_transforms, link_velocities = jax.vmap(
+            jaxsim.rbda.forward_kinematics_model, in_axes=(None,)
+        )(
+            model,
+            base_position=jnp.broadcast_to(base_position, (batch_size, 3)),
+            base_quaternion=jnp.broadcast_to(base_quaternion, (batch_size, 4)),
+            joint_positions=jnp.broadcast_to(
+                joint_positions, (batch_size, model.dofs())
+            ),
+            joint_velocities=jnp.broadcast_to(
+                joint_velocities, (batch_size, model.dofs())
+            ),
+            base_linear_velocity_inertial=jnp.broadcast_to(
+                base_linear_velocity_inertial, (batch_size, 3)
+            ),
+            base_angular_velocity_inertial=jnp.broadcast_to(
+                base_angular_velocity_inertial, (batch_size, 3)
+            ),
         )
 
+        # Adjust the output shapes.
+        if batch_size == 1:
+            link_transforms = link_transforms.reshape(self._link_transforms.shape)
+            link_velocities = link_velocities.reshape(self._link_velocities.shape)
+            joint_transforms = joint_transforms.reshape(self._joint_transforms.shape)
+
         return super().replace(
             _joint_positions=joint_positions,
             _joint_velocities=joint_velocities,
             _base_quaternion=base_quaternion,
-            _base_linear_velocity=base_linear_velocity,
-            _base_angular_velocity=base_angular_velocity,
+            _base_linear_velocity=base_linear_velocity_inertial,
+            _base_angular_velocity=base_angular_velocity_inertial,
             _base_position=base_position,
             _base_transform=base_transform,
             _joint_transforms=joint_transforms,