Sum the forces on corresponding parent links

ami-iit · Jul 23, 2024 · 28801f6 · 28801f6
1 parent cfba098
commit 28801f6
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/jaxsim/api/references.py b/src/jaxsim/api/references.py
@@ -535,7 +535,11 @@ def convert_to_link_force(
 
             return F_X_L @ W_f_F
 
-        W_f_L = jax.vmap(convert_to_link_force)(W_f_F, W_H_Fi, parent_link_idxs)
+        W_f_L_i = jax.vmap(convert_to_link_force)(W_f_F, W_H_Fi, parent_link_idxs)
+
+        # Sum the forces on the parent links.
+        mask = parent_link_idxs[:, jnp.newaxis] == jnp.arange(model.number_of_links())
+        W_f_L = mask.T @ W_f_L_i
 
         return self.apply_link_forces(
             model=model,