Add CoM Jacobian method

src/adam/casadi/computations.py

@@ -143,7 +143,7 @@ def jacobian_dot_fun(self, frame: str) -> cs.Function:
         )
 
     def CoM_position_fun(self) -> cs.Function:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Returns:
             CoM (casADi function): The CoM position
@@ -155,6 +155,19 @@ def CoM_position_fun(self) -> cs.Function:
             "CoM_pos", [base_transform, joint_positions], [com_pos.array], self.f_opts
         )
 
+    def CoM_jacobian_fun(self) -> cs.Function:
+        """Returns the CoM Jacobian
+
+        Returns:
+            J_com (casADi function): The CoM Jacobian
+        """
+        joint_positions = cs.SX.sym("s", self.NDoF)
+        base_transform = cs.SX.sym("H", 4, 4)
+        J_com = self.rbdalgos.CoM_jacobian(base_transform, joint_positions)
+        return cs.Function(
+            "J_com", [base_transform, joint_positions], [J_com.array], self.f_opts
+        )
+
     def bias_force_fun(self) -> cs.Function:
         """Returns the bias force of the floating-base dynamics equation,
         using a reduced RNEA (no acceleration and external forces)
@@ -451,7 +464,7 @@ def gravity_term(self, base_transform: cs.SX, joint_positions: cs.SX) -> cs.SX:
         ).array
 
     def CoM_position(self, base_transform: cs.SX, joint_positions: cs.SX) -> cs.SX:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (cs.SX): The homogenous transform from base to world frame
@@ -466,3 +479,20 @@ def CoM_position(self, base_transform: cs.SX, joint_positions: cs.SX) -> cs.SX:
             )
 
         return self.rbdalgos.CoM_position(base_transform, joint_positions).array
+
+    def CoM_jacobian(self, base_transform: cs.SX, joint_positions: cs.SX) -> cs.SX:
+        """Returns the CoM Jacobian
+
+        Args:
+            base_transform (cs.SX): The homogenous transform from base to world frame
+            joint_positions (cs.SX): The joints position
+
+        Returns:
+            J_com (cs.SX): The CoM Jacobian
+        """
+        if isinstance(base_transform, cs.MX) and isinstance(joint_positions, cs.MX):
+            raise ValueError(
+                "You are using casadi MX. Please use the function KinDynComputations.CoM_jacobian_fun()"
+            )
+
+        return self.rbdalgos.CoM_jacobian(base_transform, joint_positions).array

src/adam/core/rbd_algorithms.py

@@ -327,7 +327,7 @@ def jacobian_dot(
     def CoM_position(
         self, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
     ) -> npt.ArrayLike:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (T): The homogenous transform from base to world frame
@@ -347,6 +347,39 @@ def CoM_position(
         com_pos /= self.get_total_mass()
         return com_pos
 
+    def CoM_jacobian(
+        self, base_transform: npt.ArrayLike, joint_positions: npt.ArrayLike
+    ) -> npt.ArrayLike:
+        """Returns the center of mass (CoM) Jacobian using the centroidal momentum matrix.
+
+        Args:
+            base_transform (T): The homogenous transform from base to world frame
+            joint_positions (T): The joints position
+
+        Returns:
+            J_com (T): The CoM Jacobian
+        """
+        # The com velocity can be computed as dot_x * m = J_cm_mixed * nu_mixed = J_cm_body * nu_body
+        # For this reason we compute the centroidal momentum matrix in mixed representation and then we convert it to body fixed if needed
+        # Save the original frame velocity representation
+        ori_frame_velocity_representation = self.frame_velocity_representation
+        # Set the frame velocity representation to mixed and compute the centroidal momentum matrix
+        self.frame_velocity_representation = Representations.MIXED_REPRESENTATION
+        _, Jcm = self.crba(base_transform, joint_positions)
+        if (
+            ori_frame_velocity_representation
+            == Representations.BODY_FIXED_REPRESENTATION
+        ):
+            # if the frame velocity representation is body fixed we need to convert the centroidal momentum matrix to body fixed
+            # dot_x * m = J_cm_mixed * mixed_to_body * nu_body
+            X = self.math.factory.eye(6 + self.NDoF)
+            X[:6, :6] = self.math.adjoint_mixed(base_transform)
+            Jcm = Jcm @ X
+        # Reset the frame velocity representation
+        self.frame_velocity_representation = ori_frame_velocity_representation
+        # Compute the CoM Jacobian
+        return Jcm[:3, :] / self.get_total_mass()
+
     def get_total_mass(self):
         """Returns the total mass of the robot
 

src/adam/jax/computations.py

@@ -216,7 +216,7 @@ def gravity_term(
     def CoM_position(
         self, base_transform: jnp.array, joint_positions: jnp.array
     ) -> jnp.array:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (jnp.array): The homogenous transform from base to world frame
@@ -229,6 +229,22 @@ def CoM_position(
             base_transform, joint_positions
         ).array.squeeze()
 
+    def CoM_jacobian(
+        self, base_transform: jnp.array, joint_positions: jnp.array
+    ) -> jnp.array:
+        """Returns the CoM Jacobian
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            joint_positions (jnp.array): The joints position
+
+        Returns:
+            Jcom (jnp.array): The CoM Jacobian
+        """
+        return self.rbdalgos.CoM_jacobian(
+            base_transform, joint_positions
+        ).array.squeeze()
+
     def get_total_mass(self) -> float:
         """Returns the total mass of the robot
 

src/adam/numpy/computations.py

@@ -149,7 +149,7 @@ def jacobian_dot(
     def CoM_position(
         self, base_transform: np.ndarray, joint_positions: np.ndarray
     ) -> np.ndarray:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (np.ndarray): The homogenous transform from base to world frame
@@ -162,6 +162,22 @@ def CoM_position(
             base_transform, joint_positions
         ).array.squeeze()
 
+    def CoM_jacobian(
+        self, base_transform: np.ndarray, joint_positions: np.ndarray
+    ) -> np.ndarray:
+        """Returns the CoM Jacobian
+
+        Args:
+            base_transform (np.ndarray): The homogenous transform from base to world frame
+            joint_positions (np.ndarray): The joints position
+
+        Returns:
+            Jcom (np.ndarray): The CoM Jacobian
+        """
+        return self.rbdalgos.CoM_jacobian(
+            base_transform, joint_positions
+        ).array.squeeze()
+
     def bias_force(
         self,
         base_transform: np.ndarray,

src/adam/parametric/casadi/computations_parametric.py

@@ -179,7 +179,7 @@ def jacobian_dot_fun(self, frame: str) -> cs.Function:
         )
 
     def CoM_position_fun(self) -> cs.Function:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Returns:
             com (casADi function): The CoM position

src/adam/parametric/jax/computations_parametric.py

@@ -397,7 +397,7 @@ def CoM_position(
         length_multiplier: jnp.array,
         densities: jnp.array,
     ) -> jnp.array:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (jnp.array): The homogenous transform from base to world frame

src/adam/parametric/numpy/computations_parametric.py

@@ -266,7 +266,7 @@ def CoM_position(
         length_multiplier: np.ndarray,
         densities: np.ndarray,
     ) -> np.ndarray:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (np.ndarray): The homogenous transform from base to world frame

src/adam/parametric/pytorch/computations_parametric.py

@@ -271,7 +271,7 @@ def CoM_position(
         length_multiplier: torch.Tensor,
         densities: torch.Tensor,
     ) -> torch.Tensor:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (torch.tensor): The homogenous transform from base to world frame

src/adam/pytorch/computation_batch.py

@@ -406,7 +406,7 @@ def fun(base_transform, joint_positions):
     def CoM_position(
         self, base_transform: torch.Tensor, joint_positions: torch.Tensor
     ) -> torch.Tensor:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (torch.Tensor): The homogenous transform from base to world frame
@@ -418,7 +418,7 @@ def CoM_position(
         return self.CoM_position_fun()(base_transform, joint_positions)
 
     def CoM_position_fun(self):
-        """Returns the CoM positon as a pytorch function
+        """Returns the CoM position as a pytorch function
 
         Returns:
             CoM (pytorch function): The CoM position
@@ -435,6 +435,38 @@ def fun(base_transform, joint_positions):
         self.funcs["CoM_position"] = jax2torch(jit_vmapped_fun)
         return self.funcs["CoM_position"]
 
+    def CoM_jacobian(
+        self, base_transform: torch.Tensor, joint_positions: torch.Tensor
+    ) -> torch.Tensor:
+        """Returns the CoM Jacobian
+
+        Args:
+            base_transform (torch.Tensor): The homogenous transform from base to world frame
+            joint_positions (torch.Tensor): The joints position
+
+        Returns:
+            Jcom (torch.Tensor): The CoM Jacobian
+        """
+        return self.CoM_jacobian_fun()(base_transform, joint_positions)
+
+    def CoM_jacobian_fun(self):
+        """Returns the CoM Jacobian as a pytorch function
+
+        Returns:
+            Jcom (pytorch function): The CoM Jacobian
+        """
+        if self.funcs.get("CoM_jacobian") is not None:
+            return self.funcs["CoM_jacobian"]
+        print("[INFO] Compiling CoM Jacobian function")
+
+        def fun(base_transform, joint_positions):
+            return self.rbdalgos.CoM_jacobian(base_transform, joint_positions).array
+
+        vmapped_fun = jax.vmap(fun, in_axes=(0, 0))
+        jit_vmapped_fun = jax.jit(vmapped_fun)
+        self.funcs["CoM_jacobian"] = jax2torch(jit_vmapped_fun)
+        return self.funcs["CoM_jacobian"]
+
     def get_total_mass(self) -> float:
         """Returns the total mass of the robot
 

src/adam/pytorch/computations.py

@@ -157,7 +157,7 @@ def jacobian_dot(
     def CoM_position(
         self, base_transform: torch.Tensor, joint_positions: torch.Tensor
     ) -> torch.Tensor:
-        """Returns the CoM positon
+        """Returns the CoM position
 
         Args:
             base_transform (torch.tensor): The homogenous transform from base to world frame
@@ -170,6 +170,20 @@ def CoM_position(
             base_transform, joint_positions
         ).array.squeeze()
 
+    def CoM_jacobian(
+        self, base_transform: torch.Tensor, joint_positions: torch.Tensor
+    ) -> torch.Tensor:
+        """Returns the CoM Jacobian
+
+        Args:
+            base_transform (torch.tensor): The homogenous transform from base to world frame
+            joint_positions (torch.tensor): The joints position
+
+        Returns:
+            Jcom (torch.tensor): The CoM Jacobian
+        """
+        return self.rbdalgos.CoM_jacobian(base_transform, joint_positions).array
+
     def bias_force(
         self,
         base_transform: torch.Tensor,

tests/conftest.py

@@ -27,6 +27,7 @@ class IDynFunctionValues:
     mass_matrix: np.ndarray
     centroidal_momentum_matrix: np.ndarray
     CoM_position: np.ndarray
+    CoM_jacobian: np.ndarray
     total_mass: float
     jacobian: np.ndarray
     jacobian_non_actuated: np.ndarray
@@ -184,6 +185,11 @@ def compute_idyntree_values(
     # CoM position
     idyn_com = kin_dyn.getCenterOfMassPosition().toNumPy()
 
+    # Com jacobian
+    idyn_com_jacobian = idyntree.MatrixDynSize(3, kin_dyn.model().getNrOfDOFs() + 6)
+    kin_dyn.getCenterOfMassJacobian(idyn_com_jacobian)
+    idyn_com_jacobian = idyn_com_jacobian.toNumPy()
+
     # total mass
     total_mass = kin_dyn.model().getTotalMass()
 
@@ -277,6 +283,7 @@ def compute_idyntree_values(
         mass_matrix=idyn_mass_matrix,
         centroidal_momentum_matrix=idyn_cmm,
         CoM_position=idyn_com,
+        CoM_jacobian=idyn_com_jacobian,
         total_mass=total_mass,
         jacobian=idyn_jacobian,
         jacobian_non_actuated=idyn_jacobian_non_actuated,

tests/test_casadi.py

@@ -43,6 +43,17 @@ def test_CoM_pos(setup_test):
     assert adam_com - idyn_com == pytest.approx(0.0, abs=1e-5)
 
 
+def test_CoM_jacobian(setup_test):
+    adam_kin_dyn, robot_cfg, state = setup_test
+    idyn_com_jacobian = robot_cfg.idyn_function_values.CoM_jacobian
+    adam_com_jacobian = cs.DM(
+        adam_kin_dyn.CoM_jacobian_fun()(state.H, state.joints_pos)
+    )
+    assert adam_com_jacobian - idyn_com_jacobian == pytest.approx(0.0, abs=1e-5)
+    adam_com_jacobian = cs.DM(adam_kin_dyn.CoM_jacobian(state.H, state.joints_pos))
+    assert adam_com_jacobian - idyn_com_jacobian == pytest.approx(0.0, abs=1e-5)
+
+
 def test_total_mass(setup_test):
     adam_kin_dyn, robot_cfg, state = setup_test
     idyn_total_mass = robot_cfg.idyn_function_values.total_mass

tests/test_idyntree_conversion.py

@@ -46,6 +46,13 @@ def test_CoM_pos(setup_test):
     adam_com = cs.DM(adam_kin_dyn.CoM_position_fun()(state.H, state.joints_pos))
     assert adam_com - idyn_com == pytest.approx(0.0, abs=1e-5)
 
+def test_CoM_jacobian(setup_test):
+    adam_kin_dyn, robot_cfg, state = setup_test
+    idyn_com_jac = robot_cfg.idyn_function_values.CoM_jacobian
+    adam_com_jac = cs.DM(adam_kin_dyn.CoM_jacobian_fun()(state.H, state.joints_pos))
+    assert adam_com_jac - idyn_com_jac == pytest.approx(0.0, abs=1e-5)
+    adam_com_jac = cs.DM(adam_kin_dyn.CoM_jacobian_fun()(state.H, state.joints_pos))
+    assert adam_com_jac - idyn_com_jac == pytest.approx(0.0, abs=1e-5)
 
 def test_total_mass(setup_test):
     adam_kin_dyn, robot_cfg, state = setup_test

tests/test_jax.py

@@ -37,6 +37,13 @@ def test_CoM_pos(setup_test):
     assert adam_com - idyn_com == pytest.approx(0.0, abs=1e-5)
 
 
+def test_CoM_jacobian(setup_test):
+    adam_kin_dyn, robot_cfg, state = setup_test
+    idyn_com_jacobian = robot_cfg.idyn_function_values.CoM_jacobian
+    adam_com_jacobian = adam_kin_dyn.CoM_jacobian(state.H, state.joints_pos)
+    assert adam_com_jacobian - idyn_com_jacobian == pytest.approx(0.0, abs=1e-5)
+
+
 def test_total_mass(setup_test):
     adam_kin_dyn, robot_cfg, state = setup_test
     idyn_total_mass = robot_cfg.idyn_function_values.total_mass

tests/test_numpy.py

@@ -34,6 +34,13 @@ def test_CoM_pos(setup_test):
     assert adam_com - idyn_com == pytest.approx(0.0, abs=1e-5)
 
 
+def test_CoM_jacobian(setup_test):
+    adam_kin_dyn, robot_cfg, state = setup_test
+    idyn_com_jacobian = robot_cfg.idyn_function_values.CoM_jacobian
+    adam_com_jacobian = adam_kin_dyn.CoM_jacobian(state.H, state.joints_pos)
+    assert adam_com_jacobian - idyn_com_jacobian == pytest.approx(0.0, abs=1e-5)
+
+
 def test_total_mass(setup_test):
     adam_kin_dyn, robot_cfg, state = setup_test
     idyn_total_mass = robot_cfg.idyn_function_values.total_mass

tests/test_pytorch.py

@@ -44,6 +44,13 @@ def test_CoM_pos(setup_test):
     assert adam_com.numpy() - idyn_com == pytest.approx(0.0, abs=1e-5)
 
 
+def test_CoM_jacobian(setup_test):
+    adam_kin_dyn, robot_cfg, state = setup_test
+    idyn_com_jacobian = robot_cfg.idyn_function_values.CoM_jacobian
+    adam_com_jacobian = adam_kin_dyn.CoM_jacobian(state.H, state.joints_pos)
+    assert adam_com_jacobian.numpy() - idyn_com_jacobian == pytest.approx(0.0, abs=1e-5)
+
+
 def test_total_mass(setup_test):
     adam_kin_dyn, robot_cfg, state = setup_test
     idyn_total_mass = robot_cfg.idyn_function_values.total_mass