diff --git a/README.md b/README.md index 7ce04532..df556232 100644 --- a/README.md +++ b/README.md @@ -111,29 +111,28 @@ pip install .[selected-interface] #### Installation from conda-forge package - - Install **CasADi** interface: ```bash - mamba create -n adamenv -c conda-forge adam-robotics-casadi + conda create -n adamenv -c conda-forge adam-robotics-casadi ``` - + - Install **Jax** interface (warning: not available on Windows): ```bash - mamba create -n adamenv -c conda-forge adam-robotics-jax + conda create -n adamenv -c conda-forge adam-robotics-jax ``` - Install **PyTorch** interface (warning: not available on Windows): ```bash - mamba create -n adamenv -c conda-forge adam-robotics-pytorch + conda create -n adamenv -c conda-forge adam-robotics-pytorch ``` - Install **ALL** interfaces (warning: not available on Windows): ```bash - mamba create -n adamenv -c conda-forge adam-robotics-all + conda create -n adamenv -c conda-forge adam-robotics-all ``` > [!NOTE] @@ -146,31 +145,31 @@ Install in a conda environment the required dependencies: - **Jax** interface dependencies: ```bash - mamba create -n adamenv -c conda-forge jax numpy lxml prettytable matplotlib urdfdom-py + conda create -n adamenv -c conda-forge jax numpy lxml prettytable matplotlib urdfdom-py ``` - **CasADi** interface dependencies: ```bash - mamba create -n adamenv -c conda-forge casadi numpy lxml prettytable matplotlib urdfdom-py + conda create -n adamenv -c conda-forge casadi numpy lxml prettytable matplotlib urdfdom-py ``` - **PyTorch** interface dependencies: ```bash - mamba create -n adamenv -c conda-forge pytorch numpy lxml prettytable matplotlib urdfdom-py jax2torch + conda create -n adamenv -c conda-forge pytorch numpy lxml prettytable matplotlib urdfdom-py jax2torch ``` - **ALL** interfaces dependencies: ```bash - mamba create -n adamenv -c conda-forge jax casadi pytorch numpy lxml prettytable matplotlib urdfdom-py jax2torch + conda create -n adamenv -c conda-forge jax casadi pytorch numpy lxml prettytable matplotlib urdfdom-py jax2torch ``` Activate the environment, clone the repo and install the library: ```bash -mamba activate adamenv +conda activate adamenv git clone https://github.com/ami-iit/ADAM.git cd adam pip install --no-deps . @@ -352,6 +351,46 @@ M = kinDyn.mass_matrix(w_H_b_batch, joints_batch) w_H_f = kinDyn.forward_kinematics('frame_name', w_H_b_batch, joints_batch) ``` +### Inverse Kinematics + +adam provides an interface for solving inverse kinematics problems using CasADi. The solver supports + +- position, orientation, and full pose constraints +- frame-to-frame constraints (ball, fixed) +- optional joint limit constraints + +```python +import casadi as cs +import numpy as np +import adam +from adam.casadi import KinDynComputations +from adam.casadi.inverse_kinematics import InverseKinematics, TargetType + +# Load your robot model +import icub_models +model_path = icub_models.get_model_file("iCubGazeboV2_5") +# The joint list +joints_name_list = ... +# Create IK solver +ik = InverseKinematics(model_path, joints) +# Add a pose target on a frame (e.g., the left sole) +ik.add_target("l_sole", target_type=TargetType.POSE, as_soft_constraint=True, weight=1.0) +ik.add_ball_constraint(frame_1, frame_2, as_soft_constraint=True) + +# Update the target to a desired pose +desired_position = np.array([0.3, 0.2, 1.0]) +desired_orientation = np.eye(3) +ik.update_target("l_sole", (desired_position, desired_orientation)) + +# Solve +ik.solve() + +# Retrieve solution +w_H_b_sol, q_sol = ik.get_solution() +print("Base pose:\n", w_H_b_sol) +print("Joint values:\n", q_sol) +``` + ## 🦸‍♂️ Contributing **adam** is an open-source project. Contributions are very welcome! diff --git a/ci_env.yml b/ci_env.yml index 9073a8e6..c07afa39 100644 --- a/ci_env.yml +++ b/ci_env.yml @@ -20,3 +20,4 @@ dependencies: - pytorch - jax2torch - requests + - liecasadi diff --git a/ci_env_win.yml b/ci_env_win.yml index 58370404..7f3281f4 100644 --- a/ci_env_win.yml +++ b/ci_env_win.yml @@ -17,3 +17,4 @@ dependencies: - icub-models - idyntree >=11.0.0 - requests + - liecasadi diff --git a/pyproject.toml b/pyproject.toml index a1f20ff9..77083b0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ test = [ "black", "jax2torch", "requests", + "liecasadi", ] conversion = ["idyntree"] all = ["jax", "jaxlib", "casadi>=3.6", "torch", "jax2torch"] @@ -57,7 +58,7 @@ all = ["jax", "jaxlib", "casadi>=3.6", "torch", "jax2torch"] "Source" = "https://github.com/ami-iit/adam" [tool.setuptools] -package-dir = {"" = "src"} # keeps the src/ layout +package-dir = { "" = "src" } # keeps the src/ layout [tool.setuptools.packages.find] where = ["src"] diff --git a/src/adam/casadi/__init__.py b/src/adam/casadi/__init__.py index 583c07a7..03e467ff 100644 --- a/src/adam/casadi/__init__.py +++ b/src/adam/casadi/__init__.py @@ -2,3 +2,4 @@ from .casadi_like import CasadiLike from .computations import KinDynComputations +from .inverse_kinematics import InverseKinematics, TargetType diff --git a/src/adam/casadi/inverse_kinematics.py b/src/adam/casadi/inverse_kinematics.py new file mode 100644 index 00000000..1cdba473 --- /dev/null +++ b/src/adam/casadi/inverse_kinematics.py @@ -0,0 +1,558 @@ +from __future__ import annotations +import dataclasses +from enum import Enum, auto +from typing import Any +import casadi as cs +import numpy as np +from adam.casadi import KinDynComputations +from liecasadi import SO3, SE3 + + +class TargetType(Enum): + """Type of IK target supported by the solver.""" + + POSITION = auto() + ROTATION = auto() + POSE = auto() + + +class FramesConstraint(Enum): + """Type of constraint for the frames in the IK problem.""" + + BALL = auto() # Ball constraint (position only) + FIXED = auto() # Fixed constraint (position + rotation) + + +@dataclasses.dataclass +class Target: + """Dataclass to store target information.""" + + target_type: TargetType + frame: str + weight: float = 1.0 + as_soft_constraint: bool = True + forward_kinematics_function: cs.Function = None + param_pos: cs.SX = None + param_rot: cs.SX = None + + +class InverseKinematics: + def __init__( + self, + urdf_path: str, + joints_list: list[str], + joint_limits_active: bool = True, + solver_settings: dict[str, Any] = None, + ): + """Initialize the InverseKinematics solver. + + Args: + urdf_path (str): Path to the URDF file. + joints_list (list[str]): List of joint names. + joint_limits_active (bool, optional): If True, enforces joint limits. Defaults to True. + solver_settings (dict[str, ], optional): Settings for the solver. Defaults to None. + """ + self.kd = KinDynComputations(urdf_path, joints_list) + self.ndof = len(joints_list) + self.joints_list = joints_list + + # Opti variables -------------------------------------------------- + self.opti = cs.Opti() + self.base_pos = self.opti.variable(3) # translation (x,y,z) + self.base_quat = self.opti.variable(4) # unit quaternion (x,y,z,w) + self.joint_var = self.opti.variable(self.ndof) + self.opti.set_initial(self.base_pos, [0, 0, 0]) # default to origin + self.opti.set_initial(self.base_quat, [0, 0, 0, 1]) # identity quaternion + self.opti.subject_to(cs.sumsqr(self.base_quat) == 1) # enforce unit quaternion + self.opti.set_initial(self.joint_var, np.zeros(self.ndof)) # default to zero + + self.targets: dict[str, Target] = {} + self.cost_terms: list[cs.MX] = [] + self._problem_built = False + self._cached_sol = None + solver_settings = ( + { + "ipopt": {"print_level": 0}, + } + if solver_settings is None + else solver_settings + ) + self.joint_limits_active = joint_limits_active + self.set_default_joint_limit_constraints() + + self.opti.solver("ipopt", solver_settings) + + def set_solver(self, solver_name: str, solver_settings: dict[str, Any]): + """Set the solver for the optimization problem. + + Args: + solver_name (str): The name of the solver to use. + solver_settings (dict[str, Any]): The settings for the solver. + """ + self.opti.solver(solver_name, solver_settings) + self._cached_sol = None + + def get_opti(self) -> cs.Opti: + """Get the CasADi Opti object for this IK solver. This can be used to add custom constraints or objectives. + + Returns: + cs.Opti: The CasADi Opti object. + """ + return self.opti + + def get_kin_dyn_computations(self) -> KinDynComputations: + """Get the KinDynComputations object used by this IK solver. + + Returns: + KinDynComputations: The KinDynComputations object. + """ + return self.kd + + def set_default_joint_limit_constraints(self): + """Set joint limit constraints if active.""" + if self.joint_limits_active: + # Add joint limits as constraints + for i, joint_name in enumerate(self.joints_list): + joint_limit = self.kd.rbdalgos.model.joints[joint_name].limit + if joint_limit is not None: + self.opti.subject_to( + self.opti.bounded( + joint_limit.lower, self.joint_var[i], joint_limit.upper + ) + ) + + def set_joint_limits(self, lower_bounds: np.ndarray, upper_bounds: np.ndarray): + """Set custom joint limits for the optimization problem. + + Args: + lower_bounds (np.ndarray): Lower bounds for each joint. + upper_bounds (np.ndarray): Upper bounds for each joint. + """ + self._ensure_graph_modifiable() + if len(lower_bounds) != self.ndof or len(upper_bounds) != self.ndof: + raise ValueError("Bounds must match the number of degrees of freedom") + for i in range(self.ndof): + self.opti.subject_to( + self.opti.bounded(lower_bounds[i], self.joint_var[i], upper_bounds[i]) + ) + + def add_target_position( + self, + frame: str, + *, + as_soft_constraint: bool = True, + weight: float = 1.0, + ): + """Add a target position for the IK solver. + + Args: + frame (str): The name of the frame to target. + as_soft_constraint (bool): If False, adds the target as a hard constraint. If True, adds it as a soft constraint (cost term). + weight (float): The weight of the target. Defaults to 1.0. + + Raises: + ValueError: If the target frame already exists. + """ + self._ensure_graph_modifiable() + if frame in self.targets: + raise ValueError(f"Target '{frame}' already exists") + + p_des = self.opti.parameter(3) + self.opti.set_value(p_des, np.zeros(3)) # default to origin + p_fk = self.kd.forward_kinematics_fun(frame)( + self.base_transform(), self.joint_var + )[:3, 3] + + if as_soft_constraint: + self.cost_terms.append(weight * cs.sumsqr(p_fk - p_des)) + else: + self.opti.subject_to(p_fk == p_des) + + self.targets[frame] = Target( + target_type=TargetType.POSITION, + frame=frame, + weight=weight, + as_soft_constraint=as_soft_constraint, + param_pos=p_des, + forward_kinematics_function=self.kd.forward_kinematics_fun(frame), + ) + + def add_frames_constraint( + self, + parent_frame: str, + child_frame: str, + constraint_type: FramesConstraint, + as_soft_constraint: bool = True, + weight: float = 1e5, + ): + """Add a constraint between two frames. + + Args: + parent_frame (str): The name of the parent frame. + child_frame (str): The name of the child frame. + constraint_type (FramesConstraint): Type of constraint to apply. + as_soft_constraint (bool): If False, adds the constraint as a hard constraint. If True, adds it as a soft constraint (cost term). + weight (float): Weight for the constraint in the cost function. + """ + self._ensure_graph_modifiable() + # check that frames are different + if parent_frame == child_frame: + raise ValueError("Parent and child frames must be different") + H_parent = self.kd.forward_kinematics_fun(parent_frame)( + self.base_transform(), self.joint_var + ) + H_child = self.kd.forward_kinematics_fun(child_frame)( + self.base_transform(), self.joint_var + ) + p_parent, p_child = H_parent[:3, 3], H_child[:3, 3] + R_parent, R_child = H_parent[:3, :3], H_child[:3, :3] + if constraint_type is FramesConstraint.BALL: + # Ball constraint: child frame must be positioned relative to parent frame + if as_soft_constraint: + slack = self.opti.variable(1) + self.opti.subject_to(cs.sumsqr(p_child - p_parent) <= slack) + self.opti.subject_to( + self.opti.bounded(0, slack, 1e-3) + ) # small slack to allow for numerical stability + self.cost_terms.append(cs.sumsqr(slack) * weight) + else: + self.opti.subject_to(p_child == p_parent) + + elif constraint_type is FramesConstraint.FIXED: + # Fixed constraint: child frame must be aligned and positioned with parent frame + if as_soft_constraint: + slack = self.opti.variable(1) + rot_err_sq = cs.power(1 - (cs.trace(R_parent.T @ R_child) - 1) / 2, 2) + self.opti.subject_to(rot_err_sq <= slack) + self.opti.subject_to(self.opti.bounded(0, slack, 1e-3)) + self.cost_terms.append(cs.sumsqr(slack) * weight) + else: + self.opti.subject_to(p_child == p_parent) + self.opti.subject_to(R_child == R_parent) + else: + raise ValueError(f"Unsupported constraint type: {constraint_type.name}") + + def add_ball_constraint( + self, + parent_frame: str, + child_frame: str, + as_soft_constraint: bool = True, + weight: float = 1e5, + ): + """Add a ball constraint between two frames. + + Args: + parent_frame (str): The name of the parent frame. + child_frame (str): The name of the child frame. + as_soft_constraint (bool): If False, adds the constraint as a hard constraint. If True, adds it as a soft constraint (cost term). + weight (float): Weight for the constraint in the cost function, ignored if `as_soft_constraint` is `False` + """ + self.add_frames_constraint( + parent_frame, child_frame, FramesConstraint.BALL, as_soft_constraint, weight + ) + + def add_fixed_constraint( + self, + parent_frame: str, + child_frame: str, + as_soft_constraint: bool = True, + weight: float = 1e5, + ): + """Add a fixed constraint between two frames. + + Args: + parent_frame (str): The name of the parent frame. + child_frame (str): The name of the child frame. + as_soft_constraint (bool): If False, adds the constraint as a hard constraint. If True, adds it as a soft constraint (cost term). + weight (float): Weight for the constraint in the cost function, ignored if `as_soft_constraint` is `False` + """ + self.add_frames_constraint( + parent_frame, + child_frame, + FramesConstraint.FIXED, + as_soft_constraint, + weight, + ) + + def add_min_distance_constraint(self, frames_list: list[str], distance: float): + """Add a distance constraint between frames. + + Args: + frames_list (list[str]): List of frame names to apply the distance constraint. + distance (float): Minimum distance between consecutive frames. + """ + self._ensure_graph_modifiable() + if len(frames_list) < 2: + raise ValueError("At least two frames are required for distance constraint") + for i in range(len(frames_list) - 1): + print( + f"Adding distance constraint between {frames_list[i]} and {frames_list[i + 1]}" + ) + frame_i, frame_j = frames_list[i], frames_list[i + 1] + p_i = self.kd.forward_kinematics_fun(frame_i)( + self.base_transform(), self.joint_var + )[:3, 3] + p_j = self.kd.forward_kinematics_fun(frame_j)( + self.base_transform(), self.joint_var + )[:3, 3] + dist_sq = cs.sumsqr(p_i - p_j) + self.opti.subject_to(dist_sq >= distance**2) + + def add_target_orientation( + self, frame: str, *, as_soft_constraint: bool = True, weight: float = 1.0 + ): + """Add an orientation target for a frame. + + Args: + frame (str): The name of the frame to target. + as_soft_constraint (bool): If False, adds the target as a hard constraint. If True, adds it as a soft constraint (cost term). + weight (float): Weight for the target in the cost function. + + Raises: + ValueError: If the target frame already exists. + """ + self._ensure_graph_modifiable() + if frame in self.targets: + raise ValueError(f"Target '{frame}' already exists") + + R_des = self.opti.parameter(3, 3) + self.opti.set_value(R_des, np.eye(3)) # default to identity rotation + H_fk = self.kd.forward_kinematics_fun(frame)( + self.base_transform(), self.joint_var + ) + R_fk = H_fk[:3, :3] + # proxy for rotation error: trace(R) = 1 + 2 * cos(theta), where theta is the angle of rotation + rot_err_sq = cs.power(1 - (cs.trace(R_des.T @ R_fk) - 1) / 2, 2) + + if as_soft_constraint: + self.cost_terms.append(weight * rot_err_sq) + else: + self.opti.subject_to(rot_err_sq == 0) + + self.targets[frame] = Target( + target_type=TargetType.ROTATION, + frame=frame, + weight=weight, + as_soft_constraint=as_soft_constraint, + param_rot=R_des, + forward_kinematics_function=self.kd.forward_kinematics_fun(frame), + ) + + def add_target_pose( + self, frame: str, *, as_soft_constraint: bool = True, weight: float = 1.0 + ): + """Add a pose target for a frame. + + Args: + frame (str): The name of the frame to target. + as_soft_constraint (bool): If False, adds the target as a hard constraint. If True, adds it as a soft constraint (cost term). + weight (float): Weight for the target in the cost function. + Raises: + ValueError: If the target frame already exists. + """ + self._ensure_graph_modifiable() + if frame in self.targets: + raise ValueError(f"Target '{frame}' already exists") + + p_des = self.opti.parameter(3) + R_des = self.opti.parameter(3, 3) + self.opti.set_value(p_des, np.zeros(3)) # default to origin + self.opti.set_value(R_des, np.eye(3)) # default to identity rotation + H_fk = self.kd.forward_kinematics_fun(frame)( + self.base_transform(), self.joint_var + ) + p_fk, R_fk = H_fk[:3, 3], H_fk[:3, :3] + + pos_err_sq = cs.sumsqr(p_fk - p_des) + rot_err_sq = cs.power(1 - (cs.trace(R_des.T @ R_fk) - 1) / 2, 2) + + if as_soft_constraint: + self.cost_terms.append(weight * (pos_err_sq + rot_err_sq)) + else: + self.opti.subject_to(p_fk == p_des) + self.opti.subject_to(rot_err_sq == 0) + + self.targets[frame] = Target( + target_type=TargetType.POSE, + frame=frame, + weight=weight, + as_soft_constraint=as_soft_constraint, + param_pos=p_des, + param_rot=R_des, + forward_kinematics_function=self.kd.forward_kinematics_fun(frame), + ) + + def add_target( + self, + frame: str, + target_type: TargetType, + *, + as_soft_constraint: bool = True, + weight: float = 1.0, + ): + """Add a target for a frame. + + Args: + frame (str): The name of the frame to target. + target_type (TargetType): The type of target (position, rotation, or pose). + as_soft_constraint (bool): If False, adds the target as a hard constraint. If True, adds it as a soft constraint (cost term). + weight (float): Weight for the target in the cost function. + """ + if target_type is TargetType.POSITION: + self.add_target_position( + frame, weight=weight, as_soft_constraint=as_soft_constraint + ) + elif target_type is TargetType.ROTATION: + self.add_target_orientation( + frame, weight=weight, as_soft_constraint=as_soft_constraint + ) + elif target_type is TargetType.POSE: + self.add_target_pose( + frame, weight=weight, as_soft_constraint=as_soft_constraint + ) + else: + raise ValueError("Unsupported target type") + + def update_target_position(self, frame: str, position: np.ndarray): + """Update the target position for a frame. + + Args: + frame (str): The name of the frame to update. + position (np.ndarray): The new target position. + """ + self._check_target_type(frame, TargetType.POSITION) + self.opti.set_value(self.targets[frame].param_pos, position) + + def update_target_orientation(self, frame: str, rotation: np.ndarray): + """Update the target orientation for a frame. + + Args: + frame (str): The name of the frame to update. + rotation (np.ndarray): The new target rotation. + """ + self._check_target_type(frame, TargetType.ROTATION) + self.opti.set_value(self.targets[frame].param_rot, rotation) + + def update_target_pose( + self, frame: str, position: np.ndarray, rotation: np.ndarray + ): + """Update the target pose for a frame. + + Args: + frame (str): The name of the frame to update. + position (np.ndarray): The new target position. + rotation (np.ndarray): The new target rotation. + """ + self._check_target_type(frame, TargetType.POSE) + self.opti.set_value(self.targets[frame].param_pos, position) + self.opti.set_value(self.targets[frame].param_rot, rotation) + + def update_target( + self, frame: str, target: np.ndarray | tuple[np.ndarray, np.ndarray] + ): + """Update the target for a frame. + + Args: + frame (str): The name of the frame to update. + target (Union[np.ndarray, tuple[np.ndarray, np.ndarray]]): The new target position or pose. + + Raises: + ValueError: If the target is of an invalid type. + ValueError: If the frame is not a pose target. + RuntimeError: If the optimization problem has already been built. + """ + if frame not in self.targets: + raise ValueError(f"Unknown target '{frame}'") + ttype = self.targets[frame].target_type + if ttype is TargetType.POSITION: + self.update_target_position(frame, target) # type: ignore[arg-type] + elif ttype is TargetType.ROTATION: + self.update_target_orientation(frame, target) # type: ignore[arg-type] + elif ttype is TargetType.POSE: + if not (isinstance(target, (list, tuple)) and len(target) == 2): + raise ValueError("Pose update expects (position, rotation)") + self.update_target_pose(frame, *target) # type: ignore[arg-type] + else: + raise RuntimeError("Unknown target type") + + def set_initial_guess(self, base_transform: np.ndarray, joint_values: np.ndarray): + """Set the initial guess for the optimization problem. + + Args: + base_transform (np.ndarray): 4x4 transformation matrix for the floating base. + joint_values (np.ndarray): Initial joint values for the robot. + """ + pos, rot = base_transform[:3, 3], base_transform[:3, :3] + quat = SO3.from_matrix(rot).as_quat().coeffs() # (x,y,z,w) + self.opti.set_initial(self.base_pos, pos) + self.opti.set_initial(self.base_quat, quat) + self.opti.set_initial(self.joint_var, joint_values) + + def solve(self): + """Solve the optimization problem.""" + if not self._problem_built: + self._finalize_problem() + if self._cached_sol is not None: + # Warm‑start + self.opti.set_initial(self.base_pos, self._cached_sol.value(self.base_pos)) + self.opti.set_initial( + self.base_quat, self._cached_sol.value(self.base_quat) + ) + self.opti.set_initial( + self.joint_var, self._cached_sol.value(self.joint_var) + ) + self._cached_sol = self.opti.solve() + + def get_solution(self): + """Get the solution of the optimization problem. + + Raises: + RuntimeError: If the optimization problem has not been solved yet. + + Returns: + Tuple[np.ndarray, np.ndarray]: The 4x4 transformation matrix and joint values. + """ + if self._cached_sol is None: + raise RuntimeError("No solution yet - call solve() first") + pos = self._cached_sol.value(self.base_pos) + quat = self._cached_sol.value(self.base_quat) + joints = self._cached_sol.value(self.joint_var) + H = SE3.from_position_quaternion(pos, quat).as_matrix() # 4x4 numpy + return np.array(H), np.array(joints) + + def base_transform(self): + """Symbolic 4x4 transform of the floating base (SX).""" + return SE3.from_position_quaternion(self.base_pos, self.base_quat).as_matrix() + + def _ensure_graph_modifiable(self): + if self._problem_built: + raise RuntimeError( + "Cannot add new targets after the problem has been assembled (first solve())." + ) + + def _finalize_problem(self): + """Finalize the optimization problem. + + Raises: + RuntimeError: If the optimization problem has already been built. + """ + if self.cost_terms: + self.opti.minimize(cs.sum(cs.vertcat(*self.cost_terms))) + else: + raise RuntimeError( + "No targets added to the problem. Please add at least one target." + ) + self._problem_built = True + + def _check_target_type(self, frame: str, expected: TargetType): + """Check the target type for a given frame. + + Args: + frame (str): The name of the frame to check. + expected (TargetType): The expected target type. + + Raises: + ValueError: If the target type does not match the expected type. + """ + if self.targets[frame].target_type is not expected: + raise ValueError(f"Target '{frame}' is not of type {expected.name}")