Minor updated to kin_dyn_parameters

ami-iit · Nov 14, 2024 · b3c59ea · b3c59ea
1 parent b76d7a7
commit b3c59ea
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/jaxsim/api/kin_dyn_parameters.py b/src/jaxsim/api/kin_dyn_parameters.py
@@ -438,7 +438,9 @@ def joint_transforms_and_motion_subspaces(
     # Helpers to update parameters
     # ============================
 
-    def set_link_mass(self, link_index: int, mass: jtp.FloatLike) -> KynDynParameters:
+    def set_link_mass(
+        self, link_index: jtp.IntLike, mass: jtp.FloatLike
+    ) -> KynDynParameters:
         """
         Set the mass of a link.
 
@@ -457,7 +459,7 @@ def set_link_mass(self, link_index: int, mass: jtp.FloatLike) -> KynDynParameter
         return self.replace(link_parameters=link_parameters)
 
     def set_link_inertia(
-        self, link_index: int, inertia: jtp.MatrixLike
+        self, link_index: jtp.IntLike, inertia: jtp.MatrixLike
     ) -> KynDynParameters:
         r"""
         Set the inertia tensor of a link.
@@ -593,10 +595,10 @@ def build_from_spatial_inertia(index: jtp.IntLike, M: jtp.Matrix) -> LinkParamet
         """
 
         # Extract the link parameters from the 6D spatial inertia.
-        m, L_p_CoM, I = Inertia.to_params(M=M)
+        m, L_p_CoM, I_CoM = Inertia.to_params(M=M)
 
         # Extract only the necessary elements of the inertia tensor.
-        inertia_elements = I[jnp.triu_indices(3)]
+        inertia_elements = I_CoM[jnp.triu_indices(3)]
 
         return LinkParameters(
             index=jnp.array(index).squeeze().astype(int),