Merge pull request #117 from ami-iit/add_com_jacobian_method

Add CoM Jacobian method

src/adam/casadi/computations.py

@@ -147,7 +147,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
@@ -159,6 +159,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)
@@ -455,7 +468,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
@@ -470,3 +483,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

@@ -325,7 +325,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
@@ -345,6 +345,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

@@ -220,7 +220,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
@@ -233,6 +233,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

@@ -152,7 +152,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
@@ -165,6 +165,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

@@ -183,7 +183,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

@@ -401,7 +401,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

@@ -270,7 +270,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

@@ -275,7 +275,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

@@ -414,7 +414,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
@@ -426,7 +426,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
@@ -443,6 +443,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

@@ -161,7 +161,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
@@ -174,6 +174,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

@@ -47,6 +47,15 @@ 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_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
     idyn_total_mass = robot_cfg.idyn_function_values.total_mass

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