Make Cross.vx_star compatible with batches

ami-iit · Nov 22, 2024 · ff5752a · ff5752a
1 parent 1393a82
commit ff5752a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/src/jaxsim/math/cross.py b/src/jaxsim/math/cross.py
@@ -20,9 +20,8 @@ def vx(velocity_sixd: jtp.Vector) -> jtp.Matrix:
         Raises:
             ValueError: If the input vector does not have a size of 6.
         """
-        velocity_sixd = velocity_sixd.reshape(-1, 6)
 
-        v, ω = jnp.split(velocity_sixd, 2, axis=-1)
+        v, ω = jnp.split(velocity_sixd.reshape(-1, 6), 2, axis=-1)
 
         v_cross = jnp.concatenate(
             [
@@ -50,5 +49,5 @@ def vx_star(velocity_sixd: jtp.Vector) -> jtp.Matrix:
         Raises:
             ValueError: If the input vector does not have a size of 6.
         """
-        v_cross_star = -Cross.vx(velocity_sixd).T
+        v_cross_star = -Cross.vx(velocity_sixd).swapaxes(-1, -2)
         return v_cross_star
diff --git a/src/jaxsim/rbda/rnea.py b/src/jaxsim/rbda/rnea.py
@@ -121,7 +121,7 @@ def rnea(
         # Force applied to the base link that produce the base acceleration w/o gravity.
         f_0 = (
             M[0] @ a[0]
-            + Cross.vx_star(v[0]) @ M[0] @ v[0]
+            + Cross.vx_star(v[0].squeeze()) @ M[0] @ v[0]
             - W_X_B.T @ jnp.vstack(W_f[0])
         )
         f = f.at[0].set(f_0)