diff --git a/src/adam/casadi/inverse_kinematics.py b/src/adam/casadi/inverse_kinematics.py index b13b0974..e52487d1 100644 --- a/src/adam/casadi/inverse_kinematics.py +++ b/src/adam/casadi/inverse_kinematics.py @@ -60,11 +60,14 @@ def __init__( 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.joint_pos = self.opti.variable(self.ndof) + self.base_homogeneous = SE3.from_position_quaternion( + self.base_pos, self.base_quat + ).as_matrix() 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.opti.set_initial(self.joint_pos, np.zeros(self.ndof)) # default to zero self.targets: dict[str, Target] = {} self.cost_terms: list[cs.MX] = [] @@ -82,6 +85,25 @@ def __init__( self.opti.solver("ipopt", solver_settings) + def get_opti_variables(self) -> dict[str, cs.SX]: + """Get the optimization variables of the IK solver. + + Returns: + dict[str, cs.SX]: Dictionary containing the optimization variables. + + - "base_pos": Position of the base (3D vector). + - "base_quat": Quaternion representing the base orientation (4D vector with x, y, z, w). + - "joint_pos": Joint variables (Vector with length equal to the number of non-fixed joints). + - "base_homogeneous": Homogeneous transformation matrix of the base (4x4 matrix, from position and quaternion). + + """ + return { + "base_pos": self.base_pos, + "base_quat": self.base_quat, + "joint_pos": self.joint_pos, + "base_homogeneous": self.base_homogeneous, + } + def set_solver(self, solver_name: str, solver_settings: dict[str, Any]): """Set the solver for the optimization problem. @@ -117,7 +139,7 @@ def set_default_joint_limit_constraints(self): if joint_limit is not None: self.opti.subject_to( self.opti.bounded( - joint_limit.lower, self.joint_var[i], joint_limit.upper + joint_limit.lower, self.joint_pos[i], joint_limit.upper ) ) @@ -133,7 +155,7 @@ def set_joint_limits(self, lower_bounds: np.ndarray, upper_bounds: np.ndarray): 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]) + self.opti.bounded(lower_bounds[i], self.joint_pos[i], upper_bounds[i]) ) def add_target_position( @@ -160,7 +182,7 @@ def add_target_position( 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 + self.base_transform(), self.joint_pos )[:3, 3] if as_soft_constraint: @@ -199,10 +221,10 @@ def add_frames_constraint( 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 + self.base_transform(), self.joint_pos ) H_child = self.kd.forward_kinematics_fun(child_frame)( - self.base_transform(), self.joint_var + self.base_transform(), self.joint_pos ) p_parent, p_child = H_parent[:3, 3], H_child[:3, 3] R_parent, R_child = H_parent[:3, :3], H_child[:3, :3] @@ -290,10 +312,10 @@ def add_min_distance_constraint(self, frames_list: list[str], distance: float): ) 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 + self.base_transform(), self.joint_pos )[:3, 3] p_j = self.kd.forward_kinematics_fun(frame_j)( - self.base_transform(), self.joint_var + self.base_transform(), self.joint_pos )[:3, 3] dist_sq = cs.sumsqr(p_i - p_j) self.opti.subject_to(dist_sq >= distance**2) @@ -318,7 +340,7 @@ def add_target_orientation( 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 + self.base_transform(), self.joint_pos ) R_fk = H_fk[:3, :3] # proxy for rotation error: trace(R) = 1 + 2 * cos(theta), where theta is the angle of rotation @@ -359,7 +381,7 @@ def add_target_pose( 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 + self.base_transform(), self.joint_pos ) p_fk, R_fk = H_fk[:3, 3], H_fk[:3, :3] @@ -486,7 +508,7 @@ def set_initial_guess(self, base_transform: np.ndarray, joint_values: np.ndarray 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) + self.opti.set_initial(self.joint_pos, joint_values) def solve(self): """Solve the optimization problem.""" @@ -499,7 +521,7 @@ def solve(self): 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.joint_pos, self._cached_sol.value(self.joint_pos) ) self._cached_sol = self.opti.solve() @@ -516,7 +538,7 @@ def get_solution(self): 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) + joints = self._cached_sol.value(self.joint_pos) H = SE3.from_position_quaternion(pos, quat).as_matrix() # 4x4 numpy return np.array(H), np.array(joints) @@ -556,3 +578,20 @@ def _check_target_type(self, frame: str, expected: TargetType): """ if self.targets[frame].target_type is not expected: raise ValueError(f"Target '{frame}' is not of type {expected.name}") + + def add_cost(self, cost: cs.MX): + """Add a custom cost term to the optimization problem. + + This method appends the provided cost term to the `cost_terms` list. + During the `_finalize_problem` step, all cost terms in this list are + combined using `cs.sum(cs.vertcat(*self.cost_terms))` and minimized + as part of the objective function. + + Note: This method ensures that the optimization graph is still modifiable + before adding the cost term. Once the problem is finalized (after the + first call to `solve()`), no new cost terms can be added. + Args: + cost (cs.MX): The cost term to add. + """ + self._ensure_graph_modifiable() + self.cost_terms.append(cost)