Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Cyclic Coordinate Descent (CCD) Inverse Kinematics solver.""" | |
| import torch | |
| import gem.utils.matrix as matrix | |
| from gem.utils.geo_transform import qbetween | |
| from gem.utils.quaternion import qinv, qmul, qrot, qslerp | |
| from gem.utils.rotation_conversions import matrix_to_quaternion | |
| class CCD_IK: | |
| def __init__( | |
| self, | |
| local_mat, | |
| parent, | |
| target_ind, | |
| target_pos=None, | |
| target_rot=None, | |
| kinematic_chain=None, | |
| max_iter=2, | |
| threshold=0.001, | |
| pos_weight=1.0, | |
| rot_weight=0.0, | |
| ): | |
| if kinematic_chain is None: | |
| kinematic_chain = range(local_mat.shape[-3]) | |
| global_mat = matrix.forward_kinematics(local_mat, parent) | |
| # Extract kinematic chain and assign root as global mat (do not modify root during IK) | |
| local_mat = local_mat.clone() | |
| local_mat = local_mat[..., kinematic_chain, :, :] | |
| local_mat[..., 0, :, :] = global_mat[..., kinematic_chain[0], :, :] | |
| parent = [i - 1 for i in range(len(kinematic_chain))] | |
| self.local_mat = local_mat | |
| self.global_mat = matrix.forward_kinematics(local_mat, parent) # (*, J, 4, 4) | |
| self.parent = parent | |
| self.target_ind = target_ind | |
| self.target_pos = target_pos # (*, O, 3) or None | |
| if target_rot is not None: | |
| self.target_q = matrix_to_quaternion(target_rot) # (*, O, 4) | |
| else: | |
| self.target_q = None | |
| self.threshold = threshold | |
| self.J_N = self.local_mat.shape[-3] | |
| self.target_N = len(target_ind) | |
| self.max_iter = max_iter | |
| self.pos_weight = pos_weight | |
| self.rot_weight = rot_weight | |
| def solve(self): | |
| for _ in range(self.max_iter): | |
| self._optimize(1) | |
| return self.local_mat | |
| def _optimize(self, i): | |
| if i == self.J_N - 1: | |
| return | |
| pos = matrix.get_position(self.global_mat)[..., i, :] # (*, 3) | |
| rot = matrix.get_rotation(self.global_mat)[..., i, :, :] # (*, 3, 3) | |
| quat = matrix_to_quaternion(rot) # (*, 4) | |
| x_vec = torch.zeros((quat.shape[:-1] + (3,)), device=quat.device) | |
| x_vec[..., 0] = 1.0 | |
| x_vec_sum = torch.zeros_like(x_vec) | |
| y_vec = torch.zeros((quat.shape[:-1] + (3,)), device=quat.device) | |
| y_vec[..., 1] = 1.0 | |
| y_vec_sum = torch.zeros_like(y_vec) | |
| count = 0 | |
| for target_i, j in enumerate(self.target_ind): | |
| if i >= j: | |
| continue | |
| end_pos = matrix.get_position(self.global_mat)[..., j, :] # (*, 3) | |
| end_rot = matrix.get_rotation(self.global_mat)[..., j, :, :] # (*, 3, 3) | |
| end_quat = matrix_to_quaternion(end_rot) # (*, 4) | |
| if self.target_pos is not None: | |
| target_pos = self.target_pos[..., target_i, :] # (*, 3) | |
| solved_pos_target_quat = qslerp( | |
| quat, | |
| qmul(qbetween(end_pos - pos, target_pos - pos), quat), | |
| self._get_weight(i), | |
| ) | |
| x_vec_sum += qrot(solved_pos_target_quat, x_vec) | |
| y_vec_sum += qrot(solved_pos_target_quat, y_vec) | |
| if self.pos_weight > 0: | |
| count += 1 | |
| if self.target_q is not None: | |
| if target_i < self.target_N - 1: | |
| continue | |
| target_q = self.target_q[..., target_i, :] # (*, 4) | |
| solved_q_target_quat = qslerp( | |
| quat, | |
| qmul(qmul(target_q, qinv(end_quat)), quat), | |
| self._get_weight(i), | |
| ) | |
| x_vec_sum += qrot(solved_q_target_quat, x_vec) * self.rot_weight | |
| y_vec_sum += qrot(solved_q_target_quat, y_vec) * self.rot_weight | |
| if self.rot_weight > 0: | |
| count += 1 | |
| if count > 0: | |
| x_vec_avg = matrix.normalize(x_vec_sum / count) | |
| y_vec_avg = matrix.normalize(y_vec_sum / count) | |
| z_vec_avg = torch.cross(x_vec_avg, y_vec_avg, dim=-1) | |
| solved_rot = torch.stack([x_vec_avg, y_vec_avg, z_vec_avg], dim=-1) | |
| parent_rot = matrix.get_rotation(self.global_mat)[..., self.parent[i], :, :] | |
| solved_local_rot = matrix.get_mat_BtoA(parent_rot, solved_rot) | |
| self.local_mat[..., i, :-1, :-1] = solved_local_rot | |
| self.global_mat = matrix.forward_kinematics(self.local_mat, self.parent) | |
| self._optimize(i + 1) | |
| def _get_weight(self, i): | |
| return (i + 1) / self.J_N | |