From e20e14b334d27b9454014259b15c40eb68a190ac Mon Sep 17 00:00:00 2001 From: rexlee Date: Tue, 2 Apr 2024 18:37:49 -0700 Subject: [PATCH] add joint mimic --- dex_retargeting/optimizer.py | 132 ++++++++++++++++++++++++++++++++++- 1 file changed, 131 insertions(+), 1 deletion(-) diff --git a/dex_retargeting/optimizer.py b/dex_retargeting/optimizer.py index f6ff21c..ee5af7c 100644 --- a/dex_retargeting/optimizer.py +++ b/dex_retargeting/optimizer.py @@ -23,6 +23,7 @@ def __init__( self.wrist_link_name = wrist_link_name joint_names = [joint.get_name() for joint in robot.get_active_joints()] + # print(joint_names) target_joint_index = [] for target_joint_name in target_joint_names: if target_joint_name not in joint_names: @@ -32,6 +33,7 @@ def __init__( self.target_joint_indices = np.array(target_joint_index) self.fixed_joint_indices = np.array([i for i in range(robot.dof) if i not in target_joint_index], dtype=int) self.opt = nlopt.opt(nlopt.LD_SLSQP, len(target_joint_index)) + self.opt.set_maxeval(100) self.dof = len(target_joint_index) # Target @@ -41,12 +43,47 @@ def __init__( link_names = [link.get_name() for link in self.robot.get_links()] self.has_free_joint = len([name for name in link_names if "dummy" in name]) >= 6 + # mimic joints + self.parent_joint_indices = np.array([]) + self.child_joint_indices = np.array([]) + self.mimic_factors = np.array([]) + def set_joint_limit(self, joint_limits: np.ndarray): if joint_limits.shape != (self.dof, 2): raise ValueError(f"Expect joint limits have shape: {(self.dof, 2)}, but get {joint_limits.shape}") self.opt.set_lower_bounds(joint_limits[:, 0].tolist()) self.opt.set_upper_bounds(joint_limits[:, 1].tolist()) + def set_joint_mimic(self, parent_joint_names, child_joint_names, mimic_factors, tol=1e-2): + if len(parent_joint_names) != len(child_joint_names) or len(parent_joint_names) != len(mimic_factors): + raise ValueError("Expect lists to be of the same length") + n_mimic = len(mimic_factors) + tols = np.ones(n_mimic) * tol + joint_names = [joint.get_name() for joint in self.robot.get_active_joints()] + parent_joint_indices = [] + child_joint_indices = [] + for parent_joint_name in parent_joint_names: + if parent_joint_name not in joint_names: + raise ValueError(f"Joint {parent_joint_name} given does not appear to be in robot XML.") + parent_joint_indices.append(joint_names.index(parent_joint_name)) + for child_joint_name in child_joint_names: + if child_joint_name not in joint_names: + raise ValueError(f"Joint {child_joint_name} given does not appear to be in robot XML.") + child_joint_indices.append(joint_names.index(child_joint_name)) + self.parent_joint_indices = np.array(parent_joint_indices, dtype=int) + self.child_joint_indices = np.array(child_joint_indices, dtype=int) + self.mimic_factors = np.array(mimic_factors) + + def constraints(result: np.ndarray, x: np.ndarray, grad: np.array): + if grad.size > 0: + for i in range(n_mimic): + grad[i, parent_joint_indices[i]] = mimic_factors[i] + grad[i, child_joint_indices[i]] = -1 + for i in range(n_mimic): + result[i] = mimic_factors[i] * x[parent_joint_indices[i]] - x[child_joint_indices[i]] + + self.opt.add_equality_mconstraint(constraints, tols) + def get_last_result(self): return self.opt.last_optimize_result() @@ -91,7 +128,8 @@ def __init__( super().__init__(robot, wrist_link_name, target_joint_names, target_link_human_indices) self.body_names = target_link_names self.huber_loss = torch.nn.SmoothL1Loss(beta=huber_delta) - self.norm_delta = norm_delta + self.norm_delta = np.float64(norm_delta) + print(self.norm_delta) # Sanity check and cache link indices self.target_link_indices = self.get_link_indices(target_link_names) @@ -255,6 +293,13 @@ def objective(x: np.ndarray, grad: np.ndarray) -> float: for index in self.robot_link_indices ] + if len(self.mimic_factors) != 0: + for i, m in enumerate(self.mimic_factors): + p_idx = self.parent_joint_indices[i] + c_idx = self.child_joint_indices[i] + jacobians[:, :, p_idx] = jacobians[:, :, p_idx] + m * jacobians[:, :, c_idx] + jacobians[:, :, c_idx] = 0 + huber_distance.backward() grad_pos = torch_body_pos.grad.cpu().numpy()[:, None, :] grad_qpos = np.matmul(grad_pos, np.array(jacobians)) @@ -263,6 +308,8 @@ def objective(x: np.ndarray, grad: np.ndarray) -> float: grad_qpos += 2 * self.norm_delta * (x - last_qpos) grad[:] = grad_qpos[:] + if len(self.mimic_factors) != 0: + grad[self.child_joint_indices] = self.mimic_factors * grad[self.parent_joint_indices] return result @@ -494,3 +541,86 @@ def retarget(self, ref_value, fixed_qpos, last_qpos=None): ref_value, fixed_qpos, np.array(last_qpos).astype(np.float32) ) return self.optimize(objective_fn, last_qpos) + + +# This is a simple IK calculator +# for the convenience of sequential retargeting pipeline +class IKOptimizer(Optimizer): + retargeting_type = "ik" + + def __init__( + self, + robot: sapien.Articulation, + target_joint_names: List[str], + wrist_link_name: str, + huber_delta=0.02, + norm_delta=4e-3, + ): + super().__init__(robot, wrist_link_name, target_joint_names, []) + self.huber_loss = torch.nn.SmoothL1Loss(beta=huber_delta) + self.norm_delta = norm_delta + self.target_link_index = self.get_link_indices([wrist_link_name])[0] + self.use_sparse_jacobian = True + self.opt.set_ftol_abs(1e-5) + + def _get_objective_function(self, target_pos: np.ndarray, fixed_qpos: np.ndarray, last_qpos: np.ndarray): + qpos = np.zeros(self.robot_dof) + qpos[self.fixed_joint_indices] = fixed_qpos + torch_target_pos = torch.as_tensor(target_pos[0:3]) + torch_target_pos.requires_grad_(False) + torch_target_rot = torch.as_tensor(target_pos[3:7]) + torch_target_rot.requires_grad_(False) + + def objective(x: np.ndarray, grad: np.ndarray) -> float: + qpos[self.target_joint_indices] = x + self.model.compute_forward_kinematics(qpos) + target_link_pose = self.model.get_link_pose(self.target_link_index) + body_pos = target_link_pose.p + body_rot = target_link_pose.q + + # Torch computation for accurate loss and grad + torch_body_pos = torch.as_tensor(body_pos) + torch_body_pos.requires_grad_() + torch_body_rot = torch.as_tensor(body_rot) + torch_body_rot.requires_grad_() + + # Loss term for kinematics retargeting based on 3D position error + pos_distance = self.huber_loss(torch_body_pos, torch_target_pos) + rot_distance = 1 - torch.sum(torch_body_rot * torch_target_rot) ** 2 + huber_distance = pos_distance + # huber_distance = pos_distance + rot_distance + # huber_distance = torch.norm(torch_body_pos - torch_target_pos, dim=1).mean() + result = huber_distance.cpu().detach().item() + + if grad.size > 0: + link_spatial_jacobian = self.model.compute_single_link_local_jacobian(qpos, self.target_link_index)[ + :3, self.target_joint_indices + ] + link_rot = self.model.get_link_pose(self.target_link_index).to_transformation_matrix()[:3, :3] + jacobian = link_rot @ link_spatial_jacobian + + huber_distance.backward() + grad_pos = torch_body_pos.grad.cpu().numpy()[None, :] + grad_qpos = np.matmul(grad_pos, np.array(jacobian)) + grad_qpos = grad_qpos.mean(1).sum(0) + + grad_qpos += 2 * self.norm_delta * (x - last_qpos) + + grad[:] = grad_qpos[:] + + return result + + return objective + + def retarget(self, ref_value, fixed_qpos, last_qpos=None): + if len(fixed_qpos) != len(self.fixed_joint_indices): + raise ValueError( + f"Optimizer has {len(self.fixed_joint_indices)} joints but non_target_qpos {fixed_qpos} is given" + ) + if last_qpos is None: + last_qpos = np.zeros(self.dof) + if isinstance(last_qpos, np.ndarray): + last_qpos = last_qpos.astype(np.float32) + last_qpos = list(last_qpos) + objective_fn = self._get_objective_function(ref_value, fixed_qpos, np.array(last_qpos).astype(np.float32)) + return self.optimize(objective_fn, last_qpos)