Merge pull request ami-iit#217 from ami-iit/patch-reference_frame_force

Fix `apply_frame_forces` in `references` module
flferretti · Jul 29, 2024 · 1145d1f · 1145d1f
2 parents 14b401c + ff85e66
commit 1145d1f
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/jaxsim/api/references.py b/src/jaxsim/api/references.py
@@ -493,9 +493,9 @@ def apply_frame_forces(
 
         # Extract the frame indices.
         frame_idxs = js.frame.names_to_idxs(frame_names=frame_names, model=model)
-        parent_link_idxs = js.frame.idx_of_parent_link(
-            model=model, frame_idxs=frame_idxs
-        )
+        parent_link_idxs = jax.vmap(
+            lambda frame_idx: js.frame.idx_of_parent_link, in_axes=(None,)
+        )(model, frame_idx=frame_idxs)
 
         exceptions.raise_value_error_if(
             condition=jnp.logical_not(data.valid(model=model)),
@@ -522,7 +522,7 @@ def to_inertial(f_F: jtp.MatrixLike, W_H_F: jtp.MatrixLike) -> jtp.Matrix:
                 W_f_F = f_F
 
             case VelRepr.Body | VelRepr.Mixed:
-                W_f_F = jax.vmap(to_inertial)(f_F, W_H_Fi, frame_idxs)
+                W_f_F = jax.vmap(to_inertial)(f_F, W_H_Fi)
 
             case _:
                 raise ValueError("Invalid velocity representation.")