Fix opw solver

docs/source/overview/sim/solvers/srs_solver.md

@@ -51,7 +51,7 @@ cfg = SRSSolverCfg(
     end_link_name="left_ee",
     root_link_name="left_arm_base",
     dh_params=arm_params.dh_params,
-    qpos_limits=arm_params.qpos_limits,
+    user_qpos_limit=arm_params.qpos_limits,
     T_e_oe=arm_params.T_e_oe,
     T_b_ob=arm_params.T_b_ob,
     link_lengths=arm_params.link_lengths,

docs/source/tutorial/solver.rst

@@ -95,7 +95,7 @@ API Reference
            """Compute the Jacobian matrix for the given joint positions."""
 
 - **set_ik_nearst_weight**: Set weights for IK nearest neighbor search.
-- **set_position_limits / get_position_limits**: Set or get joint position limits.
+- **set_qpos_limits / get_qpos_limits**: Set or get joint position limits.
 - **set_tcp / get_tcp**: Set or get the tool center point (TCP) transformation.
 
 Configuration

embodichain/lab/sim/objects/robot.py

@@ -934,6 +934,9 @@ def init_solver(self, cfg: Union[SolverCfg, Dict[str, SolverCfg]]) -> None:
                     ):
                         solver_cfg.joint_names = self.cfg.control_parts[part_name]
                     self._solvers[name] = solver_cfg.init_solver(device=self.device)
+                    joint_ids = self.get_joint_ids(name=part_name)
+                    joint_limits = self._data.qpos_limits[0][joint_ids]
+                    self._solvers[name].update_with_robot_limit(joint_limits)
 
     def get_solver(self, name: str | None = None) -> BaseSolver | None:
         """Get the kinematic solver for a specific control part.

embodichain/lab/sim/robots/cobotmagic.py

@@ -115,13 +115,15 @@ def _build_default_cfgs() -> Dict[str, Any]:
                     tcp=np.array(
                         [[-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]]
                     ),
+                    safe_margin=5.0 * np.pi / 180.0,
                 ),
                 "right_arm": OPWSolverCfg(
                     end_link_name="right_link6",
                     root_link_name="right_arm_base",
                     tcp=np.array(
                         [[-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]]
                     ),
+                    safe_margin=5.0 * np.pi / 180.0,
                 ),
             },
             "min_position_iters": 8,

embodichain/lab/sim/robots/dexforce_w1/cfg.py

@@ -159,7 +159,7 @@ def _build_default_solver_cfg(is_industrial: bool) -> SolverCfg:
                 end_link_name="right_ee",
                 root_link_name="right_arm_base",
                 dh_params=w1_right_arm_params.dh_params,
-                qpos_limits=w1_right_arm_params.qpos_limits,
+                user_qpos_limits=w1_right_arm_params.qpos_limits,
                 T_e_oe=w1_right_arm_params.T_e_oe,
                 T_b_ob=w1_right_arm_params.T_b_ob,
                 link_lengths=w1_right_arm_params.link_lengths,
@@ -170,7 +170,7 @@ def _build_default_solver_cfg(is_industrial: bool) -> SolverCfg:
                 end_link_name="left_ee",
                 root_link_name="left_arm_base",
                 dh_params=w1_left_arm_params.dh_params,
-                qpos_limits=w1_left_arm_params.qpos_limits,
+                user_qpos_limits=w1_left_arm_params.qpos_limits,
                 T_e_oe=w1_left_arm_params.T_e_oe,
                 T_b_ob=w1_left_arm_params.T_b_ob,
                 link_lengths=w1_left_arm_params.link_lengths,

embodichain/lab/sim/solvers/base_solver.py

@@ -72,6 +72,13 @@ class SolverCfg:
     when multiple solutions are available.
     """
 
+    user_qpos_limits: List[float] | None = None
+    """
+        User defined Joint position limits [2, DOF] for the solver. 
+        If not provided (None), this value will replace by joint limits defined in urdf when solver init from robot.
+        If provided, the solver will use the intersection of user defined limits and urdf limits as the final joint limits.
+    """
+
     @abstractmethod
     def init_solver(self, device: torch.device, **kwargs) -> "BaseSolver":
         pass
@@ -165,6 +172,8 @@ def __init__(self, cfg: SolverCfg = None, device: str = None, **kwargs):
                 device=self.device,
             )
 
+        self._init_qpos_limits()
+
     def set_ik_nearest_weight(
         self, ik_weight: np.ndarray, joint_ids: np.ndarray | None = None
     ) -> bool:
@@ -223,51 +232,106 @@ def get_ik_nearest_weight(self):
         """
         return self.ik_nearest_weight
 
-    def set_position_limits(
+    def _init_qpos_limits(self):
+        self.lower_qpos_limits = None
+        self.upper_qpos_limits = None
+        if self.cfg.user_qpos_limits is not None:
+            # robot qpos limits from config, expected shape [DOF, 2]
+            user_qpos_limits = torch.tensor(
+                self.cfg.user_qpos_limits, dtype=torch.float32, device=self.device
+            )
+            if user_qpos_limits.shape == (2, self.dof):
+                self.set_qpos_limits(
+                    lower_qpos_limits=user_qpos_limits[0],
+                    upper_qpos_limits=user_qpos_limits[1],
+                )
+            elif user_qpos_limits.shape == (self.dof, 2):
+                self.set_qpos_limits(
+                    lower_qpos_limits=user_qpos_limits[:, 0],
+                    upper_qpos_limits=user_qpos_limits[:, 1],
+                )
+            else:
+                logger.log_error(
+                    f"user_qpos_limits must have shape (2, {self.dof}) or ({self.dof}, 2), but got {user_qpos_limits.shape}."
+                )
+        elif self.pk_serial_chain is not None:
+            self.set_qpos_limits(
+                lower_qpos_limits=self.pk_serial_chain.low,
+                upper_qpos_limits=self.pk_serial_chain.high,
+            )
+
+    def update_with_robot_limit(self, robot_qpos_limits: torch.Tensor):
+        """Update with robot joint limits.
+            Make sure the solver's joint limits are within the robot's joint limits.
+
+        Args:
+            robot_qpos_limits (torch.Tensor): [DOF, 2] tensor of joint limits from the robot data
+        """
+        robot_lower_limits = robot_qpos_limits[:, 0]
+        robot_upper_limits = robot_qpos_limits[:, 1]
+        if self.lower_qpos_limits is not None:
+            if torch.any(self.lower_qpos_limits < robot_lower_limits):
+                logger.log_warning(
+                    "Solver lower_qpos_limits are smaller than robot limits. Clamping to robot limits."
+                )
+                self.lower_qpos_limits = torch.max(
+                    self.lower_qpos_limits, robot_lower_limits
+                )
+        else:
+            self.lower_qpos_limits = robot_lower_limits
+        if self.upper_qpos_limits is not None:
+            if torch.any(self.upper_qpos_limits > robot_upper_limits):
+                logger.log_warning(
+                    "Solver upper_qpos_limits are larger than robot limits. Clamping to robot limits."
+                )
+                self.upper_qpos_limits = torch.min(
+                    self.upper_qpos_limits, robot_upper_limits
+                )
+        else:
+            self.upper_qpos_limits = robot_upper_limits
+
+    def set_qpos_limits(
         self,
-        lower_position_limits: List[float],
-        upper_position_limits: List[float],
+        lower_qpos_limits: List[float],
+        upper_qpos_limits: List[float],
     ) -> bool:
         r"""Sets the upper and lower joint position limits.
 
         Parameters:
-            lower_position_limits (List[float]): A list of lower limits for each joint.
-            upper_position_limits (List[float]): A list of upper limits for each joint.
+            lower_qpos_limits (List[float]): A list of lower limits for each joint.
+            upper_qpos_limits (List[float]): A list of upper limits for each joint.
 
         Returns:
             bool: True if limits are successfully set, False if the input is invalid.
         """
-        if (
-            len(lower_position_limits) != self.model.nq
-            or len(upper_position_limits) != self.model.nq
-        ):
-            logger.log_warning("Length of limits must match the number of joints.")
-            return False
 
         if any(
-            lower > upper
-            for lower, upper in zip(lower_position_limits, upper_position_limits)
+            lower > upper for lower, upper in zip(lower_qpos_limits, upper_qpos_limits)
         ):
             logger.log_warning(
                 "Each lower limit must be less than or equal to the corresponding upper limit."
             )
             return False
 
-        self.lower_position_limits = np.array(lower_position_limits)
-        self.upper_position_limits = np.array(upper_position_limits)
+        self.lower_qpos_limits = torch.tensor(
+            lower_qpos_limits, dtype=float, device=self.device
+        )
+        self.upper_qpos_limits = torch.tensor(
+            upper_qpos_limits, dtype=float, device=self.device
+        )
         return True
 
-    def get_position_limits(self) -> dict:
+    def get_qpos_limits(self) -> dict:
         r"""Returns the current joint position limits.
 
         Returns:
             dict: A dictionary containing:
-                - lower_position_limits (List[float]): The current lower limits for each joint.
-                - upper_position_limits (List[float]): The current upper limits for each joint.
+                - lower_qpos_limits (List[float]): The current lower limits for each joint.
+                - upper_qpos_limits (List[float]): The current upper limits for each joint.
         """
         return {
-            "lower_position_limits": self.lower_position_limits.tolist(),
-            "upper_position_limits": self.upper_position_limits.tolist(),
+            "lower_qpos_limits": self.lower_qpos_limits.tolist(),
+            "upper_qpos_limits": self.upper_qpos_limits.tolist(),
         }
 
     def set_tcp(self, xpos: np.ndarray):

embodichain/lab/sim/solvers/opw_solver.py

@@ -29,7 +29,7 @@
     OPWparam,
     opw_fk_kernel,
     opw_ik_kernel,
-    opw_best_ik_kernel,
+    opw_ik_select_kernel,
     wp_vec6f,
 )
 from embodichain.utils.device_utils import standardize_device_string
@@ -72,6 +72,9 @@ class OPWSolverCfg(SolverCfg):
     # Parameters for the inverse-kinematics method.
     ik_params: dict | None = None
 
+    # safe margin for joint limits, in radians
+    safe_margin: float = 5.0 * np.pi / 180.0
+
     def init_solver(
         self, device: torch.device = torch.device("cpu"), **kwargs
     ) -> "OPWSolver":
@@ -247,50 +250,87 @@ def get_ik_warp(
         N_SOL = 8
         DOF = 6
         n_sample = target_xpos.shape[0]
+        kernel_device = standardize_device_string(self.device)
 
         if target_xpos.shape == (4, 4):
-            target_xpos_batch = target_xpos[None, :, :]
+            target_xpos_batch = target_xpos[None, :, :].to(kernel_device)
         else:
-            target_xpos_batch = target_xpos
+            target_xpos_batch = target_xpos.to(kernel_device)
         target_xpos_wp = wp.from_torch(target_xpos_batch.reshape(-1))
 
         all_qpos_wp = wp.zeros(
             n_sample * N_SOL * DOF,
             dtype=float,
-            device=standardize_device_string(self.device),
+            device=standardize_device_string(kernel_device),
         )
         all_ik_valid_wp = wp.zeros(
-            n_sample * N_SOL, dtype=int, device=standardize_device_string(self.device)
+            n_sample * N_SOL, dtype=int, device=standardize_device_string(kernel_device)
         )
 
         # TODO: whether require gradient
+        offsets_ = self.offsets.to(standardize_device_string(kernel_device))
+        sign_corrections_ = self.sign_corrections.to(
+            standardize_device_string(kernel_device)
+        )
+        lower_limits_ = wp_vec6f(
+            self.lower_qpos_limits[0],
+            self.lower_qpos_limits[1],
+            self.lower_qpos_limits[2],
+            self.lower_qpos_limits[3],
+            self.lower_qpos_limits[4],
+            self.lower_qpos_limits[5],
+        )
+        upper_limits_ = wp_vec6f(
+            self.upper_qpos_limits[0],
+            self.upper_qpos_limits[1],
+            self.upper_qpos_limits[2],
+            self.upper_qpos_limits[3],
+            self.upper_qpos_limits[4],
+            self.upper_qpos_limits[5],
+        )
         wp.launch(
             kernel=opw_ik_kernel,
             dim=(n_sample),
             inputs=(
                 target_xpos_wp,
                 self._tcp_inv_warp,
                 self.params,
-                self.offsets,
-                self.sign_corrections,
+                offsets_,
+                sign_corrections_,
+                lower_limits_,
+                upper_limits_,
+                self.cfg.safe_margin,
             ),
             outputs=[all_qpos_wp, all_ik_valid_wp],
-            device=standardize_device_string(self.device),
+            device=standardize_device_string(kernel_device),
         )
 
         if return_all_solutions:
             all_qpos = wp.to_torch(all_qpos_wp).reshape(n_sample, N_SOL, DOF)
             all_ik_valid = wp.to_torch(all_ik_valid_wp).reshape(n_sample, N_SOL)
             return all_ik_valid, all_qpos
-
         if qpos_seed is not None:
-            qpos_seed_wp = wp.from_torch(qpos_seed.reshape(-1))
+            if qpos_seed.shape == (
+                n_sample,
+                DOF,
+            ):
+                qpos_seed_ = qpos_seed.to(kernel_device)
+            elif qpos_seed.shape == (DOF,):
+                qpos_seed_ = (
+                    qpos_seed.unsqueeze(0).repeat(n_sample, 1).to(kernel_device)
+                )
+            else:
+                logger.log_error(
+                    f"Invalid shape for qpos_seed: {qpos_seed.shape}. Expected ({n_sample}, {DOF}) or ({DOF},)."
+                )
+            qpos_seed_wp = wp.from_torch(qpos_seed_)
         else:
-            qpos_seed_wp = wp.zeros(
-                n_sample * DOF,
-                dtype=float,
-                device=standardize_device_string(self.device),
+            qpos_seed = torch.zeros(
+                (n_sample, DOF), dtype=torch.float32, device=kernel_device
             )
+            qpos_seed_wp = wp.from_torch(qpos_seed)
+        all_qpos_wp = all_qpos_wp.reshape((n_sample, N_SOL, DOF))
+        all_ik_valid_wp = all_ik_valid_wp.reshape((n_sample, N_SOL))
         joint_weight = kwargs.get("joint_weight", torch.ones(size=(DOF,), dtype=float))
         joint_weight_wp = wp_vec6f(
             joint_weight[0],
@@ -301,25 +341,31 @@ def get_ik_warp(
             joint_weight[5],
         )
         best_ik_result_wp = wp.zeros(
-            n_sample * 6, dtype=float, device=standardize_device_string(self.device)
+            (n_sample, 6), dtype=float, device=standardize_device_string(kernel_device)
         )
         best_ik_valid_wp = wp.zeros(
-            n_sample, dtype=int, device=standardize_device_string(self.device)
+            n_sample, dtype=int, device=standardize_device_string(kernel_device)
         )
         wp.launch(
-            kernel=opw_best_ik_kernel,
+            kernel=opw_ik_select_kernel,
             dim=(n_sample),
             inputs=[
                 all_qpos_wp,
                 all_ik_valid_wp,
                 qpos_seed_wp,
                 joint_weight_wp,
             ],
-            outputs=[best_ik_result_wp, best_ik_valid_wp],
-            device=standardize_device_string(self.device),
+            outputs=[
+                best_ik_result_wp,
+                best_ik_valid_wp,
+            ],
+            device=standardize_device_string(kernel_device),
+        )
+
+        best_ik_result = (
+            wp.to_torch(best_ik_result_wp).reshape(n_sample, 1, 6).to(self.device)
         )
-        best_ik_result = wp.to_torch(best_ik_result_wp).reshape(n_sample, 1, 6)
-        best_ik_valid = wp.to_torch(best_ik_valid_wp)
+        best_ik_valid = wp.to_torch(best_ik_valid_wp).to(self.device)
         return best_ik_valid, best_ik_result
 
     def _calculate_dynamic_weights(