diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index dba811321..49e4a9620 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -151,8 +151,39 @@ class RobotCfg(ArticulationCfg): - `dh_params`, `link_lengths`, `rotation_directions`, `T_b_ob`, `T_e_oe`: kinematic model params. - `sort_ik`: whether to rank solutions by distance to seed. +- `search_mode`: `"seeded"` computes the seed's geometric shoulder-elbow-wrist + arm angle and searches redundancy angles radially around it; if the configured + radial step cannot produce `num_samples` distinct angles in one revolution, + the incomplete radial prefix is replaced by a complete seed-centered uniform + full-circle grid. `"full"` samples the complete `[-pi, pi)` interval directly. +- `redundancy_step`: angular increment used by seed-centered search. +- Requesting all solutions always uses full-space redundancy sampling. +- CPU and CUDA derive the reference plane in the base frame and use the same + signed arm-angle, shoulder-azimuth degeneracy rule, and periodic + nearest-solution formulas. Shoulder azimuth is set to zero only when the + shoulder-to-wrist projection onto the XY plane is near zero. +- At shoulder or wrist Euler singularities, both analytical backends preserve + the seed's free coupled joint and solve the remaining coupled angle, avoiding + arbitrary equivalent-angle jumps near singular configurations. +- Candidate revolute angles are shifted by integer multiples of `2*pi` into + the configured joint limits, choosing the representation nearest the seed. +- Runtime `set_tcp()` and `set_ik_nearest_weight()` calls synchronize the CPU + and Warp analytical-backend caches immediately. +- CPU target/reference-plane geometry is precomputed per target and elbow + branch. CUDA derives target/config/angle indices directly from the Warp + thread id, and all-solution sorting uses device-side tensor sorting rather + than a serial quadratic Warp sort. +- Warp arm-angle and IK scratch arrays are reused by shape within a solver + instance to avoid repeated device allocations during steady-state calls. +- Periodic-equivalent all-solutions candidates are greedily deduplicated + against retained representatives on CPU before indexing the original device + tensor, preserving order without allocating quadratic GPU scratch space. - Requires `num_envs` in `init_solver()`. +Focused performance and accuracy validation is available at +`scripts/benchmark/robotics/kinematic_solver/srs_solver.py`; it compares CPU +and available CUDA backends in seeded and full redundancy-search modes. + ### OPWSolver-specific - `a1, a2, b, c1–c4, offsets, flip_axes, has_parallelogram`: OPW kinematic parameters. diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index 23c742fab..aad8cce48 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -16,26 +16,26 @@ from __future__ import annotations -import torch +from itertools import product +from typing import TYPE_CHECKING, Literal + import numpy as np +import torch import warp as wp -from itertools import product -from typing import Union, Tuple, Any, Literal, TYPE_CHECKING -from embodichain.utils import configclass, logger -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.solvers import BaseSolver, SolverCfg +from embodichain.utils import configclass, logger +from embodichain.utils.device_utils import standardize_device_string from embodichain.utils.warp.kinematics.srs_solver import ( - transform_pose_kernel, + check_success_kernel, + compute_arm_angle_kernel, compute_ik_kernel, - sort_ik_kernel, nearest_ik_kernel, - check_success_kernel, + transform_pose_kernel, ) -from embodichain.utils.device_utils import standardize_device_string if TYPE_CHECKING: - from typing import Self - from embodichain.lab.sim.robots.dexforce_w1.params import W1ArmKineParams + pass __all__ = ["SRSSolver", "SRSSolverCfg"] @@ -68,6 +68,16 @@ class SRSSolverCfg(SolverCfg): num_samples: int = 100 """Number of samples for elbow angle during IK computation.""" + search_mode: Literal["seeded", "full"] = "seeded" + """Redundancy search strategy. + + ``"seeded"`` searches the seed arm angle first and then expands radially; + ``"full"`` samples the complete ``[-pi, pi)`` interval. + """ + + redundancy_step: float = np.pi / 18.0 + """Angular step in radians for seed-centered redundancy search.""" + sort_ik: bool = True """Whether to sort IK solutions based on proximity to seed joint positions.""" @@ -122,6 +132,124 @@ def _parse_params(self): self.link_lengths_np = np.asarray(self.cfg.link_lengths) self.rotation_directions_np = np.asarray(self.cfg.rotation_directions) + if self.cfg.num_samples < 1: + raise ValueError("num_samples must be at least 1") + if self.cfg.search_mode not in ("seeded", "full"): + raise ValueError("search_mode must be 'seeded' or 'full'") + if ( + not np.isfinite(self.cfg.redundancy_step) + or self.cfg.redundancy_step <= 0 + or self.cfg.redundancy_step > np.pi + ): + raise ValueError("redundancy_step must be finite and in the range (0, pi]") + + def _sample_elbow_angles( + self, + qpos_seed: torch.Tensor, + *, + force_full: bool = False, + ) -> torch.Tensor: + """Build redundancy samples for every target. + + Seeded sampling follows a radial order ``0, +step, -step, ...`` so a + nearby redundancy branch is evaluated before distant branches. + """ + batch_size = qpos_seed.shape[0] + if force_full or self.cfg.search_mode == "full": + samples = torch.arange( + self.cfg.num_samples, + dtype=qpos_seed.dtype, + device=self.device, + ) + samples = -torch.pi + samples * (2.0 * torch.pi / self.cfg.num_samples) + return samples.unsqueeze(0).expand(batch_size, -1).contiguous() + + offsets = [0.0] + layer = 1 + while len(offsets) < self.cfg.num_samples: + offset = layer * self.cfg.redundancy_step + if offset > np.pi + 1e-12: + break + offsets.append(offset) + # +pi and -pi normalize to the same angle, so keep only one. + if offset < np.pi - 1e-12 and len(offsets) < self.cfg.num_samples: + offsets.append(-offset) + layer += 1 + + # A fixed radial step has only finitely many distinct positions on one + # revolution. If it cannot provide num_samples entries, replace the + # incomplete prefix with a complete seed-centered uniform grid. Mixing + # off-grid radial offsets into a fixed-size grid would necessarily omit + # some grid points and leave gaps in the redundancy search space. + if len(offsets) < self.cfg.num_samples: + offsets = [0.0] + uniform_step = 2.0 * np.pi / self.cfg.num_samples + layer = 1 + while len(offsets) < self.cfg.num_samples: + offset = layer * uniform_step + offsets.append(offset) + if len(offsets) < self.cfg.num_samples and offset < np.pi - 1e-12: + offsets.append(-offset) + layer += 1 + + offset_tensor = torch.tensor(offsets, dtype=qpos_seed.dtype, device=self.device) + seed_arm_angles = self._get_seed_arm_angles(qpos_seed) + angles = seed_arm_angles.unsqueeze(1) + offset_tensor.unsqueeze(0) + return torch.remainder(angles + torch.pi, 2.0 * torch.pi) - torch.pi + + def _get_seed_arm_angles(self, qpos_seed: torch.Tensor) -> torch.Tensor: + """Compute geometric arm angles for seed joint configurations.""" + raise NotImplementedError + + @staticmethod + def _wrap_to_limits( + joints: np.ndarray, + limits: np.ndarray, + seed: np.ndarray, + ) -> np.ndarray | None: + """Map revolute joints to equivalent values inside limits near the seed.""" + wrapped = np.empty_like(joints) + two_pi = 2.0 * np.pi + for index, value in enumerate(joints): + lower, upper = limits[index] + k_min = int(np.ceil((lower - value) / two_pi)) + k_max = int(np.floor((upper - value) / two_pi)) + if k_min > k_max: + return None + nearest_k = int(np.floor((seed[index] - value) / two_pi + 0.5)) + nearest_k = min(max(nearest_k, k_min), k_max) + wrapped[index] = value + nearest_k * two_pi + return wrapped + + @staticmethod + def _deduplicate_solutions( + solutions: torch.Tensor, tolerance: float = 1e-5 + ) -> torch.Tensor: + """Greedily retain periodic-unique representatives in input order.""" + if solutions.shape[0] < 2: + return solutions + + cpu_solutions = solutions.detach().cpu().numpy() + retained = np.empty_like(cpu_solutions) + retained_indices = np.empty(cpu_solutions.shape[0], dtype=np.int64) + retained_count = 0 + for index, candidate in enumerate(cpu_solutions): + if retained_count: + difference = candidate - retained[:retained_count] + wrapped = np.remainder(difference + np.pi, 2.0 * np.pi) - np.pi + if np.any(np.max(np.abs(wrapped), axis=1) <= tolerance): + continue + retained[retained_count] = candidate + retained_indices[retained_count] = index + retained_count += 1 + + index_tensor = torch.as_tensor( + retained_indices[:retained_count], + dtype=torch.long, + device=solutions.device, + ) + return solutions[index_tensor] + class _CPUSRSSolverImpl(_BaseSRSSolverImpl): """CPU implementation of the SRS inverse kinematics solver.""" @@ -139,13 +267,6 @@ def _parse_params(self): np.array([x, y, z]) for x, y, z in product([1.0, -1.0], repeat=3) ] - # Generate a set of elbow angles sampled uniformly from -π to π. - # The number of samples is determined by self.cfg.num_samples. - # These angles are used for searching possible IK solutions. - self.elbow_angles = torch.linspace( - -torch.pi, torch.pi, self.cfg.num_samples, device=self.device - ) - # Convert ik_nearest_weight to a tensor for efficient computation. self.ik_nearest_weight_tensor = torch.tensor( self.cfg.ik_nearest_weight, dtype=torch.float32, device=self.device @@ -189,6 +310,70 @@ def _get_fk(self, target_joint: np.ndarray) -> np.ndarray: return pose + def _get_model_fk( + self, target_joint: np.ndarray, end_joint_index: int = 6 + ) -> np.ndarray: + """Compute DH-model FK without base, end-effector, or TCP transforms.""" + pose = np.eye(4) + for index in range(end_joint_index + 1): + d, alpha, a, theta_offset = self.dh_params[index] + theta = ( + theta_offset + target_joint[index] * self.rotation_directions_np[index] + ) + pose = pose @ self._dh_transform(d, alpha, a, theta) + return pose + + def _get_seed_arm_angles(self, qpos_seed: torch.Tensor) -> torch.Tensor: + """Compute seed arm angles from shoulder-elbow-wrist geometry.""" + arm_angles = np.zeros(qpos_seed.shape[0], dtype=np.float32) + for target_index, seed in enumerate(qpos_seed.detach().cpu().numpy()): + full_pose = self._get_model_fk(seed) + elbow_pose = self._get_model_fk(seed, end_joint_index=2) + shoulder = np.array([0.0, 0.0, self.link_lengths_np[0]]) + wrist_offset = np.array([0.0, 0.0, self.dh_params_np[6, 0]]) + wrist = full_pose[:3, 3] - full_pose[:3, :3] @ wrist_offset + shoulder_to_wrist = wrist - shoulder + distance = np.linalg.norm(shoulder_to_wrist) + if distance < 1e-10: + continue + + elbow_model = ( + self.dh_params_np[3, 3] + seed[3] * self.rotation_directions_np[3] + ) + elbow_config = -1.0 if elbow_model < 0.0 else 1.0 + _, _, reference_joints = self._compute_reference_plane( + full_pose, elbow_config + ) + if reference_joints is None: + continue + + reference_pose = np.eye(4) + for index in range(3): + reference_pose = reference_pose @ self._dh_transform( + self.dh_params_np[index, 0], + self.dh_params_np[index, 1], + self.dh_params_np[index, 2], + reference_joints[index], + ) + + axis = shoulder_to_wrist / distance + reference_upper = reference_pose[:3, 3] - shoulder + actual_upper = elbow_pose[:3, 3] - shoulder + reference_radial = reference_upper - axis * np.dot(reference_upper, axis) + actual_radial = actual_upper - axis * np.dot(actual_upper, axis) + reference_norm = np.linalg.norm(reference_radial) + actual_norm = np.linalg.norm(actual_radial) + if reference_norm < 1e-10 or actual_norm < 1e-10: + continue + + reference_radial /= reference_norm + actual_radial /= actual_norm + arm_angles[target_index] = np.arctan2( + np.dot(axis, np.cross(reference_radial, actual_radial)), + np.dot(reference_radial, actual_radial), + ) + return torch.from_numpy(arm_angles).to(self.device) + def _calculate_arm_joint_angles( self, P26: np.ndarray, @@ -220,15 +405,16 @@ def _calculate_arm_joint_angles( logger.log_debug("Elbow singularity. End effector at limit.") return False - joints[3] = elbow_config * np.arccos(elbow_cos_angle) + joints[3] = elbow_config * np.arccos(np.clip(elbow_cos_angle, -1.0, 1.0)) - if abs(P26[2]) > 1e-6: + euclidean_norm = np.hypot(P26[0], P26[1]) + if euclidean_norm > 1e-6: joints[0] = np.arctan2(P26[1], P26[0]) else: joints[0] = 0 - euclidean_norm = np.hypot(P26[0], P26[1]) - angle_phi = np.arccos((d_se**2 + norm_P26**2 - d_ew**2) / (2 * d_se * norm_P26)) + angle_phi_cos = (d_se**2 + norm_P26**2 - d_ew**2) / (2 * d_se * norm_P26) + angle_phi = np.arccos(np.clip(angle_phi_cos, -1.0, 1.0)) joints[1] = np.arctan2(euclidean_norm, P26[2]) + elbow_config * angle_phi return True @@ -308,27 +494,28 @@ def _compute_reference_plane( ): return None, None, None - T34_v = self._dh_transform( - dh_params[3, 0], dh_params[3, 1], dh_params[3, 2], joint_angles[3] - ) - P34_v = T34_v[:3, 3] - - norm_P34_P02 = np.linalg.norm(P34_v - P02) - if norm_P34_P02 > 1e-6: - v1 = (P34_v - P02) / norm_P34_P02 - else: - v1 = np.zeros_like(P34_v - P02) - v2 = (P06 - P02) / np.linalg.norm(P06 - P02) - plane_normal = np.cross(v1, v2) - - base_to_elbow_rotation = np.eye(3) + base_to_elbow_pose = np.eye(4) for i in range(3): T = self._dh_transform( dh_params[i, 0], dh_params[i, 1], dh_params[i, 2], joint_angles[i] ) - base_to_elbow_rotation = base_to_elbow_rotation @ T[:3, :3] + base_to_elbow_pose = base_to_elbow_pose @ T - return plane_normal, base_to_elbow_rotation, joint_angles + reference_upper = base_to_elbow_pose[:3, 3] - P02 + shoulder_to_wrist = P06 - P02 + upper_norm = np.linalg.norm(reference_upper) + wrist_norm = np.linalg.norm(shoulder_to_wrist) + if upper_norm < 1e-10 or wrist_norm < 1e-10: + return None, None, None + plane_normal = np.cross( + reference_upper / upper_norm, shoulder_to_wrist / wrist_norm + ) + plane_norm = np.linalg.norm(plane_normal) + if plane_norm < 1e-10: + return None, None, None + plane_normal /= plane_norm + + return plane_normal, base_to_elbow_pose[:3, :3], joint_angles def _process_all_solutions( self, @@ -336,7 +523,7 @@ def _process_all_solutions( qpos_seed: torch.Tensor, valid_mask: torch.Tensor, success_tensor: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """Returns all valid IK solutions (optionally sorted). Args: @@ -350,18 +537,33 @@ def _process_all_solutions( torch.Tensor: The IK solutions tensor (sorted if specified). """ if self.cfg.sort_ik: - weighted_diff = ( - ik_qpos_tensor - qpos_seed.unsqueeze(1) - ) * self.ik_nearest_weight_tensor - distances = torch.norm(weighted_diff, dim=2) + diff = ik_qpos_tensor - qpos_seed.unsqueeze(1) + wrapped_diff = torch.atan2(torch.sin(diff), torch.cos(diff)) + distances = torch.sum( + wrapped_diff.square() * self.ik_nearest_weight_tensor, dim=2 + ) distances[~valid_mask] = float("inf") sorted_indices = torch.argsort(distances, dim=1) sorted_ik_qpos_tensor = torch.gather( ik_qpos_tensor, 1, sorted_indices.unsqueeze(-1).expand(-1, -1, 7) ) - return success_tensor, sorted_ik_qpos_tensor - else: - return success_tensor, ik_qpos_tensor + sorted_valid_mask = torch.gather(valid_mask, 1, sorted_indices) + ik_qpos_tensor = sorted_ik_qpos_tensor + valid_mask = sorted_valid_mask + valid_qpos = [ + self._deduplicate_solutions(ik_qpos_tensor[index][valid_mask[index]]) + for index in range(ik_qpos_tensor.shape[0]) + ] + max_solutions = max(solution.shape[0] for solution in valid_qpos) + compact_qpos = torch.full( + (ik_qpos_tensor.shape[0], max_solutions, 7), + float("nan"), + dtype=ik_qpos_tensor.dtype, + device=self.device, + ) + for index, solution in enumerate(valid_qpos): + compact_qpos[index, : solution.shape[0]] = solution + return success_tensor, compact_qpos def _process_single_solution( self, @@ -369,7 +571,7 @@ def _process_single_solution( qpos_seed: torch.Tensor, valid_mask: torch.Tensor, success_tensor: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """Returns the nearest valid IK solution (optionally sorted). Args: @@ -384,10 +586,11 @@ def _process_single_solution( """ num_targets = ik_qpos_tensor.shape[0] if self.cfg.sort_ik: - weighted_diff = ( - ik_qpos_tensor - qpos_seed.unsqueeze(1) - ) * self.ik_nearest_weight_tensor - distances = torch.norm(weighted_diff, dim=2) + diff = ik_qpos_tensor - qpos_seed.unsqueeze(1) + wrapped_diff = torch.atan2(torch.sin(diff), torch.cos(diff)) + distances = torch.sum( + wrapped_diff.square() * self.ik_nearest_weight_tensor, dim=2 + ) mask = success_tensor.unsqueeze(1) & valid_mask distances[~mask] = float("inf") nearest_indices = torch.argmin(distances, dim=1) @@ -410,6 +613,10 @@ def _get_each_ik( target_pose: np.ndarray | torch.Tensor, nsparam: float, config: np.ndarray, + qpos_seed: np.ndarray, + prepared: ( + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None + ) = None, ) -> tuple[bool, np.ndarray | None]: """ Computes the inverse kinematics for a given target pose, normalization parameter, and configuration. @@ -446,23 +653,20 @@ def _get_each_ik( link_lengths = self.cfg.link_lengths rotation_directions = self.cfg.rotation_directions - # Transform target pose - target_xpos = ( - self.T_b_ob_inv_np @ target_pose @ self.tcp_inv_np @ self.T_e_oe_inv_np - ) - P_target = target_xpos[:3, 3] - R_target = target_xpos[:3, :3] - P02 = np.array([0, 0, link_lengths[0]]) # Base to shoulder - P67 = np.array([0, 0, dh_params[6, 0]]) # Hand to end-effector - P06 = P_target - R_target @ P67 - P26 = P06 - P02 - - # Calculate joint angles - joints = np.zeros(dof) - if not self._calculate_arm_joint_angles( - P26, elbow_config, joints, link_lengths - ): - return False, None + if prepared is None: + target_xpos = ( + self.T_b_ob_inv_np @ target_pose @ self.tcp_inv_np @ self.T_e_oe_inv_np + ) + P_target = target_xpos[:3, 3] + R_target = target_xpos[:3, :3] + P02 = np.array([0, 0, link_lengths[0]]) + P67 = np.array([0, 0, dh_params[6, 0]]) + P26 = P_target - R_target @ P67 - P02 + _, R03_o, joints = self._compute_reference_plane(target_xpos, elbow_config) + if R03_o is None or joints is None: + return False, None + else: + target_xpos, R_target, P26, R03_o, joints = prepared # Calculate transformations T34 = self._dh_transform( @@ -470,19 +674,12 @@ def _get_each_ik( ) R34 = T34[:3, :3] - # Calculate reference plane - V_v_to_sew, R03_o, joint_v = self._compute_reference_plane( - target_xpos, config[1] - ) - if V_v_to_sew is None: - return False, None - # Calculate shoulder joint rotation matrices usw = P26 / np.linalg.norm(P26) skew_usw = self._skew(usw) angle_psi = nsparam - s_psi = wp.sin(angle_psi) - c_psi = wp.cos(angle_psi) + s_psi = np.sin(angle_psi) + c_psi = np.cos(angle_psi) # Calculate rotation matrix R03 A_s = skew_usw @ R03_o @@ -491,9 +688,17 @@ def _get_each_ik( R03 = A_s * s_psi + B_s * c_psi + C_s # Calculate shoulder joint angles - angle1 = np.arctan2(R03[1, 1] * shoulder_config, R03[0, 1] * shoulder_config) - angle2 = np.arccos(R03[2, 1]) * shoulder_config - angle3 = np.arctan2(-R03[2, 2] * shoulder_config, -R03[2, 0] * shoulder_config) + angle2 = np.arccos(np.clip(R03[2, 1], -1.0, 1.0)) * shoulder_config + if abs(np.sin(angle2)) <= 1e-6: + angle1 = qpos_seed[0] * rotation_directions[0] + dh_params[0, 3] + angle3 = np.arctan2(R03[1, 0], R03[0, 0]) - angle1 + else: + angle1 = np.arctan2( + R03[1, 1] * shoulder_config, R03[0, 1] * shoulder_config + ) + angle3 = np.arctan2( + -R03[2, 2] * shoulder_config, -R03[2, 0] * shoulder_config + ) # Calculate wrist joint angles A_w = R34.T @ A_s.T @ R_target @@ -501,9 +706,13 @@ def _get_each_ik( C_w = R34.T @ C_s.T @ R_target R47 = A_w * s_psi + B_w * c_psi + C_w - angle5 = np.arctan2(R47[1, 2] * wrist_config, R47[0, 2] * wrist_config) - angle6 = np.arccos(R47[2, 2]) * wrist_config - angle7 = np.arctan2(R47[2, 1] * wrist_config, -R47[2, 0] * wrist_config) + angle6 = np.arccos(np.clip(R47[2, 2], -1.0, 1.0)) * wrist_config + if abs(np.sin(angle6)) <= 1e-6: + angle5 = qpos_seed[4] * rotation_directions[4] + dh_params[4, 3] + angle7 = np.arctan2(-R47[2, 0], R47[0, 0]) - angle5 + else: + angle5 = np.arctan2(R47[1, 2] * wrist_config, R47[0, 2] * wrist_config) + angle7 = np.arctan2(R47[2, 1] * wrist_config, -R47[2, 0] * wrist_config) joints_output[0] = (angle1 - dh_params[0, 3]) * rotation_directions[0] joints_output[1] = (angle2 - dh_params[1, 3]) * rotation_directions[1] @@ -513,12 +722,10 @@ def _get_each_ik( joints_output[5] = (angle6 - dh_params[5, 3]) * rotation_directions[5] joints_output[6] = (angle7 - dh_params[6, 3]) * rotation_directions[6] - # Check if the calculated joint angles are within the limits - in_range = (joints_output >= self.qpos_limits_np[:, 0]) & ( - joints_output <= self.qpos_limits_np[:, 1] + joints_output = self._wrap_to_limits( + joints_output, self.qpos_limits_np, qpos_seed ) - - if not np.all(in_range): + if joints_output is None: return False, None return True, joints_output @@ -529,7 +736,7 @@ def get_ik( qpos_seed: torch.Tensor, return_all_solutions: bool = False, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """ Compute inverse kinematics (IK) for the given target pose using CPU. @@ -542,29 +749,71 @@ def get_ik( Returns: Tuple[torch.Tensor, torch.Tensor]: Success flag and joint positions. """ + target_xpos = target_xpos.to(self.device, dtype=torch.float32).view(-1, 4, 4) num_targets = target_xpos.shape[0] # Validate and normalize qpos_seed if qpos_seed is None: qpos_seed = torch.zeros( (target_xpos.shape[0], 7), dtype=torch.float32, device=self.device ) + else: + qpos_seed = ( + qpos_seed.to(self.device, dtype=torch.float32) + .reshape(num_targets, -1, 7)[:, 0] + .contiguous() + ) # Prepare to collect results - max_possible_solutions = len(self.elbow_angles) * len(self.configs) + elbow_angles = self._sample_elbow_angles( + qpos_seed, force_full=return_all_solutions + ).cpu() + max_possible_solutions = elbow_angles.shape[1] * len(self.configs) all_solutions = np.zeros( (num_targets, max_possible_solutions, 7), dtype=np.float32 ) solution_counts = np.zeros(num_targets, dtype=np.int32) + qpos_seed_np = qpos_seed.detach().cpu().numpy() + target_xpos_np = target_xpos.detach().cpu().numpy() # Iterate over target poses - for target_idx, xpos in enumerate(target_xpos): + for target_idx in range(num_targets): + target_np = target_xpos_np[target_idx] + transformed = ( + self.T_b_ob_inv_np @ target_np @ self.tcp_inv_np @ self.T_e_oe_inv_np + ) + rotation = transformed[:3, :3] + shoulder = np.array([0.0, 0.0, self.link_lengths_np[0]]) + wrist_offset = np.array([0.0, 0.0, self.dh_params_np[6, 0]]) + shoulder_to_wrist = transformed[:3, 3] - rotation @ wrist_offset - shoulder + prepared_by_elbow = {} + for elbow_config in (1.0, -1.0): + _, reference_rotation, reference_joints = self._compute_reference_plane( + transformed, elbow_config + ) + if reference_rotation is not None and reference_joints is not None: + prepared_by_elbow[elbow_config] = ( + transformed, + rotation, + shoulder_to_wrist, + reference_rotation, + reference_joints, + ) sol_idx = 0 - for psi in self.elbow_angles: + for psi in elbow_angles[target_idx]: for config in self.configs: - success, qpos = self._get_each_ik(xpos, psi.item(), config) + prepared = prepared_by_elbow.get(config[1]) + if prepared is None: + continue + success, qpos = self._get_each_ik( + target_np, + psi.item(), + config, + qpos_seed_np[target_idx], + prepared, + ) if success: fk_xpos = self._get_fk(qpos) - if np.allclose(fk_xpos, xpos, atol=1e-4): + if np.allclose(fk_xpos, target_np, atol=1e-4, rtol=0.0): all_solutions[target_idx, sol_idx, :] = qpos sol_idx += 1 solution_counts[target_idx] = sol_idx @@ -578,7 +827,7 @@ def get_ik( return ( torch.zeros(num_targets, dtype=torch.bool, device=self.device), torch.zeros( - (num_targets, 7), + (num_targets, 0 if return_all_solutions else 1, 7), dtype=qpos_seed.dtype, device=self.device, ), @@ -598,7 +847,9 @@ def get_ik( all_solutions[target_idx, :count] ).to(self.device, dtype=qpos_seed.dtype) - valid_mask = ik_qpos_tensor.abs().sum(dim=2) > 0 # (num_targets, max_solutions) + valid_mask = torch.arange(max_solutions, device=self.device).unsqueeze( + 0 + ) < torch.from_numpy(solution_counts).to(self.device).unsqueeze(1) success_tensor = torch.from_numpy(has_solution).to(self.device) if return_all_solutions: return self._process_all_solutions( @@ -649,78 +900,48 @@ def _parse_params(self): self.configs, dtype=wp.vec3, device=standardize_device_string(self.device) ) - # Generate a set of elbow angles sampled uniformly from -π to π. - # The number of samples is determined by self.cfg.num_samples. - # These angles are used for searching possible IK solutions. - joint_reference_limits = [-wp.pi, wp.pi] - self.elbow_angles = np.linspace( - joint_reference_limits[0], joint_reference_limits[1], self.cfg.num_samples - ).tolist() - - # Convert elbow angles to Warp array for CUDA computation. - self.elbow_angles_wp = wp.array( - self.elbow_angles, + self.ik_nearest_weight_wp = wp.array( + self.cfg.ik_nearest_weight, dtype=float, device=standardize_device_string(self.device), ) - - def _sort_ik_solutions( - self, qpos_out_wp, success_wp, qpos_seed, num_targets, num_configs, num_angles - ): - """ - Sort IK solutions based on weighted distance. - - Args: - qpos_out_wp: Warp array of IK solutions (shape: [num_targets * num_configs * num_angles, 7]). - success_wp: Warp array of validity flags (shape: [num_targets * num_configs * num_angles]). - qpos_seed: Warp array of seed positions (shape: [num_targets, 7]). - num_targets: Number of targets. - num_configs: Number of configurations. - num_angles: Number of angles. - - Returns: - Tuple[wp.array, wp.array]: Sorted IK solutions and their validity flags. - """ - N = num_targets - N_SOL = num_configs * num_angles - DOF = 7 - - sorted_ik_solutions = wp.zeros( - N * N_SOL * DOF, dtype=float, device=standardize_device_string(self.device) - ) - sorted_ik_valid_flags = wp.zeros( - N * N_SOL, dtype=int, device=standardize_device_string(self.device) - ) - distances = wp.zeros( - N * N_SOL, dtype=float, device=standardize_device_string(self.device) - ) - indices = wp.zeros( - N * N_SOL, dtype=int, device=standardize_device_string(self.device) - ) - + self._temporary_workspace: dict[tuple[int, str], wp.array] = {} + + def _temporary_array(self, count: int, dtype: type, name: str) -> wp.array: + """Return a zeroed reusable Warp scratch array.""" + key = (count, name) + array = self._temporary_workspace.get(key) + if array is None: + array = wp.zeros( + count, + dtype=dtype, + device=standardize_device_string(self.device), + ) + self._temporary_workspace[key] = array + else: + array.zero_() + return array + + def _get_seed_arm_angles(self, qpos_seed: torch.Tensor) -> torch.Tensor: + """Compute seed arm angles with the Warp geometric implementation.""" + batch_size = qpos_seed.shape[0] + arm_angles_wp = self._temporary_array(batch_size, float, "arm_angles") + success_wp = self._temporary_array(batch_size, int, "arm_angle_success") wp.launch( - kernel=sort_ik_kernel, - dim=num_targets, + kernel=compute_arm_angle_kernel, + dim=batch_size, inputs=[ - qpos_out_wp, - success_wp, - qpos_seed, - wp.array( - self.cfg.ik_nearest_weight, - dtype=float, - device=standardize_device_string(self.device), - ), - distances, - indices, - N_SOL, - ], - outputs=[ - sorted_ik_solutions, - sorted_ik_valid_flags, + wp.from_torch(qpos_seed.contiguous().flatten()), + self.dh_params_wp, + self.link_lengths_wp, + self.rotation_directions_wp, ], + outputs=[arm_angles_wp, success_wp], device=standardize_device_string(self.device), ) - return sorted_ik_solutions, sorted_ik_valid_flags + arm_angles = wp.to_torch(arm_angles_wp) + success = wp.to_torch(success_wp).bool() + return torch.where(success, arm_angles, torch.zeros_like(arm_angles)) def _nearest_ik_solution( self, qpos_out_wp, success_wp, qpos_seed, num_targets, num_configs, num_angles @@ -761,11 +982,7 @@ def _nearest_ik_solution( qpos_out_wp, success_wp, qpos_seed.flatten(), - wp.array( - self.cfg.ik_nearest_weight, - dtype=float, - device=standardize_device_string(self.device), - ), + self.ik_nearest_weight_wp, N_SOL, ], outputs=[ @@ -784,7 +1001,7 @@ def _process_all_solutions( num_targets: int, num_configs: int, num_angles: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """ Process and return all valid IK solutions. @@ -801,40 +1018,40 @@ def _process_all_solutions( """ num_per_target = num_configs * num_angles + ik_solutions_tensor = wp.to_torch(qpos_out_wp).view( + num_targets, num_per_target, 7 + ) + ik_valid_flags_tensor = ( + wp.to_torch(success_wp).view(num_targets, num_per_target).bool() + ) if self.cfg.sort_ik: - sorted_ik_solutions, sorted_ik_valid_flags = self._sort_ik_solutions( - qpos_out_wp, - success_wp, - qpos_seed.flatten(), - num_targets, - num_configs, - num_angles, + diff = ik_solutions_tensor - qpos_seed.unsqueeze(1) + wrapped_diff = torch.atan2(torch.sin(diff), torch.cos(diff)) + weights = torch.as_tensor( + self.cfg.ik_nearest_weight, + dtype=ik_solutions_tensor.dtype, + device=self.device, ) - - ik_solutions_tensor = wp.to_torch(sorted_ik_solutions).view( - num_targets, num_per_target, 7 - ) - ik_valid_flags_tensor = ( - wp.to_torch(sorted_ik_valid_flags) - .view(num_targets, num_per_target) - .bool() - ) - else: - ik_solutions_tensor = wp.to_torch(qpos_out_wp).view( - num_targets, num_per_target, 7 - ) - ik_valid_flags_tensor = ( - wp.to_torch(success_wp).view(num_targets, num_per_target).bool() + distances = torch.sum(wrapped_diff.square() * weights, dim=-1) + distances.masked_fill_(~ik_valid_flags_tensor, float("inf")) + indices = torch.argsort(distances, dim=1) + ik_solutions_tensor = torch.gather( + ik_solutions_tensor, 1, indices.unsqueeze(-1).expand(-1, -1, 7) ) + ik_valid_flags_tensor = torch.gather(ik_valid_flags_tensor, 1, indices) success_flags = ik_valid_flags_tensor.any(dim=1) valid_qpos_list = [ - ik_solutions_tensor[i][ik_valid_flags_tensor[i]] for i in range(num_targets) + self._deduplicate_solutions( + ik_solutions_tensor[i][ik_valid_flags_tensor[i]] + ) + for i in range(num_targets) ] max_solutions = max(q.shape[0] for q in valid_qpos_list) - valid_qpos_tensor = torch.zeros( + valid_qpos_tensor = torch.full( (num_targets, max_solutions, 7), + float("nan"), dtype=torch.float32, device=self.device, ) @@ -851,7 +1068,7 @@ def _process_single_solution( num_targets: int, num_configs: int, num_angles: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """ Process and return the nearest valid IK solution for each target. @@ -956,58 +1173,43 @@ def _check_success_flags( def _compute_ik_solutions( self, - combinations_wp: wp.array, xpos_wp: wp.array, + qpos_seed: torch.Tensor, qpos_out_wp: wp.array, success_wp: wp.array, num_combinations: int, + num_configs: int, + num_angles: int, ) -> None: """ Compute IK solutions using the provided combinations. Args: - combinations_wp: Warp array of combinations for parallel processing. xpos_wp: Transformed target poses. qpos_out_wp: Output array for joint positions. success_wp: Output array for success flags. num_combinations: Total number of combinations to process. """ # Temporary arrays - res_arm_angles = wp.zeros( - num_combinations, dtype=int, device=standardize_device_string(self.device) - ) - joints_arm = wp.zeros( - num_combinations, - dtype=wp.vec4, - device=standardize_device_string(self.device), - ) - res_plane_normal = wp.zeros( - num_combinations, dtype=int, device=standardize_device_string(self.device) - ) - plane_normal = wp.zeros( - num_combinations, - dtype=wp.vec3, - device=standardize_device_string(self.device), - ) - base_to_elbow_rotation = wp.zeros( - num_combinations, - dtype=wp.mat33, - device=standardize_device_string(self.device), + res_arm_angles = self._temporary_array(num_combinations, int, "res_arm_angles") + joints_arm = self._temporary_array(num_combinations, wp.vec4, "joints_arm") + res_plane_normal = self._temporary_array( + num_combinations, int, "res_plane_normal" ) - joints_plane = wp.zeros( - num_combinations, - dtype=wp.vec4, - device=standardize_device_string(self.device), + plane_normal = self._temporary_array(num_combinations, wp.vec3, "plane_normal") + base_to_elbow_rotation = self._temporary_array( + num_combinations, wp.mat33, "base_to_elbow_rotation" ) + joints_plane = self._temporary_array(num_combinations, wp.vec4, "joints_plane") # Launch kernel to compute IK solutions wp.launch( kernel=compute_ik_kernel, dim=num_combinations, inputs=( - combinations_wp, xpos_wp, self.elbow_angles_wp, + wp.from_torch(qpos_seed.contiguous().flatten()), self.qpos_limits_wp, self.configs_wp, self.dh_params_wp, @@ -1019,6 +1221,8 @@ def _compute_ik_solutions( plane_normal, base_to_elbow_rotation, joints_plane, + num_configs, + num_angles, ), outputs=[success_wp, qpos_out_wp], device=standardize_device_string(self.device), @@ -1030,7 +1234,7 @@ def get_ik( qpos_seed: torch.Tensor, return_all_solutions: bool = False, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """ Compute inverse kinematics (IK) for the given target pose. @@ -1043,7 +1247,7 @@ def get_ik( Tuple[torch.Tensor, torch.Tensor]: Success flag and joint positions. """ # Prepare inputs - target_xpos = target_xpos.to(self.device) + target_xpos = target_xpos.to(self.device, dtype=torch.float32) target_xpos = target_xpos.view(-1, 4, 4) target_xpos_wp = wp.from_torch(target_xpos, dtype=wp.mat44) @@ -1068,40 +1272,29 @@ def get_ik( # Define configurations and angles if qpos_seed is None: - qpos_seed = wp.zeros( + qpos_seed = torch.zeros( (target_xpos.shape[0], 7), - dtype=float, - device=standardize_device_string(self.device), + dtype=target_xpos.dtype, + device=self.device, + ) + else: + qpos_seed = ( + qpos_seed.to(self.device, dtype=torch.float32) + .reshape(target_xpos.shape[0], -1, 7)[:, 0] + .contiguous() ) - # TODO: Currently, full-space sampling is used to temporarily address situations - # where joint space discontinuities or solution failures occur in different user scenarios. - # Future plans include reducing the sampling space and adjusting the configuration. - # - # self.configs = [wp.vec3(*np.sign(qpos_seed[[1, 3, 5]].cpu().numpy()))] # Prepare output arrays num_targets = target_xpos_wp.shape[0] num_configs = len(self.configs) - num_angles = len(self.elbow_angles) + elbow_angles = self._sample_elbow_angles( + qpos_seed, force_full=return_all_solutions + ).contiguous() + num_angles = elbow_angles.shape[1] + self.elbow_angles_wp = wp.from_torch(elbow_angles.flatten()) # num_solutions = num_configs * num_angles num_combinations = num_targets * num_configs * num_angles - # Generate combinations for parallel processing - combinations_np = np.stack( - np.meshgrid( - np.arange(num_targets), - np.arange(num_configs), - np.arange(num_angles), - indexing="ij", - ), - axis=-1, - ).reshape(-1, 3) - combinations_wp = wp.array( - combinations_np, - dtype=wp.vec3, - device=standardize_device_string(self.device), - ) - # Output arrays qpos_out_wp = wp.zeros( num_combinations * 7, @@ -1114,7 +1307,13 @@ def get_ik( # Compute IK solutions self._compute_ik_solutions( - combinations_wp, xpos_wp, qpos_out_wp, success_wp, num_combinations + xpos_wp, + qpos_seed, + qpos_out_wp, + success_wp, + num_combinations, + num_configs, + num_angles, ) # Check for successful solutions @@ -1145,7 +1344,7 @@ def get_ik( return ( torch.zeros(num_targets, dtype=torch.bool, device=self.device), torch.zeros( - (num_targets, num_targets, 7), + (num_targets, 0 if return_all_solutions else 1, 7), dtype=torch.float32, device=self.device, ), @@ -1185,7 +1384,7 @@ def __init__(self, cfg: SRSSolverCfg, num_envs: int, device: str, **kwargs): fk_dict = self.pk_serial_chain.forward_kinematics( th=torch.zeros(7, dtype=torch.float32, device=self.device), end_only=False ) - root_tf = fk_dict[list(fk_dict.keys())[0]] + root_tf = fk_dict[next(iter(fk_dict))] self.root_base_xpos = root_tf.get_matrix().cpu().numpy() # Initialize implementation based on device @@ -1205,6 +1404,37 @@ def _update_impl_qpos_limits(self): device=standardize_device_string(self.device), ) + def set_tcp(self, xpos: np.ndarray) -> None: + """Set TCP and synchronize the analytical backend caches.""" + super().set_tcp(xpos) + if hasattr(self, "impl"): + self.impl.tcp_xpos = self.tcp_xpos.copy() + self.impl.tcp_inv_np = np.linalg.inv(self.tcp_xpos) + if isinstance(self.impl, _CUDASRSSolverImpl): + self.impl.tcp_inv_wp = wp.mat44(*self.impl.tcp_inv_np.flatten()) + + def set_ik_nearest_weight( + self, ik_weight: np.ndarray, joint_ids: np.ndarray | None = None + ) -> bool: + """Set nearest-solution weights and synchronize backend caches.""" + success = super().set_ik_nearest_weight(ik_weight, joint_ids) + if not success or not hasattr(self, "impl"): + return success + weights = torch.as_tensor( + self.ik_nearest_weight, dtype=torch.float32, device=self.device + ) + self.cfg.ik_nearest_weight = weights.detach().cpu().numpy().copy() + if isinstance(self.impl, _CPUSRSSolverImpl): + self.impl.ik_nearest_weight_tensor = weights + else: + # ``wp.from_torch`` creates a non-owning view, so retain the Torch + # storage for as long as the CUDA backend may launch kernels with it. + self.impl.ik_nearest_weight_tensor = weights.contiguous() + self.impl.ik_nearest_weight_wp = wp.from_torch( + self.impl.ik_nearest_weight_tensor + ) + return True + def update_with_robot_limit(self, robot_qpos_limits): super().update_with_robot_limit(robot_qpos_limits) self._update_impl_qpos_limits() @@ -1215,7 +1445,7 @@ def get_ik( qpos_seed: torch.Tensor = None, return_all_solutions: bool = False, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """ Compute inverse kinematics (IK) for the given target pose. diff --git a/embodichain/utils/warp/kinematics/srs_solver.py b/embodichain/utils/warp/kinematics/srs_solver.py index d188b12fd..e08d31965 100644 --- a/embodichain/utils/warp/kinematics/srs_solver.py +++ b/embodichain/utils/warp/kinematics/srs_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import warp as wp @@ -42,7 +44,20 @@ def identity_mat33() -> wp.mat33: @wp.func def safe_acos(x: float) -> float: - return wp.acos(wp.clamp(x, -0.999999, 0.999999)) + return wp.acos(wp.clamp(x, -1.0, 1.0)) + + +@wp.func +def wrap_to_limit(value: float, lower: float, upper: float, seed: float) -> wp.vec2: + """Return ``(success, equivalent_value)`` nearest the seed within limits.""" + two_pi = 2.0 * wp.pi + k_min = wp.ceil((lower - value) / two_pi) + k_max = wp.floor((upper - value) / two_pi) + if k_min > k_max: + return wp.vec2(0.0, value) + nearest_k = wp.floor((seed - value) / two_pi + 0.5) + nearest_k = wp.clamp(nearest_k, k_min, k_max) + return wp.vec2(1.0, value + nearest_k * two_pi) @wp.func @@ -190,7 +205,7 @@ def calculate_arm_joint_angles( joints_val[3] = elbow_GC4 * safe_acos(elbow_cos_angle) # Compute shoulder angle - joints_val[0] = wp.atan2(y, x) if wp.abs(z) > 1e-6 else 0.0 + joints_val[0] = wp.atan2(y, x) if horizontal_distance > 1e-6 else 0.0 # Compute joint 2 angle angle_phi = safe_acos( @@ -257,19 +272,8 @@ def compute_reference_plane( joints[tid] = wp.vec4() return - # Lower arm transformation (joint 4) - T34 = dh_transform( - dh_params[3 * 4 + 0], dh_params[3 * 4 + 1], dh_params[3 * 4 + 2], 0.0 - ) - P34 = wp.vec3(T34[0, 3], T34[1, 3], T34[2, 3]) - - # Reference plane normal - v1 = wp.normalize(P34 - P02) - v2 = wp.normalize(P06 - P02) - plane_normal[tid] = wp.cross(v1, v2) - - # Compute base-to-elbow rotation - base_to_elbow_rotation[tid] = identity_mat33() + # Compute the reference shoulder-to-elbow pose in the base frame. + base_to_elbow_pose = identity_mat44() for i in range(3): base_idx = i * 4 T = dh_transform( @@ -278,13 +282,43 @@ def compute_reference_plane( dh_params[base_idx + 2], joints[tid][i], ) - # fmt: off - base_to_elbow_rotation[tid] = base_to_elbow_rotation[tid] @ wp.mat33( - T[0, 0], T[0, 1], T[0, 2], - T[1, 0], T[1, 1], T[1, 2], - T[2, 0], T[2, 1], T[2, 2], - ) - # fmt: on + base_to_elbow_pose = base_to_elbow_pose @ T + + reference_elbow = wp.vec3( + base_to_elbow_pose[0, 3], + base_to_elbow_pose[1, 3], + base_to_elbow_pose[2, 3], + ) + reference_upper = reference_elbow - P02 + shoulder_to_wrist = P06 - P02 + upper_norm = wp.length(reference_upper) + wrist_norm = wp.length(shoulder_to_wrist) + if upper_norm < 1e-10 or wrist_norm < 1e-10: + res[tid] = 0 + plane_normal[tid] = wp.vec3() + base_to_elbow_rotation[tid] = identity_mat33() + return + + normal = wp.cross(reference_upper / upper_norm, shoulder_to_wrist / wrist_norm) + normal_norm = wp.length(normal) + if normal_norm < 1e-10: + res[tid] = 0 + plane_normal[tid] = wp.vec3() + base_to_elbow_rotation[tid] = identity_mat33() + return + + plane_normal[tid] = normal / normal_norm + base_to_elbow_rotation[tid] = wp.mat33( + base_to_elbow_pose[0, 0], + base_to_elbow_pose[0, 1], + base_to_elbow_pose[0, 2], + base_to_elbow_pose[1, 0], + base_to_elbow_pose[1, 1], + base_to_elbow_pose[1, 2], + base_to_elbow_pose[2, 0], + base_to_elbow_pose[2, 1], + base_to_elbow_pose[2, 2], + ) res[tid] = 1 @@ -328,7 +362,7 @@ def compute_fk_kernel( a = dh_params[base_idx + 2] theta = dh_params[base_idx + 3] theta += joint_angles[tid * num_joints + i] * rotation_directions[i] - T = dh_transform(d, alpha, d, theta) + T = dh_transform(d, alpha, a, theta) pose = pose @ T # Apply additional transforms: base, end-effector, TCP @@ -339,22 +373,111 @@ def compute_fk_kernel( success[tid] = 1 -@wp.func -def frobenius_norm(mat: wp.mat44) -> float: - """ - Compute the Frobenius norm of a 4x4 matrix. +@wp.kernel +def compute_arm_angle_kernel( + qpos: wp.array(dtype=float), + dh_params: wp.array(dtype=float), + link_lengths: wp.array(dtype=float), + rotation_directions: wp.array(dtype=float), + arm_angles: wp.array(dtype=float), + success: wp.array(dtype=int), +): + """Compute the geometric SRS arm angle for each joint configuration.""" + tid = wp.tid() + actual_pose = identity_mat44() + actual_elbow = wp.vec3() - Args: - mat (wp.mat44): Input matrix. + for i in range(7): + base_idx = i * 4 + theta = dh_params[base_idx + 3] + qpos[tid * 7 + i] * rotation_directions[i] + actual_pose = actual_pose @ dh_transform( + dh_params[base_idx + 0], + dh_params[base_idx + 1], + dh_params[base_idx + 2], + theta, + ) + if i == 2: + actual_elbow = wp.vec3( + actual_pose[0, 3], actual_pose[1, 3], actual_pose[2, 3] + ) + + target_position = wp.vec3(actual_pose[0, 3], actual_pose[1, 3], actual_pose[2, 3]) + target_rotation = wp.mat33( + actual_pose[0, 0], + actual_pose[0, 1], + actual_pose[0, 2], + actual_pose[1, 0], + actual_pose[1, 1], + actual_pose[1, 2], + actual_pose[2, 0], + actual_pose[2, 1], + actual_pose[2, 2], + ) + shoulder = wp.vec3(0.0, 0.0, link_lengths[0]) + wrist_offset = wp.vec3(0.0, 0.0, dh_params[6 * 4 + 0]) + wrist = target_position - target_rotation @ wrist_offset + shoulder_to_wrist = wrist - shoulder + distance = wp.length(shoulder_to_wrist) + if distance < 1e-10: + arm_angles[tid] = 0.0 + success[tid] = 0 + return - Returns: - float: Frobenius norm of the matrix. - """ - norm = 0.0 - for i in range(4): - for j in range(4): - norm += wp.pow(mat[i, j], 2.0) - return wp.sqrt(norm) + elbow_model = dh_params[3 * 4 + 3] + qpos[tid * 7 + 3] * rotation_directions[3] + elbow_config = -1.0 if elbow_model < 0.0 else 1.0 + shoulder_cosine = ( + wp.pow(link_lengths[1], 2.0) + + wp.pow(distance, 2.0) + - wp.pow(link_lengths[2], 2.0) + ) / (2.0 * link_lengths[1] * distance) + horizontal_distance = wp.length(wp.vec2(shoulder_to_wrist[0], shoulder_to_wrist[1])) + q1_reference = ( + wp.atan2(shoulder_to_wrist[1], shoulder_to_wrist[0]) + if horizontal_distance > 1e-6 + else 0.0 + ) + q2_reference = wp.atan2( + horizontal_distance, + shoulder_to_wrist[2], + ) + elbow_config * safe_acos(shoulder_cosine) + + reference_pose = identity_mat44() + for i in range(3): + base_idx = i * 4 + theta = 0.0 + if i == 0: + theta = q1_reference + elif i == 1: + theta = q2_reference + reference_pose = reference_pose @ dh_transform( + dh_params[base_idx + 0], + dh_params[base_idx + 1], + dh_params[base_idx + 2], + theta, + ) + + reference_elbow = wp.vec3( + reference_pose[0, 3], reference_pose[1, 3], reference_pose[2, 3] + ) + axis = shoulder_to_wrist / distance + reference_upper = reference_elbow - shoulder + actual_upper = actual_elbow - shoulder + reference_radial = reference_upper - axis * wp.dot(reference_upper, axis) + actual_radial = actual_upper - axis * wp.dot(actual_upper, axis) + reference_norm = wp.length(reference_radial) + actual_norm = wp.length(actual_radial) + if reference_norm < 1e-10 or actual_norm < 1e-10: + arm_angles[tid] = 0.0 + success[tid] = 0 + return + + reference_radial = reference_radial / reference_norm + actual_radial = actual_radial / actual_norm + arm_angles[tid] = wp.atan2( + wp.dot(axis, wp.cross(reference_radial, actual_radial)), + wp.dot(reference_radial, actual_radial), + ) + success[tid] = 1 @wp.func @@ -415,20 +538,22 @@ def validate_fk_with_target( T = dh_transform(d, alpha, a, theta) pose = pose @ T - # Compute the Frobenius norm of the difference - pose_diff = pose - target_xpos - pose_error = frobenius_norm(pose_diff) - - # Validate against tolerance - return 1 if pose_error <= tolerance else 0 + # Match NumPy's element-wise ``allclose(..., rtol=0)`` semantics used by + # the CPU backend. A Frobenius threshold would become stricter as the + # number of matrix elements grows. + for row in range(4): + for column in range(4): + if wp.abs(pose[row, column] - target_xpos[row, column]) > tolerance: + return 0 + return 1 # TODO: automatic gradient support @wp.kernel def compute_ik_kernel( - combinations: wp.array(dtype=wp.vec3), target_xpos_list: wp.array(dtype=wp.mat44), angles_list: wp.array(dtype=float), + qpos_seed: wp.array(dtype=float), qpos_limits: wp.array(dtype=wp.vec2), configs: wp.array(dtype=wp.vec3), dh_params: wp.array(dtype=float), @@ -440,6 +565,8 @@ def compute_ik_kernel( plane_normal: wp.array(dtype=wp.vec3), base_to_elbow_rotation: wp.array(dtype=wp.mat33), joints_plane: wp.array(dtype=wp.vec4), + num_configs: int, + num_angles: int, success: wp.array(dtype=int), qpos_out: wp.array(dtype=float), ): @@ -447,10 +574,9 @@ def compute_ik_kernel( Compute inverse kinematics (IK) in parallel for multiple target poses. Args: - combinations (wp.array): Array of combinations, where each entry specifies - the indices of the target pose, configuration, and reference angle. target_xpos_list (wp.array): Array of target poses (4x4 transformation matrices). angles_list (wp.array): Array of reference angles for IK computation. + qpos_seed (wp.array): Seed joint positions used for periodic limit wrapping. qpos_limits (wp.array): Array of joint position limits (min, max) for each joint. configs (wp.array): Array of configuration vectors (shoulder, elbow, wrist). dh_params (wp.array): Denavit-Hartenberg parameters for the robot. @@ -462,6 +588,8 @@ def compute_ik_kernel( plane_normal (wp.array): Output array for computed plane normal vectors. base_to_elbow_rotation (wp.array): Output array for base-to-elbow rotation matrices. joints_plane (wp.array): Output array for computed joint angles in the plane. + num_configs (int): Number of shoulder/elbow/wrist configurations. + num_angles (int): Number of redundancy-angle samples per target. success (wp.array): Output array indicating whether IK computation was successful. qpos_out (wp.array): Output array for computed joint positions. @@ -473,14 +601,14 @@ def compute_ik_kernel( tid = wp.tid() # Thread ID (for batch processing, if needed) # Extract indices - target_idx = int(combinations[tid][0]) - config_idx = int(combinations[tid][1]) - angle_idx = int(combinations[tid][2]) + angle_idx = tid % num_angles + config_idx = (tid // num_angles) % num_configs + target_idx = tid // (num_angles * num_configs) # Load inputs target_xpos = target_xpos_list[target_idx] config = configs[config_idx] - angle_ref = angles_list[angle_idx] + angle_ref = angles_list[target_idx * num_angles + angle_idx] # Extract shoulder, elbow, wrist configurations shoulder_config, elbow_config, wrist_config = config.x, config.y, config.z @@ -497,7 +625,7 @@ def compute_ik_kernel( # Compute shoulder-to-wrist vector P02 = wp.vec3(0.0, 0.0, link_lengths[0]) - P67 = wp.vec3(0.0, 0.0, dh_params[12]) + P67 = wp.vec3(0.0, 0.0, dh_params[6 * 4 + 0]) P06 = P_target - R_target @ P67 P26 = P06 - P02 @@ -557,10 +685,15 @@ def compute_ik_kernel( + (wp.outer(usw, usw) @ R03_o) ) - # TODO: judgment shoulder singularity - q1 = wp.atan2(R03[1, 1] * shoulder_config, R03[0, 1] * shoulder_config) q2 = safe_acos(R03[2, 1]) * shoulder_config - q3 = wp.atan2(-R03[2, 2] * shoulder_config, -R03[2, 0] * shoulder_config) + q1 = float(0.0) + q3 = float(0.0) + if wp.abs(wp.sin(q2)) <= 1e-6: + q1 = qpos_seed[target_idx * 7] * rotation_directions[0] + dh_params[3] + q3 = wp.atan2(R03[1, 0], R03[0, 0]) - q1 + else: + q1 = wp.atan2(R03[1, 1] * shoulder_config, R03[0, 1] * shoulder_config) + q3 = wp.atan2(-R03[2, 2] * shoulder_config, -R03[2, 0] * shoulder_config) # Calculate wrist joint angles (q5, q6, q7) Aw = wp.transpose(R34) @ wp.transpose(As) @ R_target @@ -569,10 +702,15 @@ def compute_ik_kernel( R47 = Aw * s_psi + Bw * c_psi + Cw q4 = joints_v[3] - # TODO: judgment wrist singularity - q5 = wp.atan2(R47[1, 2] * wrist_config, R47[0, 2] * wrist_config) q6 = safe_acos(R47[2, 2]) * wrist_config - q7 = wp.atan2(R47[2, 1] * wrist_config, -R47[2, 0] * wrist_config) + q5 = float(0.0) + q7 = float(0.0) + if wp.abs(wp.sin(q6)) <= 1e-6: + q5 = qpos_seed[target_idx * 7 + 4] * rotation_directions[4] + dh_params[19] + q7 = wp.atan2(-R47[2, 0], R47[0, 0]) - q5 + else: + q5 = wp.atan2(R47[1, 2] * wrist_config, R47[0, 2] * wrist_config) + q7 = wp.atan2(R47[2, 1] * wrist_config, -R47[2, 0] * wrist_config) out_of_limits = int(0) @@ -584,27 +722,44 @@ def compute_ik_kernel( q6_val = (q6 - dh_params[23]) * rotation_directions[5] q7_val = (q7 - dh_params[27]) * rotation_directions[6] - out_of_limits = int(0) - out_of_limits = out_of_limits | ( - 1 if (q1_val < qpos_limits[0][0] or q1_val > qpos_limits[0][1]) else 0 + wrapped_q1 = wrap_to_limit( + q1_val, qpos_limits[0][0], qpos_limits[0][1], qpos_seed[target_idx * 7] ) - out_of_limits = out_of_limits | ( - 1 if (q2_val < qpos_limits[1][0] or q2_val > qpos_limits[1][1]) else 0 + wrapped_q2 = wrap_to_limit( + q2_val, qpos_limits[1][0], qpos_limits[1][1], qpos_seed[target_idx * 7 + 1] ) - out_of_limits = out_of_limits | ( - 1 if (q3_val < qpos_limits[2][0] or q3_val > qpos_limits[2][1]) else 0 + wrapped_q3 = wrap_to_limit( + q3_val, qpos_limits[2][0], qpos_limits[2][1], qpos_seed[target_idx * 7 + 2] ) - out_of_limits = out_of_limits | ( - 1 if (q4_val < qpos_limits[3][0] or q4_val > qpos_limits[3][1]) else 0 + wrapped_q4 = wrap_to_limit( + q4_val, qpos_limits[3][0], qpos_limits[3][1], qpos_seed[target_idx * 7 + 3] ) - out_of_limits = out_of_limits | ( - 1 if (q5_val < qpos_limits[4][0] or q5_val > qpos_limits[4][1]) else 0 + wrapped_q5 = wrap_to_limit( + q5_val, qpos_limits[4][0], qpos_limits[4][1], qpos_seed[target_idx * 7 + 4] ) - out_of_limits = out_of_limits | ( - 1 if (q6_val < qpos_limits[5][0] or q6_val > qpos_limits[5][1]) else 0 + wrapped_q6 = wrap_to_limit( + q6_val, qpos_limits[5][0], qpos_limits[5][1], qpos_seed[target_idx * 7 + 5] ) - out_of_limits = out_of_limits | ( - 1 if (q7_val < qpos_limits[6][0] or q7_val > qpos_limits[6][1]) else 0 + wrapped_q7 = wrap_to_limit( + q7_val, qpos_limits[6][0], qpos_limits[6][1], qpos_seed[target_idx * 7 + 6] + ) + q1_val = wrapped_q1[1] + q2_val = wrapped_q2[1] + q3_val = wrapped_q3[1] + q4_val = wrapped_q4[1] + q5_val = wrapped_q5[1] + q6_val = wrapped_q6[1] + q7_val = wrapped_q7[1] + + out_of_limits = int( + wrapped_q1[0] + * wrapped_q2[0] + * wrapped_q3[0] + * wrapped_q4[0] + * wrapped_q5[0] + * wrapped_q6[0] + * wrapped_q7[0] + < 0.5 ) # Check joint limits @@ -640,78 +795,6 @@ def compute_ik_kernel( success[tid] = 0 # Mark as failed -@wp.kernel -def sort_ik_kernel( - qpos_out: wp.array(dtype=float), # [N * N_SOL, 7] - success: wp.array(dtype=int), # [N * N_SOL] - qpos_seed: wp.array(dtype=float), # [N, 7] - ik_weight: wp.array(dtype=float), # [7] - distances: wp.array(dtype=float), # [N, N_SOL] - indices: wp.array(dtype=int), # [N, N_SOL] - N_SOL: int, - sorted_qpos: wp.array(dtype=float), # [N, N_SOL, 7] - sorted_valid: wp.array(dtype=int), # [N, N_SOL] -): - """ - Sort inverse kinematics (IK) solutions for multiple targets based on their distances - to a seed configuration. - - Args: - qpos_out (wp.array): Array of computed joint positions for all solutions - ([N * N_SOL, 7]). - success (wp.array): Array indicating whether each solution is valid ([N * N_SOL]). - qpos_seed (wp.array): Array of seed joint positions for each target ([N, 7]). - ik_weight (wp.array): Array of weights for each joint to compute distance ([7]). - distances (wp.array): Output array to store computed distances ([N, N_SOL]). - indices (wp.array): Output array to store sorted indices ([N, N_SOL]). - N_SOL (int): Number of solutions per target. - sorted_qpos (wp.array): Output array for sorted joint positions ([N, N_SOL, 7]). - sorted_valid (wp.array): Output array for sorted validity flags ([N, N_SOL]). - """ - tid = wp.tid() # target index - - # 1. compute distances - for i in range(N_SOL): - idx = tid * N_SOL + i - valid = success[idx] - dist = 0.0 - if valid: - for j in range(7): - diff = qpos_out[idx * 7 + j] - qpos_seed[tid * 7 + j] - dist += ik_weight[j] * diff * diff - else: - dist = 1e10 - - distances[idx] = dist - indices[idx] = i - - # 2. bubble sort (only sort the N_SOL solutions for the current target) - for i in range(N_SOL): - min_idx = i - for j in range(i + 1, N_SOL): - idx_a = tid * N_SOL + min_idx - idx_b = tid * N_SOL + j - if distances[idx_b] < distances[idx_a]: - min_idx = j - # Swap - if min_idx != i: - idx_i = tid * N_SOL + i - idx_min = tid * N_SOL + min_idx - tmp_dist = distances[idx_i] - distances[idx_i] = distances[idx_min] - distances[idx_min] = tmp_dist - tmp_idx = indices[idx_i] - indices[idx_i] = indices[idx_min] - indices[idx_min] = tmp_idx - - # 3. reorder qpos_out and success according to sorted indices - for i in range(N_SOL): - src_idx = tid * N_SOL + indices[tid * N_SOL + i] - for j in range(7): - sorted_qpos[(tid * N_SOL + i) * 7 + j] = qpos_out[src_idx * 7 + j] - sorted_valid[tid * N_SOL + i] = success[src_idx] - - @wp.kernel def nearest_ik_kernel( qpos_out: wp.array(dtype=float), # [N * N_SOL * 7] @@ -746,7 +829,8 @@ def nearest_ik_kernel( if success[idx]: dist = 0.0 for j in range(7): - diff = qpos_out[idx * 7 + j] - qpos_seed[tid * 7 + j] + raw_diff = qpos_out[idx * 7 + j] - qpos_seed[tid * 7 + j] + diff = wp.atan2(wp.sin(raw_diff), wp.cos(raw_diff)) dist += ik_weight[j] * diff * diff if dist < min_dist: min_dist = dist diff --git a/scripts/benchmark/robotics/kinematic_solver/srs_solver.py b/scripts/benchmark/robotics/kinematic_solver/srs_solver.py new file mode 100644 index 000000000..4d8bee322 --- /dev/null +++ b/scripts/benchmark/robotics/kinematic_solver/srs_solver.py @@ -0,0 +1,460 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Benchmark SRS analytical IK across representative workloads. + +The benchmark covers randomized nominal, wide-range, joint-boundary, +near-singular, and unreachable targets; perturbs IK seeds independently from +the FK ground truth; and reports repeated latency, throughput, classification +accuracy, FK reconstruction error, and solution distance from the seed. +Run: python -m scripts.benchmark.robotics.kinematic_solver.srs_solver +""" + +from __future__ import annotations + +import argparse +import os +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +import numpy as np +import psutil +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim.robots.dexforce_w1.params import W1ArmKineParams +from embodichain.lab.sim.robots.dexforce_w1.types import ( + DexforceW1ArmSide, + DexforceW1Version, +) +from embodichain.lab.sim.solvers.srs_solver import SRSSolverCfg +from embodichain.utils.logger import set_log_level + +DEFAULT_SIZES = (1, 16, 128) +SCENARIOS = ("nominal", "wide", "boundary", "near-singular", "unreachable") + + +@dataclass +class BenchmarkCase: + """Inputs and expected reachability for one benchmark case.""" + + target: torch.Tensor + seed: torch.Tensor + expected_reachable: bool + + +def _parse_args() -> argparse.Namespace: + """Parse benchmark controls.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--sizes", nargs="+", type=int, default=list(DEFAULT_SIZES)) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--random-seed", type=int, default=20260824) + parser.add_argument("--arms", choices=("left", "right", "both"), default="both") + parser.add_argument("--devices", choices=("cpu", "cuda", "both"), default="both") + parser.add_argument("--modes", choices=("seeded", "full", "both"), default="both") + parser.add_argument( + "--scenarios", nargs="+", choices=SCENARIOS, default=list(SCENARIOS) + ) + return parser.parse_args() + + +def _selected(value: str, first: str, second: str) -> tuple[str, ...]: + """Expand a two-choice CLI selector.""" + return (first, second) if value == "both" else (value,) + + +def _make_solver(device: torch.device, search_mode: str, arm: str, size: int): + """Construct one W1 arm SRS solver.""" + side = DexforceW1ArmSide.LEFT if arm == "left" else DexforceW1ArmSide.RIGHT + prefix = "LEFT" if arm == "left" else "RIGHT" + params = W1ArmKineParams(arm_side=side, version=DexforceW1Version.V021) + cfg = SRSSolverCfg( + urdf_path=get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf"), + joint_names=[f"{prefix}_J{i + 1}" for i in range(7)], + root_link_name=f"{arm}_arm_base", + end_link_name=f"{arm}_ee", + dh_params=params.dh_params, + user_qpos_limits=params.qpos_limits, + T_b_ob=params.T_b_ob, + T_e_oe=params.T_e_oe, + link_lengths=params.link_lengths, + rotation_directions=params.rotation_directions, + search_mode=search_mode, + ) + return cfg.init_solver(num_envs=size, device=device) + + +def _memory_snapshot() -> tuple[float, float]: + """Return process RSS and PyTorch CUDA allocation in MiB.""" + cpu_mb = psutil.Process(os.getpid()).memory_info().rss / 1024**2 + gpu_mb = ( + torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 + ) + return cpu_mb, gpu_mb + + +def _synchronize(device: torch.device) -> None: + """Synchronize CUDA work before observing wall-clock time.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def _joint_limits(solver) -> tuple[torch.Tensor, torch.Tensor]: + """Return solver limits on its execution device.""" + limits = solver.get_qpos_limits() + lower = torch.tensor( + limits["lower_qpos_limits"], dtype=torch.float32, device=solver.device + ) + upper = torch.tensor( + limits["upper_qpos_limits"], dtype=torch.float32, device=solver.device + ) + return lower, upper + + +def _uniform_qpos( + lower: torch.Tensor, + upper: torch.Tensor, + size: int, + low_fraction: float, + high_fraction: float, + generator: torch.Generator, +) -> torch.Tensor: + """Sample joint configurations from a fractional limit interval.""" + unit = torch.rand( + (size, 7), generator=generator, dtype=torch.float32, device="cpu" + ).to(lower.device) + low = lower + low_fraction * (upper - lower) + high = lower + high_fraction * (upper - lower) + return low + unit * (high - low) + + +def _make_case( + solver, + scenario: str, + size: int, + generator: torch.Generator, +) -> BenchmarkCase: + """Generate independent ground-truth joints, target poses, and IK seeds.""" + lower, upper = _joint_limits(solver) + span = upper - lower + if scenario == "nominal": + truth = _uniform_qpos(lower, upper, size, 0.25, 0.75, generator) + elif scenario in ("wide", "unreachable"): + truth = _uniform_qpos(lower, upper, size, 0.05, 0.95, generator) + elif scenario == "boundary": + choose_upper = ( + torch.randint(0, 2, (size, 7), generator=generator, device="cpu") + .bool() + .to(lower.device) + ) + near_lower = lower + 0.01 * span + near_upper = upper - 0.01 * span + truth = torch.where(choose_upper, near_upper, near_lower) + elif scenario == "near-singular": + truth = _uniform_qpos(lower, upper, size, 0.25, 0.75, generator) + # SRS elbow and wrist pitch approach their singular values without + # using the exact unreachable straight-arm boundary. + truth[:, 3] = torch.clamp( + torch.full_like(truth[:, 3], -1e-3), lower[3], upper[3] + ) + truth[:, 5] = torch.clamp( + torch.full_like(truth[:, 5], 1e-3), lower[5], upper[5] + ) + else: + raise ValueError(f"Unknown scenario: {scenario}") + + target = solver.get_fk(truth) + noise = torch.randn( + (size, 7), generator=generator, dtype=torch.float32, device="cpu" + ).to(lower.device) + seed = torch.clamp(truth + noise * (0.12 * span), lower, upper) + if scenario == "unreachable": + offsets = torch.tensor( + [2.0, -2.0, 2.0], dtype=target.dtype, device=target.device + ) + target = target.clone() + target[:, :3, 3] += offsets + return BenchmarkCase(target, seed, False) + return BenchmarkCase(target, seed, True) + + +def _solution_matrix(solution: torch.Tensor) -> torch.Tensor: + """Normalize solver output to one solution per target.""" + return solution[:, 0] if solution.ndim == 3 else solution + + +def _pose_error_vectors( + target: torch.Tensor, actual: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Return per-sample translation millimetres and rotation degrees.""" + target = target.to(torch.float64) + actual = actual.to(torch.float64) + translation = torch.linalg.norm(target[:, :3, 3] - actual[:, :3, 3], dim=1) + relative = target[:, :3, :3].transpose(1, 2) @ actual[:, :3, :3] + # Float32 FK matrices can be microscopically non-orthogonal. Project the + # relative matrix onto SO(3) before converting its trace to an angle. + u, _, vh = torch.linalg.svd(relative) + projected = u @ vh + negative_determinant = torch.linalg.det(projected) < 0.0 + if negative_determinant.any(): + u = u.clone() + u[negative_determinant, :, -1] *= -1.0 + projected = u @ vh + cosine = ((projected.diagonal(dim1=1, dim2=2).sum(1) - 1.0) / 2.0).clamp(-1, 1) + return translation * 1000.0, torch.rad2deg(torch.acos(cosine)) + + +def _percentile(values: torch.Tensor, quantile: float) -> float: + """Return a finite percentile or NaN for an empty tensor.""" + return float(torch.quantile(values.float(), quantile)) if values.numel() else np.nan + + +def _quality_metrics( + solver, + case: BenchmarkCase, + success: torch.Tensor, + solution: torch.Tensor, +) -> dict[str, float]: + """Compute correctness metrics without mixing failed solutions into errors.""" + success = success.bool() + solution = _solution_matrix(solution) + expected = torch.full_like(success, case.expected_reachable) + classification_accuracy = float((success == expected).float().mean()) + success_rate = float(success.float().mean()) + if not success.any(): + return { + "success_rate": success_rate, + "classification_accuracy": classification_accuracy, + "translation_mean_mm": np.nan, + "translation_p95_mm": np.nan, + "translation_max_mm": np.nan, + "rotation_mean_deg": np.nan, + "rotation_p95_deg": np.nan, + "seed_distance_mean_rad": np.nan, + } + valid_solution = solution[success] + actual = solver.get_fk(valid_solution) + translation, rotation = _pose_error_vectors(case.target[success], actual) + delta = valid_solution - case.seed[success] + wrapped_delta = torch.atan2(torch.sin(delta), torch.cos(delta)) + seed_distance = torch.linalg.vector_norm(wrapped_delta, dim=1) + return { + "success_rate": success_rate, + "classification_accuracy": classification_accuracy, + "translation_mean_mm": float(translation.mean()), + "translation_p95_mm": _percentile(translation, 0.95), + "translation_max_mm": float(translation.max()), + "rotation_mean_deg": float(rotation.mean()), + "rotation_p95_deg": _percentile(rotation, 0.95), + "seed_distance_mean_rad": float(seed_distance.mean()), + } + + +def _measure_case( + solver, + case: BenchmarkCase, + repeats: int, + warmup: int, +) -> tuple[dict[str, float], dict[str, float]]: + """Measure repeated solve latency/memory and return final quality metrics.""" + for _ in range(warmup): + solver.get_ik(case.target, case.seed) + _synchronize(solver.device) + if solver.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(solver.device) + cpu_before, gpu_before = _memory_snapshot() + durations = [] + success = solution = None + for _ in range(repeats): + _synchronize(solver.device) + start = time.perf_counter() + success, solution = solver.get_ik(case.target, case.seed) + _synchronize(solver.device) + durations.append((time.perf_counter() - start) * 1000.0) + cpu_after, gpu_after = _memory_snapshot() + assert success is not None and solution is not None + latency = np.asarray(durations) + median_ms = float(np.median(latency)) + performance = { + "latency_median_ms": median_ms, + "latency_p95_ms": float(np.percentile(latency, 95)), + "latency_min_ms": float(latency.min()), + "throughput_targets_s": case.target.shape[0] * 1000.0 / median_ms, + "cpu_delta_mb": cpu_after - cpu_before, + "gpu_delta_mb": gpu_after - gpu_before, + "peak_gpu_mb": ( + torch.cuda.max_memory_allocated(solver.device) / 1024**2 + if solver.device.type == "cuda" + else 0.0 + ), + } + return performance, _quality_metrics(solver, case, success, solution) + + +def _format_table(rows: list[dict[str, object]]) -> list[str]: + """Format rows as one Markdown table.""" + headers = list(rows[0]) + lines = [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join(["---"] * len(headers)) + " |", + ] + lines.extend( + "| " + " | ".join(str(row[key]) for key in headers) + " |" for row in rows + ) + return lines + + +def _leaderboard(metric_rows: list[dict[str, object]]) -> list[dict[str, object]]: + """Rank every implementation by correctness, then speed.""" + grouped: dict[str, list[dict[str, object]]] = {} + for row in metric_rows: + grouped.setdefault(str(row["impl"]), []).append(row) + ranking = [] + for impl, rows in grouped.items(): + accuracy = float( + np.mean([float(row["classification_accuracy"]) for row in rows]) + ) + reachable = [row for row in rows if row["expected_reachable"]] + success = float(np.mean([float(row["success_rate"]) for row in reachable])) + ranking.append((impl, accuracy, success)) + ranking.sort(key=lambda item: (item[1], item[2]), reverse=True) + return [ + { + "rank": rank, + "algorithm": impl, + "classification_accuracy": f"{accuracy:.2%}", + "reachable_success_rate": f"{success:.2%}", + } + for rank, (impl, accuracy, success) in enumerate(ranking, 1) + ] + + +def _write_report( + perf_rows: list[dict[str, object]], + metric_rows: list[dict[str, object]], + args: argparse.Namespace, +) -> Path: + """Write a single report containing exactly three Markdown tables.""" + lines = [ + "# SRS Solver Benchmark", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + "", + f"Repeats: {args.repeats}; warm-up calls: {args.warmup}; random seed: {args.random_seed}.", + "", + "## Time & Memory", + "", + *_format_table(perf_rows), + "", + "## Success & Other Metrics", + "", + *_format_table(metric_rows), + "", + "## Leaderboard", + "", + *_format_table(_leaderboard(metric_rows)), + "", + "Errors are computed only over successful solutions. For unreachable cases, classification accuracy rewards rejection rather than success.", + ] + output_dir = Path("outputs/benchmarks") + output_dir.mkdir(parents=True, exist_ok=True) + path = output_dir / f"srs_solver_{datetime.now():%Y%m%d_%H%M%S}.md" + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path + + +def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: + """Run the configured SRS benchmark matrix and return its report path.""" + args = _parse_args() if args is None else args + set_log_level("ERROR") + if args.repeats < 1 or args.warmup < 0 or any(size < 1 for size in args.sizes): + raise ValueError( + "sizes/repeats must be positive and warmup must be non-negative" + ) + arms = _selected(args.arms, "left", "right") + modes = _selected(args.modes, "seeded", "full") + requested_devices = _selected(args.devices, "cpu", "cuda") + devices = [] + for name in requested_devices: + if name == "cuda" and not torch.cuda.is_available(): + print("Skipping CUDA: torch.cuda.is_available() is False") + continue + devices.append(torch.device(name)) + if not devices: + raise RuntimeError("No requested benchmark device is available") + + perf_rows: list[dict[str, object]] = [] + metric_rows: list[dict[str, object]] = [] + for device in devices: + for arm in arms: + for mode in modes: + for size in args.sizes: + solver = _make_solver(device, mode, arm, size) + for scenario_index, scenario in enumerate(args.scenarios): + generator = torch.Generator(device="cpu") + generator.manual_seed( + args.random_seed + scenario_index + size * 1009 + ) + case = _make_case(solver, scenario, size, generator) + performance, quality = _measure_case( + solver, case, args.repeats, args.warmup + ) + impl = f"{device.type}-{mode}-{arm}" + perf_rows.append( + { + "sample_size": size, + "scenario": scenario, + "impl": impl, + "latency_median_ms": f"{performance['latency_median_ms']:.3f}", + "latency_p95_ms": f"{performance['latency_p95_ms']:.3f}", + "throughput_targets_s": f"{performance['throughput_targets_s']:.1f}", + "cpu_delta_mb": f"{performance['cpu_delta_mb']:+.2f}", + "gpu_delta_mb": f"{performance['gpu_delta_mb']:+.2f}", + "peak_gpu_mb": f"{performance['peak_gpu_mb']:.2f}", + } + ) + metric_rows.append( + { + "sample_size": size, + "scenario": scenario, + "impl": impl, + "expected_reachable": case.expected_reachable, + "success_rate": f"{quality['success_rate']:.4f}", + "classification_accuracy": f"{quality['classification_accuracy']:.4f}", + "translation_mean_mm": f"{quality['translation_mean_mm']:.6f}", + "translation_p95_mm": f"{quality['translation_p95_mm']:.6f}", + "translation_max_mm": f"{quality['translation_max_mm']:.6f}", + "rotation_mean_deg": f"{quality['rotation_mean_deg']:.6f}", + "rotation_p95_deg": f"{quality['rotation_p95_deg']:.6f}", + "seed_distance_mean_rad": f"{quality['seed_distance_mean_rad']:.6f}", + } + ) + print( + f"{impl:>22} n={size:>4} {scenario:<13} " + f"median={performance['latency_median_ms']:>9.3f} ms " + f"success={quality['success_rate']:>7.2%} " + f"class={quality['classification_accuracy']:>7.2%}" + ) + report = _write_report(perf_rows, metric_rows, args) + print(f"Markdown report saved: {report}") + return report + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index 8c6d1bfd6..4be55f561 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -14,17 +14,21 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Move a W1 end effector along a Cartesian straight line with SRS IK.""" + from __future__ import annotations import argparse +import sys import time +import traceback + import numpy as np import torch -from IPython import embed - -from embodichain.lab.sim.objects import Robot from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import MarkerCfg +from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots import DexforceW1Cfg from embodichain.lab.visualization import ( VisualizationCfg, @@ -33,67 +37,365 @@ ) -def main(visualization: VisualizationCfg | None = None) -> None: - # Set print options for better readability +def main( + device: str, + num_steps: int, + line_offset: tuple[float, float, float], + physics_steps: int, + max_joint_step_deg: float, + visualization: VisualizationCfg | None = None, +) -> None: + """Run sequential SRS IK targets along a fixed-orientation straight line. + + Args: + device: Simulation and SRS solver device, either ``cpu`` or ``cuda``. + num_steps: Number of Cartesian waypoints, including both endpoints. + line_offset: End-point translation relative to the initial TCP pose, in meters. + physics_steps: Number of simulation steps displayed per waypoint. + max_joint_step_deg: Maximum allowed change of any joint per IK waypoint. + visualization: Optional visualization configuration. + """ + if num_steps < 2: + raise ValueError("num_steps must be at least 2") + if physics_steps < 1: + raise ValueError("physics_steps must be at least 1") + if max_joint_step_deg <= 0.0: + raise ValueError("max_joint_step_deg must be positive") + if device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested, but PyTorch cannot access a CUDA device" + ) + np.set_printoptions(precision=5, suppress=True) torch.set_printoptions(precision=5, sci_mode=False) - - # Initialize simulation - sim_device = "cpu" sim = SimulationManager( SimulationManagerCfg( - headless=False, - sim_device=sim_device, + # Keep the native window closed while planning so renderer/window + # lifecycle events cannot terminate or perturb timed CUDA IK calls. + headless=True, + sim_device=device, width=2200, height=1200, visualization=visualization or VisualizationCfg(), ) ) + # Keep pose conversion, IK, and FK error measurement deterministic. In automatic + # mode the engine thread may advance the robot base between these operations. + sim.set_manual_update(True) - sim.set_manual_update(False) + try: + arm_name = "left_arm" + robot_cfg = DexforceW1Cfg.from_dict({"uid": "dexforce_w1"}) + # The robot default intentionally ignores the elbow joint when ranking IK + # candidates. For a Cartesian-path tutorial, weight every joint to prevent + # visually disruptive branch changes between adjacent waypoints. + robot_cfg.solver_cfg[arm_name].ik_nearest_weight = np.array( + [2.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0] + ) + robot: Robot = sim.add_robot(cfg=robot_cfg) + joint_ids = robot.get_joint_ids(arm_name) + qpos_seed = torch.tensor( + [[np.pi / 6, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, np.pi / 6]], + dtype=torch.float32, + device=sim.device, + ) + reference_qpos = qpos_seed.clone() + robot.set_qpos(qpos_seed, joint_ids=joint_ids, target=False) + robot.set_qpos(qpos_seed, joint_ids=joint_ids, target=True) + sim.update(step=physics_steps) - robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) - arm_name = "left_arm" - # Set initial joint positions for left arm - qpos_fk_list = [ - torch.tensor([[0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], dtype=torch.float32), - ] - robot.set_qpos(qpos_fk_list[0], joint_ids=robot.get_joint_ids(arm_name)) + start_pose = robot.compute_fk(qpos=qpos_seed, name=arm_name, to_matrix=True) + target_poses = start_pose.repeat(num_steps, 1, 1) + offset = torch.tensor( + line_offset, dtype=start_pose.dtype, device=start_pose.device + ) + interpolation = torch.linspace( + 0.0, 1.0, num_steps, dtype=start_pose.dtype, device=start_pose.device + ) + target_poses[:, :3, 3] += interpolation.unsqueeze(1) * offset - time.sleep(0.5) + # Warm up lazy FK compilation and the selected SRS backend outside timing. + robot.compute_ik( + pose=target_poses[:1], + joint_seed=qpos_seed, + name=arm_name, + return_all_solutions=True, + ) + if device == "cuda": + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() - fk_xpos_batch = torch.cat(qpos_fk_list, dim=0) + solve_times_ms: list[float] = [] + candidate_counts: list[int] = [] + translation_errors_mm: list[float] = [] + waypoint_candidates: list[torch.Tensor] = [] + continuity_weights = torch.tensor( + [2.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0], + dtype=qpos_seed.dtype, + device=qpos_seed.device, + ) + max_joint_step = torch.deg2rad( + torch.tensor( + max_joint_step_deg, dtype=qpos_seed.dtype, device=qpos_seed.device + ) + ) + for waypoint, target_pose in enumerate(target_poses): + start_time = time.perf_counter() + success, solution = robot.compute_ik( + pose=target_pose.unsqueeze(0), + # Use one fixed seed while enumerating candidates. Path continuity + # is selected globally below instead of greedily changing the + # candidate set after every waypoint. + joint_seed=reference_qpos, + name=arm_name, + return_all_solutions=True, + ) + if device == "cuda": + torch.cuda.synchronize() + solve_times_ms.append((time.perf_counter() - start_time) * 1000.0) - fk_xpos_list = robot.compute_fk(qpos=fk_xpos_batch, name=arm_name, to_matrix=True) + if not bool(success[0]): + raise RuntimeError( + f"Trajectory planning failed at waypoint " + f"{waypoint + 1}/{num_steps}; nothing was executed." + ) - start_time = time.time() - res, ik_qpos = robot.compute_ik( - pose=fk_xpos_list, - name=arm_name, - # joint_seed=qpos_fk_list[0], - return_all_solutions=True, - ) - end_time = time.time() - print( - f"Batch IK computation time for {len(fk_xpos_list)} poses: {end_time - start_time:.6f} seconds" - ) + candidates = solution[0] + candidate_counts.append(candidates.shape[0]) + waypoint_candidates.append(candidates) + print( + f"Enumerated {waypoint + 1:03d}/{num_steps}: " + f"solve={solve_times_ms[-1]:.3f} ms, " + f"candidates={candidates.shape[0]}" + ) - if ik_qpos.dim() == 3: - first_solutions = ik_qpos[:, 0, :] - else: - first_solutions = ik_qpos - robot.set_qpos(first_solutions, joint_ids=robot.get_joint_ids(arm_name)) + # Find a minimum-cost path through the layered IK candidate graph. This + # avoids the greedy failure mode where a locally attractive arm angle has + # no continuous successor at a later waypoint. + path_costs: list[torch.Tensor] = [] + predecessors: list[torch.Tensor] = [] + first_delta = torch.atan2( + torch.sin(waypoint_candidates[0] - reference_qpos), + torch.cos(waypoint_candidates[0] - reference_qpos), + ) + first_allowed = first_delta.abs().amax(dim=1) <= max_joint_step + first_cost = (first_delta.square() * continuity_weights).sum(dim=1) + first_cost.masked_fill_(~first_allowed, float("inf")) + path_costs.append(first_cost) + predecessors.append(torch.full_like(first_cost, -1, dtype=torch.long)) - ik_xpos_list = robot.compute_fk(qpos=first_solutions, name=arm_name, to_matrix=True) + for waypoint in range(1, num_steps): + previous_candidates = waypoint_candidates[waypoint - 1] + candidates = waypoint_candidates[waypoint] + edge_delta = torch.atan2( + torch.sin(candidates[:, None, :] - previous_candidates[None, :, :]), + torch.cos(candidates[:, None, :] - previous_candidates[None, :, :]), + ) + allowed_edges = edge_delta.abs().amax(dim=2) <= max_joint_step + transition_cost = (edge_delta.square() * continuity_weights).sum(dim=2) + transition_cost.masked_fill_(~allowed_edges, float("inf")) + reference_delta = torch.atan2( + torch.sin(candidates - reference_qpos), + torch.cos(candidates - reference_qpos), + ) + node_cost = 0.05 * (reference_delta.square() * continuity_weights).sum( + dim=1 + ) + total_cost = transition_cost + path_costs[-1].unsqueeze(0) + best_cost, best_predecessor = total_cost.min(dim=1) + best_cost += node_cost + if not bool(torch.isfinite(best_cost).any()): + reachable_previous = torch.isfinite(path_costs[-1]) + reachable_edges = edge_delta[:, reachable_previous] + smallest_step = reachable_edges.abs().amax(dim=2).min() + raise RuntimeError( + f"No globally continuous IK path reaches waypoint " + f"{waypoint + 1}/{num_steps}: smallest available maximum " + f"joint step is {torch.rad2deg(smallest_step).item():.3f} deg, " + f"limit is {max_joint_step_deg:.3f} deg." + ) + path_costs.append(best_cost) + predecessors.append(best_predecessor) - print("fk_xpos_list: ", fk_xpos_list) - print("ik_xpos_list: ", ik_xpos_list) + selected_indices = [int(path_costs[-1].argmin())] + for waypoint in range(num_steps - 1, 0, -1): + selected_indices.append(int(predecessors[waypoint][selected_indices[-1]])) + selected_indices.reverse() + planned_qpos = [ + waypoint_candidates[i][selected_indices[i]].unsqueeze(0) + for i in range(num_steps) + ] - sim.capture_visualization(force=True) - embed(header="Test SRSSolver example. Press Ctrl-D to exit.") + for waypoint, (target_pose, qpos_seed) in enumerate( + zip(target_poses, planned_qpos, strict=True) + ): + actual_pose = robot.compute_fk( + qpos=qpos_seed, name=arm_name, to_matrix=True + ) + error_mm = float( + torch.linalg.vector_norm( + actual_pose[0, :3, 3] - target_pose[:3, 3] + ).item() + * 1000.0 + ) + translation_errors_mm.append(error_mm) + print( + f"Planned {waypoint + 1:03d}/{num_steps}: " f"error={error_mm:.4f} mm" + ) + + print( + f"SRS {device.upper()} planning summary: {num_steps}/{num_steps} solved, " + f"median/p95/mean solve={np.median(solve_times_ms):.3f}/" + f"{np.percentile(solve_times_ms, 95):.3f}/" + f"{np.mean(solve_times_ms):.3f} ms, " + f"max translation error=" + f"{max(translation_errors_mm, default=float('nan')):.4f} mm" + ) + print( + f"IK candidates per waypoint min/median/max: " + f"{min(candidate_counts)}/{int(np.median(candidate_counts))}/" + f"{max(candidate_counts)}" + ) + if device == "cuda": + print( + f"CUDA peak allocated memory: " + f"{torch.cuda.max_memory_allocated() / 1024**2:.2f} MiB" + ) + planned_qpos_tensor = torch.cat(planned_qpos) + wrapped_steps = torch.atan2( + torch.sin(planned_qpos_tensor[1:] - planned_qpos_tensor[:-1]), + torch.cos(planned_qpos_tensor[1:] - planned_qpos_tensor[:-1]), + ) + max_step_flat_index = wrapped_steps.abs().argmax() + max_step_waypoint = int(max_step_flat_index // wrapped_steps.shape[1]) + 2 + max_step_joint = int(max_step_flat_index % wrapped_steps.shape[1]) + 1 + print( + f"Max adjacent joint step: " + f"{torch.rad2deg(wrapped_steps.abs().max()).item():.3f} deg " + f"at waypoint {max_step_waypoint}, joint {max_step_joint}" + ) + if wrapped_steps.shape[0] > 1: + joint_step_changes = wrapped_steps[1:] - wrapped_steps[:-1] + print( + f"Max joint step change: " + f"{torch.rad2deg(joint_step_changes.abs().max()).item():.3f} deg" + ) + + sim.open_window() + marker_stride = max(1, num_steps // 25) + marker_indices = torch.arange(0, num_steps, marker_stride, device=sim.device) + if marker_indices[-1] != num_steps - 1: + marker_indices = torch.cat( + (marker_indices, marker_indices.new_tensor([num_steps - 1])) + ) + sim.draw_marker( + MarkerCfg( + name="srs_target_path", + marker_type="axis", + axis_xpos=target_poses[marker_indices], + axis_size=0.0006, + axis_len=0.008, + arena_index=0, + ) + ) + + print("Planning completed; executing the joint trajectory...") + execution_errors_mm: list[float] = [] + previous_qpos = robot.get_qpos(name=arm_name).clone() + zero_qvel = torch.zeros_like(previous_qpos) + for waypoint, qpos in enumerate(planned_qpos): + wrapped_delta = torch.atan2( + torch.sin(qpos - previous_qpos), torch.cos(qpos - previous_qpos) + ) + for substep in range(1, physics_steps + 1): + alpha = substep / physics_steps + interpolated_qpos = previous_qpos + alpha * wrapped_delta + robot.set_qpos(interpolated_qpos, joint_ids=joint_ids, target=False) + robot.set_qpos(interpolated_qpos, joint_ids=joint_ids, target=True) + robot.set_qvel(zero_qvel, joint_ids=joint_ids, target=False) + robot.set_qvel(zero_qvel, joint_ids=joint_ids, target=True) + sim.update(step=1) + previous_qpos = qpos + actual_qpos = robot.get_qpos(name=arm_name) + actual_pose = robot.compute_fk( + qpos=actual_qpos, name=arm_name, to_matrix=True + ) + execution_errors_mm.append( + float( + torch.linalg.vector_norm( + actual_pose[0, :3, 3] - target_poses[waypoint, :3, 3] + ).item() + * 1000.0 + ) + ) + if waypoint % marker_stride == 0 or waypoint == num_steps - 1: + sim.draw_marker( + MarkerCfg( + name=f"srs_executed_path_{waypoint:03d}", + marker_type="axis", + axis_xpos=actual_pose, + axis_size=0.0012, + axis_len=0.004, + arena_index=0, + ) + ) + print("Trajectory execution completed.") + print( + f"Max execution tracking error: {max(execution_errors_mm):.4f} mm. " + "Long axes show targets; short thick axes show executed samples." + ) + sim.capture_visualization(force=True) + finally: + # Do not use the default os._exit(0) cleanup path: it suppresses Python + # tracebacks raised during planning and makes failures look like clean exits. + sim.destroy(exit_process=False) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--device", choices=("cpu", "cuda"), default="cpu", help="SRS backend." + ) + parser.add_argument( + "--num-steps", type=int, default=50, help="Number of line waypoints." + ) + parser.add_argument( + "--line-offset", + type=float, + nargs=3, + metavar=("X", "Y", "Z"), + default=(0.0, -0.30, 0.0), + help="TCP line displacement in meters (default: 0 -0.30 0).", + ) + parser.add_argument( + "--physics-steps", + type=int, + default=20, + help="Simulation interpolation steps per IK waypoint (default: 20).", + ) + parser.add_argument( + "--max-joint-step-deg", + type=float, + default=15.0, + help="Reject an IK branch changing any joint by more than this angle.", + ) add_viser_args_to_parser(parser) - main(visualization=visualization_cfg_from_args(parser.parse_args())) + args = parser.parse_args() + exit_code = 0 + try: + main( + device=args.device, + num_steps=args.num_steps, + line_offset=tuple(args.line_offset), + physics_steps=args.physics_steps, + max_joint_step_deg=args.max_joint_step_deg, + visualization=visualization_cfg_from_args(args), + ) + except BaseException: + traceback.print_exc() + exit_code = 1 + finally: + # Deferred destruction is only safe after main() has unwound and no local + # Robot/solver wrappers remain live on its Python frame. + SimulationManager.flush_cleanup_queue() + sys.exit(exit_code) diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index 540814987..f760ed616 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -17,30 +17,29 @@ from __future__ import annotations import os + import torch torch._dynamo.config.cache_size_limit = 128 # recompile_limit -import pytest import numpy as np +import pytest +import warp as wp +from embodichain.data import get_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RobotCfg from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.cfg import RobotCfg, RenderCfg -from embodichain.data import get_data_path - -from embodichain.lab.sim.solvers.srs_solver import SRSSolver, SRSSolverCfg +from embodichain.lab.sim.robots.dexforce_w1.params import ( + W1ArmKineParams, +) from embodichain.lab.sim.robots.dexforce_w1.types import ( DexforceW1ArmSide, DexforceW1Version, ) -from embodichain.lab.sim.robots.dexforce_w1.params import ( - W1ArmKineParams, -) +from embodichain.lab.sim.solvers.srs_solver import SRSSolver, SRSSolverCfg class BaseSolverTest: - solver = {} - def get_arm_config(self): return [ (DexforceW1ArmSide.LEFT, "left_arm"), @@ -48,6 +47,7 @@ def get_arm_config(self): ] def setup_solver(self, solver_type: str, device: str = "cpu"): + self.solver = {} for arm_side, arm_name in self.get_arm_config(): arm_params = W1ArmKineParams( arm_side=arm_side, @@ -57,7 +57,7 @@ def setup_solver(self, solver_type: str, device: str = "cpu"): cfg = SRSSolverCfg() cfg.joint_names = [ - f"{'LEFT' if arm_side == DexforceW1ArmSide.LEFT else 'RIGHT'}_J{i+1}" + f"{'LEFT' if arm_side == DexforceW1ArmSide.LEFT else 'RIGHT'}_J{i + 1}" for i in range(7) ] cfg.end_link_name = ( @@ -75,9 +75,14 @@ def setup_solver(self, solver_type: str, device: str = "cpu"): cfg.T_b_ob = arm_params.T_b_ob cfg.link_lengths = arm_params.link_lengths cfg.rotation_directions = arm_params.rotation_directions + cfg.ik_nearest_weight = np.array([2.0, 2.0, 2.0, 0.0, 1.0, 1.0, 1.0]) self.solver[arm_name] = SRSSolver(cfg=cfg, num_envs=1, device=device) + def teardown_method(self): + """Release per-test solver instances and their backend allocations.""" + self.solver.clear() + @pytest.mark.parametrize( "arm_side, arm_name", [ @@ -105,6 +110,33 @@ def test_ik(self, arm_side: DexforceW1ArmSide, arm_name: str): fk_xpos, ik_xpos, atol=1e-3, rtol=1e-3 ), f"FK and IK results do not match for {arm_name}" + @pytest.mark.parametrize( + "return_all_solutions, expected_shape", + [(False, (2, 1, 7)), (True, (2, 0, 7))], + ) + def test_no_solution_preserves_output_rank( + self, return_all_solutions: bool, expected_shape: tuple[int, int, int] + ): + """Test an entirely failed batch still follows the IK output contract.""" + solver = self.solver[next(iter(self.solver))] + target_xpos = torch.eye(4, dtype=torch.float32, device=solver.device).repeat( + 2, 1, 1 + ) + target_xpos[:, :3, 3] = torch.tensor( + [100.0, 100.0, 100.0], dtype=torch.float32, device=solver.device + ) + qpos_seed = torch.zeros((2, 7), dtype=torch.float32, device=solver.device) + + success, solutions = solver.get_ik( + target_xpos, + qpos_seed=qpos_seed, + return_all_solutions=return_all_solutions, + ) + + assert success.shape == (2,) + assert not success.any() + assert solutions.shape == expected_shape + def test_update_with_robot_limit_intersects_existing_solver_limits(self): """Test robot limit sync only tightens solver limits and never widens them.""" solver_key = next(iter(self.solver)) @@ -174,14 +206,193 @@ def test_update_with_robot_limit_intersects_existing_solver_limits(self): atol=1e-5, ), "FAIL: robot sync did not tighten solver upper_qpos_limits" - @classmethod - def teardown_class(cls): - if cls.solver is not None: - try: - del cls.solver - print("solver destroyed successfully") - except Exception as e: - print(f"Error during solver destruction: {e}") + def test_seeded_redundancy_sampling_expands_around_geometric_arm_angle(self): + """Test redundancy samples expand around the seed's geometric arm angle.""" + solver = self.solver[next(iter(self.solver))] + solver.cfg.num_samples = 5 + seed = torch.tensor( + [[0.15, -0.35, 0.40, -0.70, 0.20, 0.30, -0.15]], + dtype=torch.float32, + device=solver.device, + ) + + seed_arm_angle = solver.impl._get_seed_arm_angles(seed) + angles = solver.impl._sample_elbow_angles(seed) + + step = solver.cfg.redundancy_step + expected = torch.stack( + ( + seed_arm_angle, + seed_arm_angle + step, + seed_arm_angle - step, + ), + dim=1, + ) + expected = torch.remainder(expected + torch.pi, 2.0 * torch.pi) - torch.pi + assert torch.allclose(angles[:, :3], expected, atol=1e-6) + assert angles.shape[1] == solver.cfg.num_samples + assert torch.all(angles >= -torch.pi) + assert torch.all(angles < torch.pi) + wrapped_delta = torch.atan2( + torch.sin(angles[:, :, None] - angles[:, None, :]), + torch.cos(angles[:, :, None] - angles[:, None, :]), + ).abs() + diagonal = torch.eye( + angles.shape[1], dtype=torch.bool, device=solver.device + ).unsqueeze(0) + assert torch.all(wrapped_delta.masked_fill(diagonal, torch.inf) > 1e-6) + + def test_underfilled_seeded_sampling_uses_complete_uniform_grid(self): + """Test an exhausted radial sequence becomes a gap-free circular grid.""" + solver = self.solver[next(iter(self.solver))] + solver.cfg.num_samples = 100 + seed = torch.tensor( + [[0.15, -0.35, 0.40, -0.70, 0.20, 0.30, -0.15]], + dtype=torch.float32, + device=solver.device, + ) + + seed_arm_angle = solver.impl._get_seed_arm_angles(seed) + angles = solver.impl._sample_elbow_angles(seed) + offsets = ( + torch.remainder( + angles - seed_arm_angle.unsqueeze(1) + torch.pi, + 2.0 * torch.pi, + ) + - torch.pi + ) + sorted_offsets = offsets.sort(dim=1).values + circular_gaps = torch.cat( + ( + sorted_offsets[:, 1:] - sorted_offsets[:, :-1], + sorted_offsets[:, :1] + 2.0 * torch.pi - sorted_offsets[:, -1:], + ), + dim=1, + ) + + assert angles.shape == (1, solver.cfg.num_samples) + assert torch.allclose(angles[:, 0], seed_arm_angle, atol=1e-6) + assert torch.allclose( + circular_gaps, + torch.full_like(circular_gaps, 2.0 * torch.pi / solver.cfg.num_samples), + atol=1e-6, + ) + + def test_horizontal_shoulder_wrist_seed_recovers_its_fk_pose(self): + """Test horizontal shoulder-wrist geometry retains its shoulder azimuth.""" + seed = torch.tensor( + [[0.0, -np.pi / 2, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], + dtype=torch.float32, + ) + + for solver in self.solver.values(): + device_seed = seed.to(solver.device) + target = solver.get_fk(device_seed) + arm_angle = solver.impl._get_seed_arm_angles(device_seed) + success, solution = solver.get_ik(target, device_seed) + + assert torch.all(torch.isfinite(arm_angle)) + assert torch.all(success) + assert torch.allclose( + solver.get_fk(solution[:, 0]), target, atol=1e-4, rtol=1e-4 + ) + + def test_redundancy_step_larger_than_pi_is_rejected(self): + """Test invalid seeded-search steps fail instead of silently using one sample.""" + solver = self.solver[next(iter(self.solver))] + cfg = solver.cfg.copy() + cfg.redundancy_step = np.pi + 1e-3 + + with pytest.raises(ValueError, match=r"range \(0, pi\]"): + cfg.init_solver(num_envs=1, device=solver.device) + + def test_periodic_joint_values_wrap_inside_limits_near_seed(self): + """Test equivalent revolute values are retained instead of rejected.""" + solver = self.solver[next(iter(self.solver))] + joints = np.array([-3.2, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + limits = np.array([[3.0, 3.2]] + [[-1.0, 1.0]] * 6) + seed = np.array([3.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + + wrapped = solver.impl._wrap_to_limits(joints, limits, seed) + + assert wrapped is not None + assert np.isclose(wrapped[0], -3.2 + 2.0 * np.pi) + assert np.all(wrapped >= limits[:, 0]) + assert np.all(wrapped <= limits[:, 1]) + + def test_periodic_joint_half_turn_tie_matches_cuda_rounding(self): + """Test exact half-turn ties use the same half-up rule as Warp.""" + solver = self.solver[next(iter(self.solver))] + joints = np.zeros(7) + limits = np.array([[-0.1, 2.0 * np.pi + 0.1]] + [[-1.0, 1.0]] * 6) + seed = np.array([np.pi, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + + wrapped = solver.impl._wrap_to_limits(joints, limits, seed) + + assert wrapped is not None + assert np.isclose(wrapped[0], 2.0 * np.pi) + + def test_all_solution_deduplication_is_periodic_and_order_preserving(self): + """Test deduplication retains the first periodic representative.""" + solver = self.solver[next(iter(self.solver))] + first = torch.tensor( + [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7], device=solver.device + ) + second = torch.tensor( + [-0.8, 0.9, -1.0, 1.1, -1.2, 1.3, -1.4], device=solver.device + ) + solutions = torch.stack( + (first, first + 2.0 * torch.pi, second, second - 2.0 * torch.pi) + ) + + unique = solver.impl._deduplicate_solutions(solutions) + + assert unique.shape == (2, 7) + assert torch.allclose(unique[0], first) + assert torch.allclose(unique[1], second) + + def test_all_solution_deduplication_compares_retained_representatives(self): + """Test a discarded chain neighbor cannot remove a later unique row.""" + solver = self.solver[next(iter(self.solver))] + tolerance = 1e-5 + solutions = torch.zeros((3, 7), device=solver.device) + solutions[:, 0] = torch.tensor( + [0.0, 0.9 * tolerance, 1.8 * tolerance], + device=solver.device, + ) + + unique = solver.impl._deduplicate_solutions(solutions, tolerance=tolerance) + + assert unique.shape == (2, 7) + assert torch.equal(unique[0], solutions[0]) + assert torch.equal(unique[1], solutions[2]) + + def test_runtime_tcp_and_weight_updates_reach_backend(self): + """Test mutable solver settings synchronize analytical backend caches.""" + solver = self.solver[next(iter(self.solver))] + tcp = np.eye(4) + tcp[:3, 3] = np.array([0.01, -0.02, 0.03]) + weights = np.arange(1.0, 8.0) + + solver.set_tcp(tcp) + assert np.allclose(solver.impl.tcp_inv_np, np.linalg.inv(tcp)) + + assert solver.set_ik_nearest_weight(weights) + assert np.allclose(solver.cfg.ik_nearest_weight, weights) + if solver.device.type == "cpu": + assert torch.allclose( + solver.impl.ik_nearest_weight_tensor.cpu(), + torch.from_numpy(weights).float(), + ) + else: + assert torch.allclose( + solver.impl.ik_nearest_weight_tensor.cpu(), + torch.from_numpy(weights).float(), + ) + assert torch.allclose( + wp.to_torch(solver.impl.ik_nearest_weight_wp).cpu(), + solver.impl.ik_nearest_weight_tensor.cpu(), + ) # Base test class for CPU and CUDA @@ -210,10 +421,10 @@ def setup_simulation(self, solver_type: str, device: str = "cpu"): cfg_dict = { "fpath": urdf, "control_parts": { - "left_arm": [f"LEFT_J{i+1}" for i in range(7)], - "right_arm": [f"RIGHT_J{i+1}" for i in range(7)], + "left_arm": [f"LEFT_J{i + 1}" for i in range(7)], + "right_arm": [f"RIGHT_J{i + 1}" for i in range(7)], "torso": ["ANKLE", "KNEE", "BUTTOCK", "WAIST"], - "head": [f"NECK{i+1}" for i in range(2)], + "head": [f"NECK{i + 1}" for i in range(2)], }, "drive_pros": { "stiffness": { @@ -342,6 +553,78 @@ class TestSRSCUDASolver(BaseSolverTest): def setup_method(self): self.setup_solver(solver_type="SRSSolver", device="cuda") + def test_cpu_cuda_backend_parity(self): + """Test CPU and CUDA select equivalent solutions for identical inputs.""" + sample_qpos = torch.tensor( + [ + [0.15, -0.35, 0.25, -0.70, 0.20, 0.30, -0.15], + [-0.20, 0.25, -0.30, -0.55, 0.35, -0.20, 0.10], + [0.0, -np.pi / 2, 0.0, -np.pi / 2, 0.0, 0.0, 0.0], + [np.pi / 6, 0.0, 0.0, -np.pi / 2, 0.4, 0.0, np.pi / 6], + ], + dtype=torch.float32, + ) + + for cuda_solver in self.solver.values(): + cpu_solver = cuda_solver.cfg.init_solver( + num_envs=sample_qpos.shape[0], device=torch.device("cpu") + ) + cuda_seed = sample_qpos.to(cuda_solver.device) + target = cuda_solver.get_fk(cuda_seed) + + cuda_arm_angle = cuda_solver.impl._get_seed_arm_angles(cuda_seed) + cpu_arm_angle = cpu_solver.impl._get_seed_arm_angles(sample_qpos) + arm_angle_delta = torch.atan2( + torch.sin(cuda_arm_angle.cpu() - cpu_arm_angle.cpu()), + torch.cos(cuda_arm_angle.cpu() - cpu_arm_angle.cpu()), + ) + assert torch.allclose( + arm_angle_delta, + torch.zeros_like(arm_angle_delta), + atol=1e-4, + rtol=1e-4, + ) + + cuda_success, cuda_qpos = cuda_solver.get_ik(target, cuda_seed) + cpu_success, cpu_qpos = cpu_solver.get_ik(target.cpu(), sample_qpos) + + assert torch.equal(cuda_success.cpu(), cpu_success.cpu()) + assert torch.all(cuda_success) + + cuda_solution = cuda_qpos[:, 0].cpu() + cpu_solution = cpu_qpos[:, 0].cpu() + wrapped_delta = torch.atan2( + torch.sin(cuda_solution - cpu_solution), + torch.cos(cuda_solution - cpu_solution), + ) + dh_params = np.asarray(cuda_solver.cfg.dh_params) + directions = np.asarray(cuda_solver.cfg.rotation_directions) + model_q2 = sample_qpos[:, 1].numpy() * directions[1] + dh_params[1, 3] + model_q6 = sample_qpos[:, 5].numpy() * directions[5] + dh_params[5, 3] + singular = torch.from_numpy( + (np.abs(np.sin(model_q2)) <= 1e-6) | (np.abs(np.sin(model_q6)) <= 1e-6) + ) + assert torch.allclose( + wrapped_delta[~singular], + torch.zeros_like(wrapped_delta[~singular]), + atol=1e-4, + rtol=1e-4, + ) + + # At an Euler singularity, coupled joints can have substantially + # different parameterizations for the same pose. The redundancy + # step bounds arm-angle sampling, not individual joint deltas, so + # singular solutions are compared through FK below. + cuda_reconstructed = cuda_solver.get_fk(cuda_qpos[:, 0]).cpu() + cpu_reconstructed = cpu_solver.get_fk(cpu_qpos[:, 0]).cpu() + assert torch.allclose( + cuda_reconstructed, target.cpu(), atol=1e-4, rtol=1e-4 + ) + assert torch.allclose(cpu_reconstructed, target.cpu(), atol=1e-4, rtol=1e-4) + assert torch.allclose( + cuda_reconstructed, cpu_reconstructed, atol=1e-4, rtol=1e-4 + ) + class TestSRSCPURobotSolver(BaseRobotSolverTest): def setup_method(self):