Avoid unnecessary jax.vmap in RelaxedRigidContacts

src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -490,16 +490,12 @@ def continuing_criterion(carry: OptimizationCarry) -> jtp.Bool:
         CW_fl_C = solution.reshape(-1, 3)
 
         # Convert the contact forces from mixed to inertial-fixed representation.
-        W_f_C = jax.vmap(
-            lambda CW_fl_C, W_H_C: (
-                ModelDataWithVelocityRepresentation.other_representation_to_inertial(
-                    array=jnp.zeros(6).at[0:3].set(CW_fl_C),
-                    transform=W_H_C,
-                    other_representation=VelRepr.Mixed,
-                    is_force=True,
-                )
-            ),
-        )(CW_fl_C, W_H_C)
+        W_f_C = ModelDataWithVelocityRepresentation.other_representation_to_inertial(
+            array=jnp.hstack((CW_fl_C, jnp.zeros_like(CW_fl_C))),
+            transform=W_H_C,
+            other_representation=VelRepr.Mixed,
+            is_force=True,
+        )
 
         return W_f_C, {}