add joint mimic

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)