From a8ab2426f7be6fcfa838dad12e49b2fd5cd96987 Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Mon, 24 Aug 2026 20:21:31 +0800 Subject: [PATCH 01/10] Improve SRS IK arm-angle search and CPU/CUDA parity --- agent_context/topics/ik-solvers/ik-solvers.md | 22 + embodichain/lab/sim/solvers/srs_solver.py | 662 +++++++++++------- .../utils/warp/kinematics/srs_solver.py | 250 +++++-- .../robotics/kinematic_solver/srs_solver.py | 460 ++++++++++++ tests/sim/solvers/test_srs_solver.py | 159 ++++- 5 files changed, 1248 insertions(+), 305 deletions(-) create mode 100644 scripts/benchmark/robotics/kinematic_solver/srs_solver.py diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index dba811321..194011e29 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -151,8 +151,30 @@ 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; `"full"` + samples the complete `[-pi, pi)` interval. +- `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 and periodic nearest-solution formulas. +- 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 deduplicated before return. - 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..7b5f9dcd6 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,92 @@ 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: + raise ValueError("redundancy_step must be finite and positive") + + 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 + 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.rint((seed[index] - value) / two_pi)) + 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: + """Remove periodic-equivalent solutions while preserving their order.""" + if solutions.shape[0] < 2: + return solutions + unique_indices: list[int] = [] + for index in range(solutions.shape[0]): + if not unique_indices: + unique_indices.append(index) + continue + diff = solutions[index] - solutions[unique_indices] + wrapped_diff = torch.atan2(torch.sin(diff), torch.cos(diff)) + if not torch.any(torch.amax(torch.abs(wrapped_diff), dim=1) <= tolerance): + unique_indices.append(index) + return solutions[unique_indices] + class _CPUSRSSolverImpl(_BaseSRSSolverImpl): """CPU implementation of the SRS inverse kinematics solver.""" @@ -139,13 +235,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 +278,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,7 +373,7 @@ 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: joints[0] = np.arctan2(P26[1], P26[0]) @@ -228,7 +381,8 @@ def _calculate_arm_joint_angles( 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 +462,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 + + 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_rotation, joint_angles + return plane_normal, base_to_elbow_pose[:3, :3], joint_angles def _process_all_solutions( self, @@ -336,7 +491,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 +505,32 @@ 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.zeros( + (ik_qpos_tensor.shape[0], max_solutions, 7), + 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 +538,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 +553,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 +580,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 +620,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 +641,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 @@ -492,7 +656,7 @@ def _get_each_ik( # 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 + angle2 = np.arccos(np.clip(R03[2, 1], -1.0, 1.0)) * shoulder_config angle3 = np.arctan2(-R03[2, 2] * shoulder_config, -R03[2, 0] * shoulder_config) # Calculate wrist joint angles @@ -502,7 +666,7 @@ def _get_each_ik( 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 + angle6 = np.arccos(np.clip(R47[2, 2], -1.0, 1.0)) * 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] @@ -513,12 +677,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 +691,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 +704,74 @@ 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): + transformed = ( + self.T_b_ob_inv_np + @ target_xpos_np[target_idx] + @ 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( + xpos, + 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): + target_np = xpos.detach().cpu().numpy() + if np.linalg.norm(fk_xpos - target_np) <= 1e-4: all_solutions[target_idx, sol_idx, :] = qpos sol_idx += 1 solution_counts[target_idx] = sol_idx @@ -598,7 +805,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 +858,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 +940,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 +959,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,36 +976,35 @@ 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, - ) - - 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 + 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_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( @@ -851,7 +1025,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 +1130,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 +1178,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 +1191,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 +1204,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 +1229,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 +1264,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 +1301,7 @@ def get_ik( return ( torch.zeros(num_targets, dtype=torch.bool, device=self.device), torch.zeros( - (num_targets, num_targets, 7), + (num_targets, 7), dtype=torch.float32, device=self.device, ), @@ -1185,7 +1341,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 +1361,32 @@ 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: + self.impl.ik_nearest_weight_wp = wp.from_torch(weights.contiguous()) + 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 +1397,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..ee59e492c 100644 --- a/embodichain/utils/warp/kinematics/srs_solver.py +++ b/embodichain/utils/warp/kinematics/srs_solver.py @@ -42,7 +42,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 @@ -257,19 +270,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 +280,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 +360,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,6 +371,108 @@ def compute_fk_kernel( success[tid] = 1 +@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() + + 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 + + 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) + q1_reference = wp.atan2(shoulder_to_wrist[1], shoulder_to_wrist[0]) + q2_reference = wp.atan2( + wp.length(wp.vec2(shoulder_to_wrist[0], shoulder_to_wrist[1])), + 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 def frobenius_norm(mat: wp.mat44) -> float: """ @@ -426,9 +560,9 @@ def validate_fk_with_target( # 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 +574,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 +583,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 +597,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 +610,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 +634,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 @@ -584,27 +721,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] + ) + 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 (q2_val < qpos_limits[1][0] or q2_val > qpos_limits[1][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 (q3_val < qpos_limits[2][0] or q3_val > qpos_limits[2][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 (q4_val < qpos_limits[3][0] or q4_val > qpos_limits[3][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 (q5_val < qpos_limits[4][0] or q5_val > qpos_limits[4][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 (q6_val < qpos_limits[5][0] or q6_val > qpos_limits[5][1]) else 0 + wrapped_q7 = wrap_to_limit( + q7_val, qpos_limits[6][0], qpos_limits[6][1], qpos_seed[target_idx * 7 + 6] ) - out_of_limits = out_of_limits | ( - 1 if (q7_val < qpos_limits[6][0] or q7_val > qpos_limits[6][1]) else 0 + 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 @@ -677,7 +831,8 @@ def sort_ik_kernel( dist = 0.0 if valid: 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 else: dist = 1e10 @@ -746,7 +901,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/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index 540814987..5f05234a1 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -17,25 +17,25 @@ 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 +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: @@ -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,6 +75,7 @@ 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) @@ -101,9 +102,9 @@ def test_ik(self, arm_side: DexforceW1ArmSide, arm_name: str): ik_xpos = self.solver[arm_name].get_fk(qpos=ik_qpos[:, 0, :]) - assert torch.allclose( - fk_xpos, ik_xpos, atol=1e-3, rtol=1e-3 - ), f"FK and IK results do not match for {arm_name}" + assert torch.allclose(fk_xpos, ik_xpos, atol=1e-3, rtol=1e-3), ( + f"FK and IK results do not match for {arm_name}" + ) def test_update_with_robot_limit_intersects_existing_solver_limits(self): """Test robot limit sync only tightens solver limits and never widens them.""" @@ -174,6 +175,73 @@ def test_update_with_robot_limit_intersects_existing_solver_limits(self): atol=1e-5, ), "FAIL: robot sync did not tighten solver upper_qpos_limits" + 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))] + 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_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_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(), + ) + @classmethod def teardown_class(cls): if cls.solver is not None: @@ -210,10 +278,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": { @@ -303,9 +371,9 @@ def test_robot_ik(self, arm_name: str): else: ik_xpos = self.robot.compute_fk(qpos=ik_qpos, name=arm_name, to_matrix=True) - assert torch.allclose( - fk_xpos, ik_xpos, atol=1e-4, rtol=1e-4 - ), f"FK and IK results do not match for {arm_name}" + assert torch.allclose(fk_xpos, ik_xpos, atol=1e-4, rtol=1e-4), ( + f"FK and IK results do not match for {arm_name}" + ) # test for failed xpos invalid_pose = torch.tensor( @@ -342,6 +410,61 @@ 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], + ], + 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), + ) + assert torch.allclose( + wrapped_delta, + torch.zeros_like(wrapped_delta), + atol=1e-4, + rtol=1e-4, + ) + + 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, cpu_reconstructed, atol=1e-4, rtol=1e-4 + ) + class TestSRSCPURobotSolver(BaseRobotSolverTest): def setup_method(self): From adb11b5b9ef4f6b576860c1fbb49f9f950843c7b Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Mon, 24 Aug 2026 21:03:52 +0800 Subject: [PATCH 02/10] Fix bug and apply style --- embodichain/lab/sim/solvers/srs_solver.py | 25 ++++++------ .../utils/warp/kinematics/srs_solver.py | 2 + tests/sim/solvers/test_srs_solver.py | 40 ++++++++++++++++--- 3 files changed, 49 insertions(+), 18 deletions(-) diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index 7b5f9dcd6..b7d7ce3e1 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -136,8 +136,12 @@ def _parse_params(self): 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: - raise ValueError("redundancy_step must be finite and positive") + 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, @@ -207,16 +211,13 @@ def _deduplicate_solutions( """Remove periodic-equivalent solutions while preserving their order.""" if solutions.shape[0] < 2: return solutions - unique_indices: list[int] = [] - for index in range(solutions.shape[0]): - if not unique_indices: - unique_indices.append(index) - continue - diff = solutions[index] - solutions[unique_indices] - wrapped_diff = torch.atan2(torch.sin(diff), torch.cos(diff)) - if not torch.any(torch.amax(torch.abs(wrapped_diff), dim=1) <= tolerance): - unique_indices.append(index) - return solutions[unique_indices] + pairwise_diff = solutions[:, None, :] - solutions[None, :, :] + wrapped_diff = ( + torch.remainder(pairwise_diff + torch.pi, 2.0 * torch.pi) - torch.pi + ) + equivalent = torch.amax(torch.abs(wrapped_diff), dim=-1) <= tolerance + has_equivalent_predecessor = torch.tril(equivalent, diagonal=-1).any(dim=1) + return solutions[~has_equivalent_predecessor] class _CPUSRSSolverImpl(_BaseSRSSolverImpl): diff --git a/embodichain/utils/warp/kinematics/srs_solver.py b/embodichain/utils/warp/kinematics/srs_solver.py index ee59e492c..9b303f88e 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 diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index 5f05234a1..5f88255c7 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -102,9 +102,9 @@ def test_ik(self, arm_side: DexforceW1ArmSide, arm_name: str): ik_xpos = self.solver[arm_name].get_fk(qpos=ik_qpos[:, 0, :]) - assert torch.allclose(fk_xpos, ik_xpos, atol=1e-3, rtol=1e-3), ( - f"FK and IK results do not match for {arm_name}" - ) + assert torch.allclose( + fk_xpos, ik_xpos, atol=1e-3, rtol=1e-3 + ), f"FK and IK results do not match for {arm_name}" def test_update_with_robot_limit_intersects_existing_solver_limits(self): """Test robot limit sync only tightens solver limits and never widens them.""" @@ -210,6 +210,15 @@ def test_seeded_redundancy_sampling_expands_around_geometric_arm_angle(self): ).unsqueeze(0) assert torch.all(wrapped_delta.masked_fill(diagonal, torch.inf) > 1e-6) + 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))] @@ -224,6 +233,25 @@ def test_periodic_joint_values_wrap_inside_limits_near_seed(self): assert np.all(wrapped >= limits[:, 0]) assert np.all(wrapped <= limits[:, 1]) + def test_all_solution_deduplication_is_periodic_and_order_preserving(self): + """Test vectorized 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_runtime_tcp_and_weight_updates_reach_backend(self): """Test mutable solver settings synchronize analytical backend caches.""" solver = self.solver[next(iter(self.solver))] @@ -371,9 +399,9 @@ def test_robot_ik(self, arm_name: str): else: ik_xpos = self.robot.compute_fk(qpos=ik_qpos, name=arm_name, to_matrix=True) - assert torch.allclose(fk_xpos, ik_xpos, atol=1e-4, rtol=1e-4), ( - f"FK and IK results do not match for {arm_name}" - ) + assert torch.allclose( + fk_xpos, ik_xpos, atol=1e-4, rtol=1e-4 + ), f"FK and IK results do not match for {arm_name}" # test for failed xpos invalid_pose = torch.tensor( From 182ed377ed01159752217989404ef588431196e4 Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Tue, 25 Aug 2026 17:04:00 +0800 Subject: [PATCH 03/10] prof --- agent_context/topics/ik-solvers/ik-solvers.md | 8 +- embodichain/lab/sim/solvers/srs_solver.py | 36 +++- .../utils/warp/kinematics/srs_solver.py | 84 +------- scripts/tutorials/sim/srs_solver.py | 193 +++++++++++++----- tests/sim/solvers/test_srs_solver.py | 65 +++++- 5 files changed, 246 insertions(+), 140 deletions(-) diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index 194011e29..3ec3ab44d 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -157,7 +157,9 @@ class RobotCfg(ArticulationCfg): - `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 and periodic nearest-solution formulas. + 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. - 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 @@ -168,7 +170,9 @@ class RobotCfg(ArticulationCfg): 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 deduplicated before return. +- 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 diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index b7d7ce3e1..34fb70e64 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -208,16 +208,30 @@ def _wrap_to_limits( def _deduplicate_solutions( solutions: torch.Tensor, tolerance: float = 1e-5 ) -> torch.Tensor: - """Remove periodic-equivalent solutions while preserving their order.""" + """Greedily retain periodic-unique representatives in input order.""" if solutions.shape[0] < 2: return solutions - pairwise_diff = solutions[:, None, :] - solutions[None, :, :] - wrapped_diff = ( - torch.remainder(pairwise_diff + torch.pi, 2.0 * torch.pi) - torch.pi + + 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, ) - equivalent = torch.amax(torch.abs(wrapped_diff), dim=-1) <= tolerance - has_equivalent_predecessor = torch.tril(equivalent, diagonal=-1).any(dim=1) - return solutions[~has_equivalent_predecessor] + return solutions[index_tensor] class _CPUSRSSolverImpl(_BaseSRSSolverImpl): @@ -376,12 +390,12 @@ def _calculate_arm_joint_angles( 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_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 @@ -786,7 +800,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, ), @@ -1302,7 +1316,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=torch.float32, device=self.device, ), diff --git a/embodichain/utils/warp/kinematics/srs_solver.py b/embodichain/utils/warp/kinematics/srs_solver.py index 9b303f88e..e664569dd 100644 --- a/embodichain/utils/warp/kinematics/srs_solver.py +++ b/embodichain/utils/warp/kinematics/srs_solver.py @@ -205,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( @@ -430,9 +430,14 @@ def compute_arm_angle_kernel( + wp.pow(distance, 2.0) - wp.pow(link_lengths[2], 2.0) ) / (2.0 * link_lengths[1] * distance) - q1_reference = wp.atan2(shoulder_to_wrist[1], shoulder_to_wrist[0]) + 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( - wp.length(wp.vec2(shoulder_to_wrist[0], shoulder_to_wrist[1])), + horizontal_distance, shoulder_to_wrist[2], ) + elbow_config * safe_acos(shoulder_cosine) @@ -796,79 +801,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): - 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 - 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] diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index 8c6d1bfd6..cbf0c54f6 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -14,17 +14,18 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Move a W1 end effector along a Cartesian straight line with SRS IK.""" + from __future__ import annotations import argparse import time + 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.objects import Robot from embodichain.lab.sim.robots import DexforceW1Cfg from embodichain.lab.visualization import ( VisualizationCfg, @@ -33,67 +34,159 @@ ) -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, + 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. + 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 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, + 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) - - 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)) - - time.sleep(0.5) - - fk_xpos_batch = torch.cat(qpos_fk_list, dim=0) - - fk_xpos_list = robot.compute_fk(qpos=fk_xpos_batch, name=arm_name, to_matrix=True) - - 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" - ) - - 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)) - - ik_xpos_list = robot.compute_fk(qpos=first_solutions, name=arm_name, to_matrix=True) - - print("fk_xpos_list: ", fk_xpos_list) - print("ik_xpos_list: ", ik_xpos_list) + try: + robot: Robot = sim.add_robot( + cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"}) + ) + arm_name = "left_arm" + joint_ids = robot.get_joint_ids(arm_name) + qpos_seed = torch.tensor( + [[0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], + dtype=torch.float32, + device=sim.device, + ) + robot.set_qpos(qpos_seed, joint_ids=joint_ids) + sim.update(step=physics_steps) - sim.capture_visualization(force=True) - embed(header="Test SRSSolver example. Press Ctrl-D to exit.") + 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 + + # 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=False, + ) + if device == "cuda": + torch.cuda.synchronize() + + solve_times_ms: list[float] = [] + translation_errors_mm: list[float] = [] + solved_waypoints = 0 + for waypoint, target_pose in enumerate(target_poses): + start_time = time.perf_counter() + success, solution = robot.compute_ik( + pose=target_pose.unsqueeze(0), + joint_seed=qpos_seed, + name=arm_name, + return_all_solutions=False, + ) + if device == "cuda": + torch.cuda.synchronize() + solve_times_ms.append((time.perf_counter() - start_time) * 1000.0) + + if not bool(success[0]): + print(f"Waypoint {waypoint + 1}/{num_steps}: IK failed") + break + + qpos_seed = solution.reshape(1, 7) + robot.set_qpos(qpos_seed, joint_ids=joint_ids) + # Measure the solver reconstruction before advancing physics so that + # movement of the simulated base is not counted as analytical IK error. + 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) + solved_waypoints += 1 + print( + f"Waypoint {waypoint + 1:03d}/{num_steps}: " + f"solve={solve_times_ms[-1]:.3f} ms, error={error_mm:.4f} mm" + ) + sim.update(step=physics_steps) + + print( + f"SRS {device.upper()} summary: {solved_waypoints}/{num_steps} solved, " + f"mean solve={np.mean(solve_times_ms):.3f} ms, " + f"max translation error=" + f"{max(translation_errors_mm, default=float('nan')):.4f} mm" + ) + sim.capture_visualization(force=True) + finally: + sim.destroy() 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.10, 0.0), + help="TCP line displacement in meters (default: 0 0.10 0).", + ) + parser.add_argument( + "--physics-steps", + type=int, + default=2, + help="Simulation steps displayed per waypoint.", + ) add_viser_args_to_parser(parser) - main(visualization=visualization_cfg_from_args(parser.parse_args())) + args = parser.parse_args() + main( + device=args.device, + num_steps=args.num_steps, + line_offset=tuple(args.line_offset), + physics_steps=args.physics_steps, + visualization=visualization_cfg_from_args(args), + ) diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index 5f88255c7..2f10db015 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -106,6 +106,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)) @@ -210,6 +237,25 @@ def test_seeded_redundancy_sampling_expands_around_geometric_arm_angle(self): ).unsqueeze(0) assert torch.all(wrapped_delta.masked_fill(diagonal, torch.inf) > 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))] @@ -234,7 +280,7 @@ def test_periodic_joint_values_wrap_inside_limits_near_seed(self): assert np.all(wrapped <= limits[:, 1]) def test_all_solution_deduplication_is_periodic_and_order_preserving(self): - """Test vectorized deduplication retains the first periodic representative.""" + """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 @@ -252,6 +298,22 @@ def test_all_solution_deduplication_is_periodic_and_order_preserving(self): 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))] @@ -444,6 +506,7 @@ def test_cpu_cuda_backend_parity(self): [ [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], ], dtype=torch.float32, ) From d4bd61243db4011d75bd2288d0ff7ced0a0e0c44 Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Tue, 25 Aug 2026 20:22:19 +0800 Subject: [PATCH 04/10] modify tutorials --- agent_context/topics/ik-solvers/ik-solvers.md | 3 + embodichain/lab/sim/solvers/srs_solver.py | 20 +- .../utils/warp/kinematics/srs_solver.py | 22 +- scripts/tutorials/sim/srs_solver.py | 277 +++++++++++++++--- tests/sim/solvers/test_srs_solver.py | 1 + 5 files changed, 279 insertions(+), 44 deletions(-) diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index 3ec3ab44d..73d04a042 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -160,6 +160,9 @@ class RobotCfg(ArticulationCfg): 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 diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index 34fb70e64..f95940395 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -670,9 +670,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(np.clip(R03[2, 1], -1.0, 1.0)) * shoulder_config - angle3 = np.arctan2(-R03[2, 2] * shoulder_config, -R03[2, 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 @@ -680,9 +688,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(np.clip(R47[2, 2], -1.0, 1.0)) * wrist_config - angle7 = np.arctan2(R47[2, 1] * wrist_config, -R47[2, 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] diff --git a/embodichain/utils/warp/kinematics/srs_solver.py b/embodichain/utils/warp/kinematics/srs_solver.py index e664569dd..4d187d330 100644 --- a/embodichain/utils/warp/kinematics/srs_solver.py +++ b/embodichain/utils/warp/kinematics/srs_solver.py @@ -701,10 +701,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 @@ -713,10 +718,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) diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index cbf0c54f6..4be55f561 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -19,12 +19,15 @@ from __future__ import annotations import argparse +import sys import time +import traceback import numpy as np import torch 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 ( @@ -39,6 +42,7 @@ def main( 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. @@ -48,12 +52,15 @@ def main( 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" @@ -63,7 +70,9 @@ def main( torch.set_printoptions(precision=5, sci_mode=False) sim = SimulationManager( SimulationManagerCfg( - headless=False, + # 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, @@ -75,17 +84,24 @@ def main( sim.set_manual_update(True) try: - robot: Robot = sim.add_robot( - cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"}) - ) 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( - [[0.0, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, 0.0]], + [[np.pi / 6, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, np.pi / 6]], dtype=torch.float32, device=sim.device, ) - robot.set_qpos(qpos_seed, joint_ids=joint_ids) + 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) start_pose = robot.compute_fk(qpos=qpos_seed, name=arm_name, to_matrix=True) @@ -103,34 +119,116 @@ def main( pose=target_poses[:1], joint_seed=qpos_seed, name=arm_name, - return_all_solutions=False, + return_all_solutions=True, ) if device == "cuda": torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() solve_times_ms: list[float] = [] + candidate_counts: list[int] = [] translation_errors_mm: list[float] = [] - solved_waypoints = 0 + 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), - joint_seed=qpos_seed, + # 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=False, + return_all_solutions=True, ) if device == "cuda": torch.cuda.synchronize() solve_times_ms.append((time.perf_counter() - start_time) * 1000.0) if not bool(success[0]): - print(f"Waypoint {waypoint + 1}/{num_steps}: IK failed") - break + raise RuntimeError( + f"Trajectory planning failed at waypoint " + f"{waypoint + 1}/{num_steps}; nothing was executed." + ) + + 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]}" + ) + + # 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)) + + 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) - qpos_seed = solution.reshape(1, 7) - robot.set_qpos(qpos_seed, joint_ids=joint_ids) - # Measure the solver reconstruction before advancing physics so that - # movement of the simulated base is not counted as analytical IK error. + 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) + ] + + 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 ) @@ -141,22 +239,116 @@ def main( * 1000.0 ) translation_errors_mm.append(error_mm) - solved_waypoints += 1 print( - f"Waypoint {waypoint + 1:03d}/{num_steps}: " - f"solve={solve_times_ms[-1]:.3f} ms, error={error_mm:.4f} mm" + f"Planned {waypoint + 1:03d}/{num_steps}: " f"error={error_mm:.4f} mm" ) - sim.update(step=physics_steps) print( - f"SRS {device.upper()} summary: {solved_waypoints}/{num_steps} solved, " - f"mean solve={np.mean(solve_times_ms):.3f} ms, " + 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: - sim.destroy() + # 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__": @@ -172,21 +364,38 @@ def main( type=float, nargs=3, metavar=("X", "Y", "Z"), - default=(0.0, 0.10, 0.0), - help="TCP line displacement in meters (default: 0 0.10 0).", + 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=2, - help="Simulation steps displayed per waypoint.", + 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) args = parser.parse_args() - main( - device=args.device, - num_steps=args.num_steps, - line_offset=tuple(args.line_offset), - physics_steps=args.physics_steps, - visualization=visualization_cfg_from_args(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 2f10db015..ab3ab3950 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -507,6 +507,7 @@ def test_cpu_cuda_backend_parity(self): [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, ) From abd62853db26f42f5387af6479ef9e3174dfb7e2 Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Tue, 25 Aug 2026 20:35:35 +0800 Subject: [PATCH 05/10] fix bug --- agent_context/topics/ik-solvers/ik-solvers.md | 6 ++- embodichain/lab/sim/solvers/srs_solver.py | 41 +++++++++++++++---- tests/sim/solvers/test_srs_solver.py | 2 +- 3 files changed, 39 insertions(+), 10 deletions(-) diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index 73d04a042..9da34e822 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -152,8 +152,10 @@ 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; `"full"` - samples the complete `[-pi, pi)` interval. + arm angle and searches redundancy angles radially around it; if the configured + radial step cannot produce `num_samples` distinct angles in one revolution, + uncovered intermediate angles from a uniform full-circle grid fill the remaining + slots. `"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 diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index f95940395..bffe17283 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -175,6 +175,36 @@ def _sample_elbow_angles( 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, retain the + # seed-first radial prefix and fill the remaining slots from a uniform + # full-circle grid. This preserves the configured local search order + # while honoring num_samples and covering intermediate arm angles. + if len(offsets) < self.cfg.num_samples: + uniform_offsets = np.linspace( + -np.pi, + np.pi, + self.cfg.num_samples, + endpoint=False, + dtype=np.float64, + ) + uniform_offsets = sorted( + uniform_offsets, + key=lambda value: abs((float(value) + np.pi) % (2.0 * np.pi) - np.pi), + ) + for candidate in uniform_offsets: + wrapped_candidate = (float(candidate) + np.pi) % (2.0 * np.pi) - np.pi + if any( + abs((wrapped_candidate - existing + np.pi) % (2.0 * np.pi) - np.pi) + <= 1e-12 + for existing in offsets + ): + continue + offsets.append(wrapped_candidate) + if len(offsets) == self.cfg.num_samples: + break + 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) @@ -758,12 +788,10 @@ def get_ik( 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_xpos_np[target_idx] - @ self.tcp_inv_np - @ self.T_e_oe_inv_np + 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]]) @@ -789,7 +817,7 @@ def get_ik( if prepared is None: continue success, qpos = self._get_each_ik( - xpos, + target_np, psi.item(), config, qpos_seed_np[target_idx], @@ -797,7 +825,6 @@ def get_ik( ) if success: fk_xpos = self._get_fk(qpos) - target_np = xpos.detach().cpu().numpy() if np.linalg.norm(fk_xpos - target_np) <= 1e-4: all_solutions[target_idx, sol_idx, :] = qpos sol_idx += 1 diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index ab3ab3950..fd3899c43 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -225,7 +225,7 @@ def test_seeded_redundancy_sampling_expands_around_geometric_arm_angle(self): ) 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 angles.shape[1] == solver.cfg.num_samples assert torch.all(angles >= -torch.pi) assert torch.all(angles < torch.pi) wrapped_delta = torch.atan2( From 14d7254a3d24a2d32c24c7808a66cfda0d0acb4d Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Wed, 26 Aug 2026 11:34:33 +0800 Subject: [PATCH 06/10] fix test_srs_solver --- tests/sim/solvers/test_srs_solver.py | 33 ++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index fd3899c43..d4a84ad5f 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -48,6 +48,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, @@ -334,9 +335,10 @@ def test_runtime_tcp_and_weight_updates_reach_backend(self): @classmethod def teardown_class(cls): - if cls.solver is not None: + solver = getattr(cls, "solver", None) + if solver is not None: try: - del cls.solver + solver.clear() print("solver destroyed successfully") except Exception as e: print(f"Error during solver destruction: {e}") @@ -544,15 +546,38 @@ def test_cpu_cuda_backend_parity(self): 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, - torch.zeros_like(wrapped_delta), + wrapped_delta[~singular], + torch.zeros_like(wrapped_delta[~singular]), atol=1e-4, rtol=1e-4, ) + # Euler singularities have a free coupled joint, so backends may + # choose different nearby parameterizations after full sample-space + # completion. Both choices must remain close to the seed. + for solution in (cuda_solution, cpu_solution): + seed_delta = torch.atan2( + torch.sin(solution[singular] - sample_qpos[singular]), + torch.cos(solution[singular] - sample_qpos[singular]), + ) + assert torch.all( + seed_delta.abs() <= cuda_solver.cfg.redundancy_step + 1e-4 + ) + 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 ) From e53a7418bc3422b0ac323e5faf260d5e29a87834 Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Thu, 27 Aug 2026 19:34:41 +0800 Subject: [PATCH 07/10] remove teardown_class --- agent_context/topics/ik-solvers/ik-solvers.md | 4 +- embodichain/lab/sim/solvers/srs_solver.py | 39 +++++--------- tests/sim/solvers/test_srs_solver.py | 53 ++++++++++++++----- 3 files changed, 56 insertions(+), 40 deletions(-) diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index 9da34e822..49e4a9620 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -154,8 +154,8 @@ class RobotCfg(ArticulationCfg): - `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, - uncovered intermediate angles from a uniform full-circle grid fill the remaining - slots. `"full"` samples the complete `[-pi, pi)` interval directly. + 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 diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index bffe17283..1469e1c60 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -177,33 +177,20 @@ def _sample_elbow_angles( layer += 1 # A fixed radial step has only finitely many distinct positions on one - # revolution. If it cannot provide num_samples entries, retain the - # seed-first radial prefix and fill the remaining slots from a uniform - # full-circle grid. This preserves the configured local search order - # while honoring num_samples and covering intermediate arm angles. + # 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: - uniform_offsets = np.linspace( - -np.pi, - np.pi, - self.cfg.num_samples, - endpoint=False, - dtype=np.float64, - ) - uniform_offsets = sorted( - uniform_offsets, - key=lambda value: abs((float(value) + np.pi) % (2.0 * np.pi) - np.pi), - ) - for candidate in uniform_offsets: - wrapped_candidate = (float(candidate) + np.pi) % (2.0 * np.pi) - np.pi - if any( - abs((wrapped_candidate - existing + np.pi) % (2.0 * np.pi) - np.pi) - <= 1e-12 - for existing in offsets - ): - continue - offsets.append(wrapped_candidate) - if len(offsets) == self.cfg.num_samples: - break + 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) diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index d4a84ad5f..e7c16230e 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -39,8 +39,6 @@ class BaseSolverTest: - solver = {} - def get_arm_config(self): return [ (DexforceW1ArmSide.LEFT, "left_arm"), @@ -80,6 +78,10 @@ def setup_solver(self, solver_type: str, device: str = "cpu"): 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", [ @@ -206,6 +208,7 @@ def test_update_with_robot_limit_intersects_existing_solver_limits(self): 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, @@ -238,6 +241,42 @@ def test_seeded_redundancy_sampling_expands_around_geometric_arm_angle(self): ).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( @@ -333,16 +372,6 @@ def test_runtime_tcp_and_weight_updates_reach_backend(self): torch.from_numpy(weights).float(), ) - @classmethod - def teardown_class(cls): - solver = getattr(cls, "solver", None) - if solver is not None: - try: - solver.clear() - print("solver destroyed successfully") - except Exception as e: - print(f"Error during solver destruction: {e}") - # Base test class for CPU and CUDA class BaseRobotSolverTest: From 1866a44e0a14951bfb5adf4893b793fb6716e64c Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Fri, 28 Aug 2026 10:38:07 +0800 Subject: [PATCH 08/10] fix test srs solver --- tests/sim/solvers/test_srs_solver.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index e7c16230e..6608f6a7c 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -589,18 +589,10 @@ def test_cpu_cuda_backend_parity(self): rtol=1e-4, ) - # Euler singularities have a free coupled joint, so backends may - # choose different nearby parameterizations after full sample-space - # completion. Both choices must remain close to the seed. - for solution in (cuda_solution, cpu_solution): - seed_delta = torch.atan2( - torch.sin(solution[singular] - sample_qpos[singular]), - torch.cos(solution[singular] - sample_qpos[singular]), - ) - assert torch.all( - seed_delta.abs() <= cuda_solver.cfg.redundancy_step + 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( From 919b3492de84034af22bce2bf6bbe376f64615c0 Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Fri, 28 Aug 2026 12:01:20 +0800 Subject: [PATCH 09/10] fix test --- embodichain/lab/sim/solvers/srs_solver.py | 10 +++--- .../utils/warp/kinematics/srs_solver.py | 32 +++++-------------- tests/sim/solvers/test_srs_solver.py | 12 +++++++ 3 files changed, 26 insertions(+), 28 deletions(-) diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index 1469e1c60..fde5a6643 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -216,7 +216,7 @@ def _wrap_to_limits( k_max = int(np.floor((upper - value) / two_pi)) if k_min > k_max: return None - nearest_k = int(np.rint((seed[index] - value) / two_pi)) + 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 @@ -555,8 +555,9 @@ def _process_all_solutions( for index in range(ik_qpos_tensor.shape[0]) ] max_solutions = max(solution.shape[0] for solution in valid_qpos) - compact_qpos = torch.zeros( + compact_qpos = torch.full( (ik_qpos_tensor.shape[0], max_solutions, 7), + float("nan"), dtype=ik_qpos_tensor.dtype, device=self.device, ) @@ -812,7 +813,7 @@ def get_ik( ) if success: fk_xpos = self._get_fk(qpos) - if np.linalg.norm(fk_xpos - target_np) <= 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 @@ -1048,8 +1049,9 @@ def _process_all_solutions( 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, ) diff --git a/embodichain/utils/warp/kinematics/srs_solver.py b/embodichain/utils/warp/kinematics/srs_solver.py index 4d187d330..e08d31965 100644 --- a/embodichain/utils/warp/kinematics/srs_solver.py +++ b/embodichain/utils/warp/kinematics/srs_solver.py @@ -480,24 +480,6 @@ def compute_arm_angle_kernel( success[tid] = 1 -@wp.func -def frobenius_norm(mat: wp.mat44) -> float: - """ - Compute the Frobenius norm of a 4x4 matrix. - - Args: - mat (wp.mat44): Input matrix. - - 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) - - @wp.func def validate_fk_with_target( q1: float, @@ -556,12 +538,14 @@ 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 diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index 6608f6a7c..d2ec4eaa6 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -319,6 +319,18 @@ def test_periodic_joint_values_wrap_inside_limits_near_seed(self): 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))] From 65568631149ba0b99abfc80ab38cb862493a5f98 Mon Sep 17 00:00:00 2001 From: Jietao Chen Date: Fri, 28 Aug 2026 12:06:30 +0800 Subject: [PATCH 10/10] save to ik_nearest_weight_tensor --- embodichain/lab/sim/solvers/srs_solver.py | 7 ++++++- tests/sim/solvers/test_srs_solver.py | 10 ++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index fde5a6643..aad8cce48 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -1427,7 +1427,12 @@ def set_ik_nearest_weight( if isinstance(self.impl, _CPUSRSSolverImpl): self.impl.ik_nearest_weight_tensor = weights else: - self.impl.ik_nearest_weight_wp = wp.from_torch(weights.contiguous()) + # ``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): diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index d2ec4eaa6..f760ed616 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -23,6 +23,7 @@ torch._dynamo.config.cache_size_limit = 128 # recompile_limit import numpy as np import pytest +import warp as wp from embodichain.data import get_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -383,6 +384,15 @@ def test_runtime_tcp_and_weight_updates_reach_backend(self): 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