diff --git a/docs/source/overview/sim/solvers/srs_solver.md b/docs/source/overview/sim/solvers/srs_solver.md index 3cabb57e..2b26ee6d 100644 --- a/docs/source/overview/sim/solvers/srs_solver.md +++ b/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, diff --git a/docs/source/tutorial/solver.rst b/docs/source/tutorial/solver.rst index b3c95807..61300096 100644 --- a/docs/source/tutorial/solver.rst +++ b/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 diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index e6dac158..07273e80 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/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. diff --git a/embodichain/lab/sim/robots/cobotmagic.py b/embodichain/lab/sim/robots/cobotmagic.py index 1ffdcd71..2c2885d1 100644 --- a/embodichain/lab/sim/robots/cobotmagic.py +++ b/embodichain/lab/sim/robots/cobotmagic.py @@ -115,6 +115,7 @@ 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", @@ -122,6 +123,7 @@ 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, ), }, "min_position_iters": 8, diff --git a/embodichain/lab/sim/robots/dexforce_w1/cfg.py b/embodichain/lab/sim/robots/dexforce_w1/cfg.py index c6586b4e..40f95b09 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/cfg.py +++ b/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, diff --git a/embodichain/lab/sim/solvers/base_solver.py b/embodichain/lab/sim/solvers/base_solver.py index 143e3a89..40c61af5 100644 --- a/embodichain/lab/sim/solvers/base_solver.py +++ b/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): diff --git a/embodichain/lab/sim/solvers/opw_solver.py b/embodichain/lab/sim/solvers/opw_solver.py index 4d8f9047..bdb68e34 100644 --- a/embodichain/lab/sim/solvers/opw_solver.py +++ b/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,23 +250,44 @@ 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), @@ -271,26 +295,42 @@ def get_ik_warp( 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,13 +341,13 @@ 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, @@ -315,11 +355,17 @@ def get_ik_warp( 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( diff --git a/embodichain/lab/sim/solvers/pinocchio_solver.py b/embodichain/lab/sim/solvers/pinocchio_solver.py index ec7e345a..f66f1685 100644 --- a/embodichain/lab/sim/solvers/pinocchio_solver.py +++ b/embodichain/lab/sim/solvers/pinocchio_solver.py @@ -129,9 +129,6 @@ def __init__(self, cfg: PinocchioSolverCfg, **kwargs): self.robot.model.njoints - 1 ) # Degrees of freedom of reduced robot joints - self.upper_position_limits = self.robot.model.upperPositionLimit - self.lower_position_limits = self.robot.model.lowerPositionLimit - self.ik_nearest_weight = np.ones(self.dof) # TODO: The Casadi-based solver is currently disabled due to stability issues. @@ -325,12 +322,14 @@ def qpos_to_limits( # Generate possible values for each joint dof_num = len(q) + lower_limits = self.lower_qpos_limits.to("cpu").numpy() + upper_limits = self.upper_qpos_limits.to("cpu").numpy() for i in range(dof_num): current_possible_values = [] # Calculate how many 2π fits into the adjustment to the limits - lower_adjustment = (q[i] - self.lower_position_limits[i]) // (2 * np.pi) - upper_adjustment = (self.upper_position_limits[i] - q[i]) // (2 * np.pi) + lower_adjustment = (q[i] - lower_limits[i]) // (2 * np.pi) + upper_adjustment = (upper_limits[i] - q[i]) // (2 * np.pi) # Consider the current value and its periodic adjustments for offset in range( @@ -339,15 +338,11 @@ def qpos_to_limits( adjusted_value = q[i] + offset * (2 * np.pi) # Check if the adjusted value is within limits - if ( - self.lower_position_limits[i] - <= adjusted_value - <= self.upper_position_limits[i] - ): + if lower_limits[i] <= adjusted_value <= upper_limits[i]: current_possible_values.append(adjusted_value) # Also check the original value - if self.lower_position_limits[i] <= q[i] <= self.upper_position_limits[i]: + if lower_limits[i] <= q[i] <= upper_limits[i]: current_possible_values.append(q[i]) if not current_possible_values: diff --git a/embodichain/lab/sim/solvers/pytorch_solver.py b/embodichain/lab/sim/solvers/pytorch_solver.py index cdcdc562..bfe5a080 100644 --- a/embodichain/lab/sim/solvers/pytorch_solver.py +++ b/embodichain/lab/sim/solvers/pytorch_solver.py @@ -174,9 +174,6 @@ def __init__( self.dof = self.pk_serial_chain.n_joints - self.upper_position_limits = self.pk_serial_chain.high - self.lower_position_limits = self.pk_serial_chain.low - def get_iteration_params(self) -> dict: r"""Returns the current iteration parameters. @@ -294,8 +291,8 @@ def _compute_inverse_kinematics( def _qpos_to_limits_single( q: torch.Tensor, joint_seed: torch.Tensor, - lower_position_limits: torch.Tensor, - upper_position_limits: torch.Tensor, + lower_qpos_limits: torch.Tensor, + upper_qpos_limits: torch.Tensor, ik_nearest_weight: torch.Tensor, periodic_mask: torch.Tensor = None, # Optional mask for periodic joints ) -> torch.Tensor: @@ -305,8 +302,8 @@ def _qpos_to_limits_single( Args: q (torch.Tensor): The initial joint positions. joint_seed (torch.Tensor): The seed joint positions for comparison. - lower_position_limits (torch.Tensor): The lower bounds for the joint positions. - upper_position_limits (torch.Tensor): The upper bounds for the joint positions. + lower_qpos_limits (torch.Tensor): The lower bounds for the joint positions. + upper_qpos_limits (torch.Tensor): The upper bounds for the joint positions. ik_nearest_weight (torch.Tensor): The weights for the inverse kinematics nearest calculation. periodic_mask (torch.Tensor, optional): Boolean mask indicating which joints are periodic. @@ -315,8 +312,8 @@ def _qpos_to_limits_single( """ device = q.device joint_seed = joint_seed.to(device) - lower = lower_position_limits.to(device) - upper = upper_position_limits.to(device) + lower = lower_qpos_limits.to(device) + upper = upper_qpos_limits.to(device) weight = ik_nearest_weight.to(device) # If periodic_mask is not provided, assume all joints are periodic @@ -359,7 +356,6 @@ def _qpos_to_limits( torch.Tensor: Batch of adjusted joint positions that fit within the limits, shape (M, dof), where M <= N (invalid candidates are filtered out). """ - periodic_mask = torch.ones_like( qpos_list_split[0], dtype=torch.bool, device=self.device ) @@ -368,8 +364,8 @@ def _qpos_to_limits( self._qpos_to_limits_single( q, joint_seed, - self.lower_position_limits, - self.upper_position_limits, + self.lower_qpos_limits, + self.upper_qpos_limits, self.ik_nearest_weight, periodic_mask, ) @@ -452,8 +448,6 @@ def get_ik( target_xpos = target_xpos @ torch.inverse(tcp_xpos) # Get joint limits and ensure shape matches dof - upper_limits = self.upper_position_limits.float() - lower_limits = self.lower_position_limits.float() batch_size = target_xpos.shape[0] @@ -461,7 +455,10 @@ def get_ik( num_samples=self._num_samples, dof=self.dof, device=self.device ) random_qpos_seeds = sampler.sample( - qpos_seed, lower_limits, upper_limits, batch_size + qpos_seed, + self.lower_qpos_limits, + self.upper_qpos_limits, + batch_size, ) target_xpos_repeated = sampler.repeat_target_xpos( target_xpos, self._num_samples diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index 64c4f492..d68f470b 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -51,9 +51,6 @@ class SRSSolverCfg(SolverCfg): dh_params = [] """Denavit-Hartenberg parameters for the robot's kinematic chain.""" - qpos_limits = [] - """Joint position limits for the robot.""" - T_b_ob = np.eye(4) """Base to observed base transform.""" @@ -107,9 +104,7 @@ def __init__(self, cfg: SRSSolverCfg, device: torch.device): self.device = device self.dofs = 7 self.dh_params = cfg.dh_params - self.qpos_limits = cfg.qpos_limits self.tcp_xpos = np.eye(4) - # Initialize transformation matrices self._parse_params() @@ -122,7 +117,6 @@ def _parse_params(self): # Convert configuration parameters to numpy arrays for efficient computation. self.dh_params_np = np.asarray(self.cfg.dh_params) - self.qpos_limits_np = np.asarray(self.cfg.qpos_limits) self.link_lengths_np = np.asarray(self.cfg.link_lengths) self.rotation_directions_np = np.asarray(self.cfg.rotation_directions) @@ -628,11 +622,6 @@ def _parse_params(self): dtype=float, device=standardize_device_string(self.device), ) - self.qpos_limits_wp = wp.array( - self.qpos_limits_np, - dtype=wp.vec2, - device=standardize_device_string(self.device), - ) self.link_lengths_wp = wp.array( self.link_lengths_np.flatten(), dtype=float, @@ -1197,6 +1186,21 @@ def __init__(self, cfg: SRSSolverCfg, num_envs: int, device: str, **kwargs): else: self.impl = _CPUSRSSolverImpl(cfg, self.device) + self._update_impl_qpos_limits() + + def _update_impl_qpos_limits(self): + qpos_limits = torch.vstack([self.lower_qpos_limits, self.upper_qpos_limits]).T + self.impl.qpos_limits_np = qpos_limits.cpu().numpy() + self.impl.qpos_limits_wp = wp.array( + self.impl.qpos_limits_np, + dtype=wp.vec2, + device=standardize_device_string(self.device), + ) + + def update_with_robot_limit(self, robot_qpos_limits): + super().update_with_robot_limit(robot_qpos_limits) + self._update_impl_qpos_limits() + def get_ik( self, target_xpos: torch.Tensor, diff --git a/embodichain/utils/warp/kinematics/opw_solver.py b/embodichain/utils/warp/kinematics/opw_solver.py index 1f1cf459..c152934c 100644 --- a/embodichain/utils/warp/kinematics/opw_solver.py +++ b/embodichain/utils/warp/kinematics/opw_solver.py @@ -30,6 +30,23 @@ def normalize_to_pi(angle: float) -> float: return wp.atan2(wp.sin(angle), wp.cos(angle)) +@wp.func +def normalize_in_limit(angle: float, lower: float, upper: float) -> float: + two_pi = 2.0 * wp.pi + k = wp.ceil((lower - angle) / two_pi) + result = angle + k * two_pi + return result + + +@wp.func +def is_within_limit( + angle: float, lower: float, upper: float, safe_margin: float +) -> bool: + if angle < lower + safe_margin or angle > upper - safe_margin: + return False + return True + + @wp.func def safe_acos(x: float) -> float: return wp.acos(wp.clamp(x, -1.0, 1.0)) @@ -219,6 +236,9 @@ def opw_ik_kernel( params: OPWparam, offsets: wp.array(dtype=float), sign_corrections: wp.array(dtype=float), + lower_limits: wp_vec6f, + upper_limits: wp_vec6f, + safe_margin: float, qpos: wp.array(dtype=float), ik_valid: wp.array(dtype=int), ): @@ -433,8 +453,10 @@ def opw_ik_kernel( for k in range(DOF): idx = j * DOF + k - qpos[qpos_start + k] = normalize_to_pi( - (theta[idx] + offsets[k]) * sign_corrections[k] + qpos[qpos_start + k] = normalize_in_limit( + (theta[idx] + offsets[k]) * sign_corrections[k], + lower=lower_limits[k], + upper=upper_limits[k], ) # filter invalid solutions @@ -449,42 +471,46 @@ def opw_ik_kernel( ) t_err, r_err = get_transform_err(check_ee_pose, ee_pose) # mark invalid solutions (cannot pass ik check) + ik_valid[i * N_SOL + j] = 1 + for k in range(DOF): + if not is_within_limit( + qpos[qpos_start + k], + lower_limits[k], + upper_limits[k], + safe_margin=safe_margin, + ): + ik_valid[i * N_SOL + j] = 0 + break if t_err > 1e-2 or r_err > 1e-1: ik_valid[i * N_SOL + j] = 0 - else: - ik_valid[i * N_SOL + j] = 1 @wp.kernel -def opw_best_ik_kernel( - full_ik_result: wp.array(dtype=float), - full_ik_valid: wp.array(dtype=int), - qpos_seed: wp.array(dtype=float), +def opw_ik_select_kernel( + full_ik_result: wp.array(dtype=float, ndim=3), # [n_sample, N_SOL, DOF] + full_ik_valid: wp.array(dtype=int, ndim=2), # [n_sample, N_SOL] + qpos_seed: wp.array(dtype=float, ndim=2), # [n_sample, DOF] joint_weights: wp_vec6f, - best_ik_result: wp.array(dtype=float), - best_ik_valid: wp.array(dtype=int), + best_ik_result: wp.array(dtype=float, ndim=2), # [n_sample, DOF] + best_ik_valid: wp.array(dtype=int, ndim=1), # [n_sample, ] ): - i = wp.tid() - DOF = 6 - N_SOL = 8 - + i = wp.tid() # index for sample best_weighted_dis = float(1e10) best_ids = int(-1) + DOF = 6 + N_SOL = 8 for j in range(N_SOL): - is_full_valid = full_ik_valid[i * N_SOL + j] + is_full_valid = full_ik_valid[i, j] if is_full_valid == 0: # invalid ik result continue weighted_dis = 0.0 for t in range(DOF): weighted_dis += ( - (full_ik_result[i * N_SOL * DOF + j * DOF + t] - qpos_seed[i * DOF + t]) - * joint_weights[0] - * ( - full_ik_result[i * N_SOL * DOF + j * DOF + t] - - qpos_seed[i * DOF + t] - ) - * joint_weights[0] + (full_ik_result[i, j, t] - qpos_seed[i, t]) + * joint_weights[t] + * (full_ik_result[i, j, t] - qpos_seed[i, t]) + * joint_weights[t] ) if weighted_dis < best_weighted_dis: best_weighted_dis = weighted_dis @@ -493,9 +519,7 @@ def opw_best_ik_kernel( # found best solution best_ik_valid[i] = 1 for k in range(DOF): - best_ik_result[i * DOF + k] = full_ik_result[ - i * N_SOL * DOF + best_ids * DOF + k - ] + best_ik_result[i, k] = full_ik_result[i, best_ids, k] else: # no valid solution best_ik_valid[i] = 0 diff --git a/scripts/benchmark/opw_solver.py b/scripts/benchmark/opw_solver.py index c248eaba..78f7e3d7 100644 --- a/scripts/benchmark/opw_solver.py +++ b/scripts/benchmark/opw_solver.py @@ -23,6 +23,10 @@ import time +LOWER_LIMITS = [-2.618, 0.0, -2.967, -1.745, -1.22, -2.0944] +UPPER_LIMITS = [2.618, 3.14159, 0.0, 1.745, 1.22, 2.0944] + + def get_pose_err(matrix_a: np.ndarray, matrix_b: np.ndarray) -> Tuple[float, float]: t_err = np.linalg.norm(matrix_a[:3, 3] - matrix_b[:3, 3]) relative_rot = matrix_a[:3, :3].T @ matrix_b[:3, :3] @@ -46,9 +50,13 @@ def get_poses_err( def check_opw_solver(solver_warp, solver_py_opw, n_samples=1000): DOF = 6 - qpos_np = np.random.uniform(low=-np.pi, high=np.pi, size=(n_samples, DOF)).astype( - float - ) + qpos_np = np.random.uniform( + low=np.array(LOWER_LIMITS) + + 5.1 / 180.0 * np.pi, # add a margin to avoid sampling near the joint limits + high=np.array(UPPER_LIMITS) + -5.1 / 180.0 * np.pi, + size=(n_samples, DOF), + ).astype(float) + qpos = torch.tensor(qpos_np, device=torch.device("cuda"), dtype=torch.float32) xpos = solver_warp.get_fk(qpos) qpos_seed = torch.tensor( @@ -108,7 +116,10 @@ def check_opw_solver(solver_warp, solver_py_opw, n_samples=1000): def benchmark_opw_solver(): - cfg = OPWSolverCfg() + cfg = OPWSolverCfg( + joint_names=("J1", "J2", "J3", "J4", "J5", "J6"), + user_qpos_limits=(LOWER_LIMITS, UPPER_LIMITS), + ) cfg.a1 = 400.333 cfg.a2 = -251.449 cfg.b = 0.0 @@ -127,11 +138,11 @@ def benchmark_opw_solver(): cfg.flip_axes = (True, False, True, True, False, True) cfg.has_parallelogram = False - # TODO: ignore pk_serial_chain for OPW + # TODO: Set pk_serial_chain to "" to ignore pk_serial_chain for OPW. solver_warp = cfg.init_solver(device=torch.device("cuda"), pk_serial_chain="") solver_py_opw = cfg.init_solver(device=torch.device("cpu"), pk_serial_chain="") + n_samples = [100, 1000, 10000, 100000] - # n_samples = [100] for n_sample in n_samples: # check_opw_solver(solver_warp, solver_py_opw, device=device, n_samples=n_sample) ( @@ -142,13 +153,13 @@ def benchmark_opw_solver(): py_opw_t_mean_err, py_opw_r_mean_err, ) = check_opw_solver(solver_warp, solver_py_opw, n_samples=n_sample) - print(f"===warp OPW Solver FK/IK test over {n_sample} samples:") - print(f" Warp IK time: {warp_cost_time * 1000:.6f} ms") - print(f"Translation mean error: {warp_t_mean_err*1000:.6f} mm") - print(f"Rotation mean error: {warp_r_mean_err*180/np.pi:.6f} degrees") - print(f"===Py OPW IK time: {py_opw_cost_time * 1000:.6f} ms") - print(f"Translation mean error: {py_opw_t_mean_err*1000:.6f} mm") - print(f"Rotation mean error: {py_opw_r_mean_err*180/np.pi:.6f} degrees") + print(f"*******warp cuda OPW Solver FK/IK test over {n_sample} samples:") + print(f"===Warp IK time: {warp_cost_time * 1000:.6f} ms") + print(f" Translation mean error: {warp_t_mean_err*1000:.6f} mm") + print(f" Rotation mean error: {warp_r_mean_err*180/np.pi:.6f} degrees") + print(f"===warp cpu IK time: {py_opw_cost_time * 1000:.6f} ms") + print(f" Translation mean error: {py_opw_t_mean_err*1000:.6f} mm") + print(f" Rotation mean error: {py_opw_r_mean_err*180/np.pi:.6f} degrees") if __name__ == "__main__": diff --git a/tests/sim/solvers/test_opw_solver.py b/tests/sim/solvers/test_opw_solver.py index fe04f4b4..24b91ae7 100644 --- a/tests/sim/solvers/test_opw_solver.py +++ b/tests/sim/solvers/test_opw_solver.py @@ -28,6 +28,7 @@ def grid_sample_qpos_from_limits( steps_per_joint: int = 4, device=None, max_samples: int = 4096, + safe_margin: float = 5 / 180 * np.pi, # 5 degrees in radians ) -> torch.Tensor: """Generate grid samples for qpos from qpos_limits. @@ -44,8 +45,8 @@ def grid_sample_qpos_from_limits( device = qpos_limits.device limits = qpos_limits.squeeze(0) if qpos_limits.dim() == 3 else qpos_limits - lows = limits[:, 0].to(device) - highs = limits[:, 1].to(device) + lows = limits[:, 0].to(device) + safe_margin * 1.01 + highs = limits[:, 1].to(device) - safe_margin * 1.01 # create per-joint linspaces grids = [ @@ -98,12 +99,20 @@ def setup_simulation(self, sim_device): "end_link_name": "left_link6", "root_link_name": "left_arm_base", "tcp": [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]], + "qpos_limits": [ + [-2.618, 0.0, -2.967, -1.745, -1.22, -2.0944], + [2.618, 3.14159, 0.0, 1.745, 1.22, 2.0944], + ], }, "right_arm": { "class_type": "OPWSolver", "end_link_name": "right_link6", "root_link_name": "right_arm_base", "tcp": [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.143], [0, 0, 0, 1]], + "qpos_limits": [ + [-2.618, 0.0, -2.967, -1.745, -1.22, -2.0944], + [2.618, 3.14159, 0.0, 1.745, 1.22, 2.0944], + ], }, }, } @@ -165,7 +174,7 @@ def test_ik(self, arm_name: str): device=self.robot.device, ) res, ik_qpos = self.robot.compute_ik( - pose=invalid_pose, joint_seed=ik_qpos, name=arm_name + pose=invalid_pose, joint_seed=ik_qpos[:, 0, :], name=arm_name ) dof = ik_qpos.shape[-1] assert res[0] == False diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index a4a375ed..ddb24120 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -73,7 +73,7 @@ def setup_solver(self, solver_type: str, device: str = "cpu"): ) cfg.urdf_path = urdf cfg.dh_params = arm_params.dh_params - cfg.qpos_limits = arm_params.qpos_limits + cfg.user_qpos_limits = arm_params.qpos_limits cfg.T_e_oe = arm_params.T_e_oe cfg.T_b_ob = arm_params.T_b_ob cfg.link_lengths = arm_params.link_lengths