diff --git a/src/adam/casadi/computations.py b/src/adam/casadi/computations.py index ff67fda2..1d9191c5 100644 --- a/src/adam/casadi/computations.py +++ b/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 diff --git a/src/adam/core/rbd_algorithms.py b/src/adam/core/rbd_algorithms.py index 4c67d980..a10e8105 100644 --- a/src/adam/core/rbd_algorithms.py +++ b/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 diff --git a/src/adam/jax/computations.py b/src/adam/jax/computations.py index 82d3611f..e2d66185 100644 --- a/src/adam/jax/computations.py +++ b/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 diff --git a/src/adam/numpy/computations.py b/src/adam/numpy/computations.py index 9ddd2d26..6dccdca0 100644 --- a/src/adam/numpy/computations.py +++ b/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, diff --git a/src/adam/parametric/casadi/computations_parametric.py b/src/adam/parametric/casadi/computations_parametric.py index 4149662d..76ed45a8 100644 --- a/src/adam/parametric/casadi/computations_parametric.py +++ b/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 diff --git a/src/adam/parametric/jax/computations_parametric.py b/src/adam/parametric/jax/computations_parametric.py index 7920006a..5104029a 100644 --- a/src/adam/parametric/jax/computations_parametric.py +++ b/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 diff --git a/src/adam/parametric/numpy/computations_parametric.py b/src/adam/parametric/numpy/computations_parametric.py index f7030376..a1cd3d98 100644 --- a/src/adam/parametric/numpy/computations_parametric.py +++ b/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 diff --git a/src/adam/parametric/pytorch/computations_parametric.py b/src/adam/parametric/pytorch/computations_parametric.py index 63568d45..33e3c6e0 100644 --- a/src/adam/parametric/pytorch/computations_parametric.py +++ b/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 diff --git a/src/adam/pytorch/computation_batch.py b/src/adam/pytorch/computation_batch.py index b5e96a10..6964e6f5 100644 --- a/src/adam/pytorch/computation_batch.py +++ b/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 diff --git a/src/adam/pytorch/computations.py b/src/adam/pytorch/computations.py index 3fc70805..71981046 100644 --- a/src/adam/pytorch/computations.py +++ b/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, diff --git a/tests/conftest.py b/tests/conftest.py index 31c582bf..a49b765a 100644 --- a/tests/conftest.py +++ b/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, diff --git a/tests/test_casadi.py b/tests/test_casadi.py index 4f213a4d..f136bfa8 100644 --- a/tests/test_casadi.py +++ b/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 diff --git a/tests/test_idyntree_conversion.py b/tests/test_idyntree_conversion.py index ceb5f779..063126a0 100644 --- a/tests/test_idyntree_conversion.py +++ b/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 diff --git a/tests/test_jax.py b/tests/test_jax.py index a74449e6..7a35681c 100644 --- a/tests/test_jax.py +++ b/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 diff --git a/tests/test_numpy.py b/tests/test_numpy.py index ba3e608c..2b2aa169 100644 --- a/tests/test_numpy.py +++ b/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 diff --git a/tests/test_pytorch.py b/tests/test_pytorch.py index 2af3bbc4..4976cd0b 100644 --- a/tests/test_pytorch.py +++ b/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 diff --git a/tests/test_pytorch_batch.py b/tests/test_pytorch_batch.py index 4839b207..e8509967 100644 --- a/tests/test_pytorch_batch.py +++ b/tests/test_pytorch_batch.py @@ -73,6 +73,20 @@ def test_CoM_pos(setup_test): assert adam_com.shape == (n_samples, 3) +def test_CoM_jacobian(setup_test): + adam_kin_dyn, robot_cfg, state, n_samples = 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) + try: + adam_com_jacobian.sum().backward() + except: + raise ValueError(adam_com_jacobian) + assert adam_com_jacobian[0].detach().numpy() - idyn_com_jacobian == pytest.approx( + 0.0, abs=1e-4 + ) + assert adam_com_jacobian.shape == (n_samples, 3, robot_cfg.n_dof + 6) + + def test_jacobian(setup_test): adam_kin_dyn, robot_cfg, state, n_samples = setup_test idyn_jacobian = robot_cfg.idyn_function_values.jacobian