From d155b4592298758e17a70d27632759662a17b282 Mon Sep 17 00:00:00 2001 From: matafela Date: Thu, 10 Sep 2026 11:03:51 +0800 Subject: [PATCH] fix ur solver allocate memory --- agent_context/topics/ik-solvers/ik-solvers.md | 9 + .../embodichain.lab.sim.motion.solvers.rst | 21 + embodichain/compute/kinematics/_warp/ur.py | 149 ++++++-- .../lab/sim/motion/solvers/ur_solver.py | 74 ++-- .../robotics/kinematic_solver/ur_solver.py | 359 ++++++++++++++++++ tests/sim/motion/solvers/test_ur_solver.py | 249 +++++++++++- 6 files changed, 813 insertions(+), 48 deletions(-) create mode 100644 scripts/benchmark/robotics/kinematic_solver/ur_solver.py diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index bff1d9c56..99aaf4c63 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -130,5 +130,14 @@ OPW, SRS, and UR Warp implementations live in `lab.sim.motion.solvers` classes own configuration, state, device buffers, and solver interfaces. The compute kernels do not import simulation modules. `utils/warp/kinematics/*_solver.py` are compatibility aliases. + +`URSolver.get_ik(return_all_solutions=False)` uses `ur_ik_nearest_kernel` to +generate, validate and select the seed-weighted nearest candidate in local +storage, returning joints `(N, 6)` and validity `(N,)`. It retains the eight +analytical branches and 64 periodic combinations per branch without writing +the full candidate tensor. `return_all_solutions=True` uses `ur_ik_kernel` and +preserves the ordered joints `(N, 512, 6)` and validity `(N, 512)`, including +repeated representatives when no shifted value fits the joint limits. + Validate kernel import/compilation with `tests/compute/test_imports.py` and solver behavior with the corresponding `tests/sim/motion/solvers/` tests. diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst index c8e58bd27..eb09b1b1f 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst @@ -129,6 +129,27 @@ SRS Solver UR Solver --------- +``URSolverCfg`` selects the analytical DH parameters for UR3, UR5, UR10 and +their e-series variants. ``URSolver.get_ik()`` returns validity flags followed +by joint positions. By default, a dedicated Warp kernel selects the nearest +valid candidate using the joint seed and per-joint weights, returning shapes +``(N,)`` and ``(N, 6)`` without allocating the complete candidate tensor. + +With ``return_all_solutions=True``, the solver returns shapes ``(N, 512)`` and +``(N, 512, 6)``. These candidates contain eight analytical branches expanded +over 64 combinations of periodic joint representatives, with validity flags +for FK agreement and joint limits. When no shifted representative fits a +joint's limits, its base value is repeated. + +.. currentmodule:: embodichain.lab.sim.motion.solvers.ur_solver + +.. autosummary:: + + URSolverCfg + URSolver + +.. currentmodule:: embodichain.lab.sim.motion.solvers + .. autoclass:: URSolverCfg :members: :exclude-members: __init__, copy, replace, to_dict, validate diff --git a/embodichain/compute/kinematics/_warp/ur.py b/embodichain/compute/kinematics/_warp/ur.py index d886750c5..08272fab6 100644 --- a/embodichain/compute/kinematics/_warp/ur.py +++ b/embodichain/compute/kinematics/_warp/ur.py @@ -16,10 +16,20 @@ from __future__ import annotations -from typing import Tuple +from typing import Any, Tuple import warp as wp +__all__ = [ + "wp_vec6f", + "wp_vec48f", + "normalize_to_pi", + "URParam", + "ur_single_fk", + "ur_ik_kernel", + "ur_ik_nearest_kernel", +] + wp_vec6f = wp.types.vector(length=6, dtype=float) wp_vec48f = wp.types.vector(length=48, dtype=float) @@ -227,31 +237,12 @@ def _shift_to_limit(q: float, lo: float, hi: float) -> float: return q -@wp.kernel -def ur_ik_kernel( - xpos: wp.array(dtype=float), # [n_sample * 16] row-major 4x4 target poses - params: URParam, - lower_qpos_limits_wp: wp.array(dtype=float), # [6] lower joint limits - upper_qpos_limits_wp: wp.array(dtype=float), # [6] upper joint limits - qpos: wp.array(dtype=float), # [n_sample * 512 * DOF] output joint solutions - ik_valid: wp.array(dtype=int), # [n_sample * 512] output validity flags -): - """Compute expanded analytical IK solutions for a batch of UR poses. - - Each thread handles one target pose. The 8 base analytical solutions are - expanded to ``8 * 2**6 = 512`` candidates: for every base solution and every - joint, a second FK-equivalent value shifted by +/- 2*pi is generated when it - falls inside the joint limits (UR joints are 2*pi-periodic, so this preserves - the end-effector pose). When no shifted value fits the limits, the joint's own - value is repeated. Each candidate is flagged valid only if the base FK matches - the target *and* every joint lies within its limits. - """ - i = wp.tid() - DOF = int(6) - N_SOL = int(8) - N_SHIFT = int(64) # 2**6 per-joint +/- 2*pi shift combinations +@wp.func +def _ur_ik_branches( + xpos: wp.array(dtype=float), params: URParam, i: int +) -> Tuple[wp_vec48f, wp.mat44f]: + """Compute the eight ordered analytical branches and their target pose.""" base = i * 16 - # Load rotation and translation from the row-major 4x4 target pose. r11 = xpos[base + 0] r12 = xpos[base + 1] @@ -424,6 +415,34 @@ def ur_ik_kernel( t6_5nb, # sol 7 ) + return theta, target_pose + + +@wp.kernel +def ur_ik_kernel( + xpos: wp.array(dtype=float), # [n_sample * 16] row-major 4x4 target poses + params: URParam, + lower_qpos_limits_wp: wp.array(dtype=float), # [6] lower joint limits + upper_qpos_limits_wp: wp.array(dtype=float), # [6] upper joint limits + qpos: wp.array(dtype=float), # [n_sample * 512 * DOF] output joint solutions + ik_valid: wp.array(dtype=int), # [n_sample * 512] output validity flags +): + """Compute expanded analytical IK solutions for a batch of UR poses. + + Each thread handles one target pose. The 8 base analytical solutions are + expanded to ``8 * 2**6 = 512`` candidates: for every base solution and every + joint, a second FK-equivalent value shifted by +/- 2*pi is generated when it + falls inside the joint limits (UR joints are 2*pi-periodic, so this preserves + the end-effector pose). When no shifted value fits the limits, the joint's own + value is repeated. Each candidate is flagged valid only if the base FK matches + the target *and* every joint lies within its limits. + """ + i = wp.tid() + DOF = int(6) + N_SOL = int(8) + N_SHIFT = int(64) # 2**6 per-joint +/- 2*pi shift combinations + theta, target_pose = _ur_ik_branches(xpos, params, i) + # Expand each of the 8 base solutions into 2**6 = 64 per-joint +/- 2*pi shift # variants, yielding 8 * 64 = 512 candidates total. Shifting is FK-equivalent # and only applied when it lands inside the joint limits; otherwise the joint's @@ -498,3 +517,83 @@ def ur_ik_kernel( ): valid = int(0) ik_valid[i * N_SOL * N_SHIFT + j * N_SHIFT + k] = valid + + +@wp.kernel(enable_backward=False) +def ur_ik_nearest_kernel( + xpos: wp.array(dtype=float), + params: URParam, + lower_limits: wp.array(dtype=float), + upper_limits: wp.array(dtype=float), + qpos_seed: wp.array(dtype=Any, ndim=2), + joint_weights: wp.array(dtype=Any), + qpos: wp.array(dtype=float, ndim=2), + ik_valid: wp.array(dtype=int), +): + """Generate, validate and select one UR solution per target in local storage. + + Args: + xpos: Flattened target transforms, shape ``(N * 16,)``. + params: UR Denavit-Hartenberg parameters. + lower_limits: Six lower joint limits in radians. + upper_limits: Six upper joint limits in radians. + qpos_seed: Joint seeds, shape ``(N, 6)``. + joint_weights: Six weights, with the same dtype as the seeds. + qpos: Output joint values, shape ``(N, 6)``. + ik_valid: Output validity flags, shape ``(N,)``. + + Candidates follow the all-solutions kernel's branch and periodic order. + Strict improvement retains the first equal-distance candidate. When all + candidates are invalid, the first candidate is returned with a false flag, + matching the previous masked ``torch.norm`` / ``argmin`` selection. + """ + i = wp.tid() + theta, target_pose = _ur_ik_branches(xpos, params, i) + best_q = wp_vec6f() + best_valid = int(0) + best_distance = qpos_seed.dtype(wp.inf) + tol = float(1e-9) + + for j in range(8): + base_q = wp_vec6f() + shifted_q = wp_vec6f() + for t in range(6): + base_q[t] = normalize_to_pi(theta[j * 6 + t]) + shifted_q[t] = _shift_to_limit(base_q[t], lower_limits[t], upper_limits[t]) + + fk_result = ur_single_fk( + base_q[0], base_q[1], base_q[2], base_q[3], base_q[4], base_q[5], params + ) + t_err, r_err = _ur_transform_err(fk_result, target_pose) + fk_ok = int(1) + if t_err > float(1e-2) or r_err > float(1e-1): + fk_ok = int(0) + + for k in range(64): + candidate = wp_vec6f() + valid = fk_ok + squared_distance = qpos_seed.dtype(0.0) + for t in range(6): + q = base_q[t] if (k & (1 << t)) == 0 else shifted_q[t] + candidate[t] = q + if q < lower_limits[t] - tol or q > upper_limits[t] + tol: + valid = int(0) + delta = (qpos_seed.dtype(q) - qpos_seed[i, t]) * joint_weights[t] + squared_distance = squared_distance + delta * delta + + if j == 0 and k == 0: + best_q = candidate + best_valid = valid + # Retain the norm (including its rounding) for equal-distance ties. + distance = wp.sqrt(squared_distance) + if valid != 0 and ( + distance < best_distance + or (wp.isnan(distance) and not wp.isnan(best_distance)) + ): + best_distance = distance + best_q = candidate + best_valid = valid + + for t in range(6): + qpos[i, t] = best_q[t] + ik_valid[i] = best_valid diff --git a/embodichain/lab/sim/motion/solvers/ur_solver.py b/embodichain/lab/sim/motion/solvers/ur_solver.py index 127e04074..05396db31 100644 --- a/embodichain/lab/sim/motion/solvers/ur_solver.py +++ b/embodichain/lab/sim/motion/solvers/ur_solver.py @@ -25,10 +25,13 @@ from embodichain.compute.kinematics._warp.ur import ( URParam, ur_ik_kernel, + ur_ik_nearest_kernel, ) import math from embodichain.utils.device_utils import standardize_device_string +__all__ = ["URSolverCfg", "URSolver"] + @configclass class URSolverCfg(SolverCfg): @@ -140,19 +143,24 @@ def get_ik( qpos_seed: torch.Tensor | None = None, return_all_solutions: bool = False, **kwargs, - ): - """Compute target joint positions using OPW inverse kinematics. + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute target joint positions using UR inverse kinematics. Args: target_xpos (torch.Tensor): Current end-effector pose, shape (n_sample, 4, 4). - qpos_seed (torch.Tensor): Current joint positions, shape (n_sample, num_joints). - return_all_solutions (bool, optional): Whether to return all IK solutions or just the best one. Defaults to False. + qpos_seed (torch.Tensor): Current joint positions, shape (n_sample, 6) + or (1, 6). Defaults to the joint-limit midpoint. + return_all_solutions (bool, optional): Whether to return all 512 + candidates. False selects the nearest valid candidate inside + the Warp kernel without allocating the full candidate tensor. **kwargs: Additional keyword arguments for future extensions. Returns: Tuple[torch.Tensor, torch.Tensor]: - - target_joints (torch.Tensor): Computed target joint positions, shape (n_sample, n_solution, num_joints). - - success (torch.Tensor): Boolean tensor indicating IK solution validity for each environment, shape (n_sample,). + - success (torch.Tensor): Boolean validity, shape (n_sample,) + or (n_sample, 512) when all solutions are requested. + - target_joints (torch.Tensor): Joint positions, shape + (n_sample, 6) or (n_sample, 512, 6), respectively. """ N_SOL = 512 DOF = 6 @@ -164,7 +172,7 @@ def get_ik( target_xpos_batch = target_xpos_batch @ tcp_inv[None, :, :] n_sample = target_xpos_batch.shape[0] - if qpos_seed is None: + if qpos_seed is None and not return_all_solutions: # A missing seed previously crashed at the nearest-solution step; # default to the feasibility-safe joint-range midpoint. qpos_seed = ( @@ -178,10 +186,44 @@ def get_ik( wp_device = standardize_device_string(self.device) # Flatten target poses to a 1-D float array for the Warp kernel. xpos_wp = wp.from_torch(target_xpos_batch.reshape(-1)) - all_qpos_wp = wp.zeros(n_sample * N_SOL * DOF, dtype=float, device=wp_device) - all_ik_valid_wp = wp.zeros(n_sample * N_SOL, dtype=int, device=wp_device) lower_qpos_limits_wp = wp.from_torch(self.lower_qpos_limits) upper_qpos_limits_wp = wp.from_torch(self.upper_qpos_limits) + + if not return_all_solutions: + # Match PyTorch's promotion in weight * (float32 candidates - seed), + # including double-precision seeds and broadcast/noncontiguous seeds. + selection_dtype = torch.promote_types(torch.float32, qpos_seed.dtype) + selection_dtype = torch.promote_types( + selection_dtype, self.ik_nearest_weight.dtype + ) + seed = ( + qpos_seed.to(device=device, dtype=selection_dtype) + .expand(n_sample, DOF) + .contiguous() + ) + weights = self.ik_nearest_weight.to( + device=device, dtype=selection_dtype + ).contiguous() + best_qpos_wp = wp.empty((n_sample, DOF), dtype=float, device=wp_device) + best_valid_wp = wp.empty(n_sample, dtype=int, device=wp_device) + wp.launch( + kernel=ur_ik_nearest_kernel, + dim=n_sample, + inputs=[ + xpos_wp, + self._ur_params, + lower_qpos_limits_wp, + upper_qpos_limits_wp, + wp.from_torch(seed), + wp.from_torch(weights), + ], + outputs=[best_qpos_wp, best_valid_wp], + device=wp_device, + ) + return wp.to_torch(best_valid_wp).bool(), wp.to_torch(best_qpos_wp) + + all_qpos_wp = wp.zeros(n_sample * N_SOL * DOF, dtype=float, device=wp_device) + all_ik_valid_wp = wp.zeros(n_sample * N_SOL, dtype=int, device=wp_device) wp.launch( kernel=ur_ik_kernel, dim=(n_sample,), @@ -207,19 +249,7 @@ def get_ik( .to(device=device) ) - if return_all_solutions: - return all_solutions_validity, all_solutions - # Select ik qpos based on the closest distance to the seed qpos - qpos_seed_expanded = qpos_seed.unsqueeze(1).expand(-1, N_SOL, -1) - distances = torch.norm( - self.ik_nearest_weight * (all_solutions - qpos_seed_expanded), dim=-1 - ) - # fill invalid solutions with inf distance - distances[~all_solutions_validity] = float("inf") - closest_indices = torch.argmin(distances, dim=1) - ik_qpos = all_solutions[torch.arange(n_sample), closest_indices] - ik_validity = all_solutions_validity[torch.arange(n_sample), closest_indices] - return ik_validity, ik_qpos + return all_solutions_validity, all_solutions @staticmethod def dh_matrix(theta_i, d_i, a_i, alpha_i): diff --git a/scripts/benchmark/robotics/kinematic_solver/ur_solver.py b/scripts/benchmark/robotics/kinematic_solver/ur_solver.py new file mode 100644 index 000000000..c2183d007 --- /dev/null +++ b/scripts/benchmark/robotics/kinematic_solver/ur_solver.py @@ -0,0 +1,359 @@ +# ---------------------------------------------------------------------------- +# 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 UR single-solution, legacy selection, and all-solution IK paths. + +Run: python -m scripts.benchmark.robotics.kinematic_solver.ur_solver + +Targets use independent DH forward kinematics, avoiding robot assets, simulation, +and compiled FK setup. Timings exclude warmup and include CUDA synchronization. +Memory reports distinguish measured PyTorch allocation from the logical Warp +output-buffer payload; neither is a measurement of total device peak memory. +""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time +from collections.abc import Callable +from datetime import datetime +from pathlib import Path + +import psutil +import torch +import warp as wp + +from embodichain.lab.sim.motion.solvers import URSolver, URSolverCfg + +__all__ = ["run_all_benchmarks"] + +_MODES = ("single", "legacy", "all") +_MIB = 1024**2 +_Result = tuple[torch.Tensor, torch.Tensor] + + +def _parse_args() -> argparse.Namespace: + """Parse workload, repetition, and report controls.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda") + parser.add_argument("--counts", nargs="+", type=int, default=[1000, 10000, 100000]) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--seed", type=int, default=20260910) + parser.add_argument("--output-dir", type=Path, default=Path("outputs/benchmarks")) + parser.add_argument("--json", type=Path, help="Optional JSON results path.") + args = parser.parse_args() + if min(args.counts) < 1 or args.repeats < 1 or args.warmup < 1: + parser.error("counts, repeats, and warmup must all be positive") + return args + + +def _synchronize(device: torch.device) -> None: + """Wait for both Torch and Warp work on the requested CUDA device.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def _forward_kinematics(qpos: torch.Tensor, cfg: URSolverCfg) -> torch.Tensor: + """Compute batch UR DH transforms independently from the IK implementation.""" + qpos = qpos.to(dtype=torch.float64) + poses = torch.eye(4, dtype=qpos.dtype, device=qpos.device).repeat(len(qpos), 1, 1) + parameters = ( + (cfg.d1, 0.0, math.pi / 2), + (0.0, cfg.a2, 0.0), + (0.0, cfg.a3, 0.0), + (cfg.d4, 0.0, math.pi / 2), + (cfg.d5, 0.0, -math.pi / 2), + (cfg.d6, 0.0, 0.0), + ) + for joint, (d, a, alpha) in enumerate(parameters): + c, s = torch.cos(qpos[:, joint]), torch.sin(qpos[:, joint]) + ca, sa = math.cos(alpha), math.sin(alpha) + transform = torch.zeros_like(poses) + transform[:, 0, 0] = c + transform[:, 0, 1] = -s * ca + transform[:, 0, 2] = s * sa + transform[:, 0, 3] = a * c + transform[:, 1, 0] = s + transform[:, 1, 1] = c * ca + transform[:, 1, 2] = -c * sa + transform[:, 1, 3] = a * s + transform[:, 2, 1] = sa + transform[:, 2, 2] = ca + transform[:, 2, 3] = d + transform[:, 3, 3] = 1.0 + poses = poses @ transform + return poses + + +def _solve( + solver: URSolver, target: torch.Tensor, seed: torch.Tensor, mode: str +) -> _Result: + """Run one API mode or reconstruct the previous Torch nearest selection.""" + valid, joints = solver.get_ik( + target, qpos_seed=seed, return_all_solutions=mode != "single" + ) + if mode != "legacy": + return valid, joints + distances = torch.norm( + solver.ik_nearest_weight * (joints - seed.unsqueeze(1)), dim=-1 + ) + distances[~valid] = float("inf") + indices = distances.argmin(dim=1) + rows = torch.arange(len(target), device=target.device) + return valid[rows, indices], joints[rows, indices] + + +def _measure( + operation: Callable[[], _Result], + device: torch.device, + count: int, + mode: str, + warmup: int, + repeats: int, +) -> tuple[dict[str, object], _Result]: + """Measure synchronized calls, then profile one separate warmed call.""" + for _ in range(warmup): + result = operation() + _synchronize(device) + del result + elapsed: list[float] = [] + for _ in range(repeats): + _synchronize(device) + start = time.perf_counter() + result = operation() + _synchronize(device) + elapsed.append(time.perf_counter() - start) + del result + + process = psutil.Process() + rss_before = process.memory_info().rss + gpu_before = 0 + if device.type == "cuda": + gpu_before = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + result = operation() + _synchronize(device) + gpu_delta, gpu_peak = 0, 0 + if device.type == "cuda": + gpu_delta = torch.cuda.memory_allocated(device) - gpu_before + gpu_peak = torch.cuda.max_memory_allocated(device) - gpu_before + median = statistics.median(elapsed) + # Current kernels allocate six float32 joints and one int32 validity per + # output. Torch conversion shares joints but allocates boolean validity. + warp_bytes = count * (1 if mode == "single" else 512) * (6 * 4 + 4) + return { + "sample_size": count, + "impl": mode, + "cost_time_ms": median * 1000, + "total_repeats_ms": sum(elapsed) * 1000, + "targets_per_second": count / median, + "cpu_delta_mb": (process.memory_info().rss - rss_before) / _MIB, + "gpu_delta_mb": gpu_delta / _MIB, + "peak_gpu_mb": gpu_peak / _MIB, + "warp_output_payload_mb": warp_bytes / _MIB, + }, result + + +def _quality( + result: _Result, + reference: _Result, + target: torch.Tensor, + cfg: URSolverCfg, + mode: str, +) -> dict[str, object]: + """Check validity parity and reconstruction of one valid result per pose.""" + valid, joints = result + ref_valid, ref_joints = reference + if valid.ndim == 2: + first = valid.to(dtype=torch.int32).argmax(dim=1) + rows = torch.arange(len(target), device=target.device) + joints = joints[rows, first] + valid = valid.any(dim=1) + torch.testing.assert_close(valid, ref_valid) + max_joint_delta = None + if mode != "all" and valid.any(): + max_joint_delta = (joints[valid] - ref_joints[valid]).abs().max().item() + torch.testing.assert_close( + joints[valid], ref_joints[valid], atol=1e-4, rtol=1e-5 + ) + translation_max = rotation_max = None + if valid.any(): + reconstructed = _forward_kinematics(joints[valid], cfg) + expected = target[valid].to(dtype=torch.float64) + translation_max = ( + torch.linalg.vector_norm( + reconstructed[:, :3, 3] - expected[:, :3, 3], dim=-1 + ) + .max() + .item() + ) + relative = reconstructed[:, :3, :3].transpose(1, 2) @ expected[:, :3, :3] + skew = torch.stack( + [ + relative[:, 2, 1] - relative[:, 1, 2], + relative[:, 0, 2] - relative[:, 2, 0], + relative[:, 1, 0] - relative[:, 0, 1], + ], + dim=-1, + ) + sine = torch.linalg.vector_norm(skew, dim=-1) / 2 + cosine = (relative.diagonal(dim1=1, dim2=2).sum(dim=1) - 1) / 2 + rotation_max = torch.atan2(sine, cosine).max().item() + return { + "sample_size": len(target), + "impl": mode, + "success_rate": valid.float().mean().item(), + "translation_max_m": translation_max, + "rotation_max_rad": rotation_max, + "max_joint_delta_vs_legacy": max_joint_delta, + "validity_matches_legacy": True, + } + + +def _write_report(payload: dict[str, object], output_dir: Path) -> Path: + """Write the three required benchmark tables to one Markdown artifact.""" + output_dir.mkdir(parents=True, exist_ok=True) + report = output_dir / f"ur_solver_{datetime.now():%Y%m%d_%H%M%S}.md" + lines = ["# UR solver benchmark", "", str(payload["environment"]), ""] + lines.extend(str(note) for note in payload["notes"]) + for title, key in ( + ("Time & Memory", "performance"), + ("Success & Other Metrics", "quality"), + ("Leaderboard", "leaderboard"), + ): + rows = payload[key] + headers = list(rows[0]) + lines.extend( + [ + "", + f"## {title}", + "", + "| " + " | ".join(headers) + " |", + "| " + " | ".join(["---"] * len(headers)) + " |", + ] + ) + for row in rows: + values = [ + f"{row[key]:.6g}" if isinstance(row[key], float) else str(row[key]) + for key in headers + ] + lines.append("| " + " | ".join(values) + " |") + report.write_text("\n".join(lines) + "\n", encoding="utf-8") + return report + + +@torch.inference_mode() +def run_all_benchmarks() -> None: + """Compare all three UR5 IK paths and write timing and accuracy reports.""" + args = _parse_args() + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + print("Skipped: CUDA is unavailable; use --device cpu for a CPU run.") + return + wp.init() + cfg = URSolverCfg( + ur_type="ur5", + joint_names=[f"joint_{joint}" for joint in range(6)], + user_qpos_limits=[[-2 * math.pi] * 6, [2 * math.pi] * 6], + ik_nearest_weight=[1.0, 0.5, 2.0, 1.5, 0.75, 3.0], + ) + # IK uses only DH parameters. A supplied placeholder chain bypasses unrelated + # URDF loading and torch.compile FK setup; FK validation above is independent. + solver = cfg.init_solver(device=device, pk_serial_chain=object()) + generator = torch.Generator(device="cpu").manual_seed(args.seed) + performance, quality = [], [] + print("UR5 analytic IK benchmark (warmup excluded; synchronized wall time)") + for count in args.counts: + qpos = (torch.rand((count, 6), generator=generator) * 2 - 1).to(device) + target = _forward_kinematics(qpos * math.pi, cfg).float() + seed = ( + (torch.rand((count, 6), generator=generator) * 2 - 1) * (2 * math.pi) + ).to(device) + reference = _solve(solver, target, seed, "legacy") + for mode in _MODES: + row, result = _measure( + lambda: _solve(solver, target, seed, mode), + device, + count, + mode, + args.warmup, + args.repeats, + ) + performance.append(row) + quality.append(_quality(result, reference, target, cfg, mode)) + del result + print( + f"n={count:>7d} {mode:>6s}: {row['cost_time_ms']:>9.3f} ms " + f"| {row['targets_per_second']:>12.0f} targets/s " + f"| Torch peak delta={row['peak_gpu_mb']:.3f} MiB " + f"| Warp output payload={row['warp_output_payload_mb']:.3f} MiB" + ) + del reference, target, seed, qpos + + leaderboard = [] + for mode in _MODES: + rows = [row for row in quality if row["impl"] == mode] + total = sum(row["sample_size"] for row in rows) + success = sum(row["sample_size"] * row["success_rate"] for row in rows) + leaderboard.append({"algorithm": mode, "overall_success_rate": success / total}) + leaderboard.sort(key=lambda row: row["overall_success_rate"], reverse=True) + payload = { + "environment": { + "device": str(device), + "gpu": ( + torch.cuda.get_device_name(device) if device.type == "cuda" else None + ), + "torch": torch.__version__, + "warp": wp.__version__, + "warmup": args.warmup, + "repeats": args.repeats, + "seed": args.seed, + }, + "notes": [ + "single uses return_all_solutions=False; all uses True; legacy expands " + "all solutions then runs the original weighted Torch norm and argmin.", + "cost_time_ms is the median; total_repeats_ms sums measured calls only. " + "Input generation, warmup, memory profiling, and quality checks are excluded.", + "Memory uses MiB. cpu_delta_mb is process RSS delta. gpu_delta_mb and " + "peak_gpu_mb are measured PyTorch allocation increases for a separate " + "warmed call; they exclude Warp allocations and are zero on CPU.", + "warp_output_payload_mb is the logical payload of direct Warp output " + "buffers (6 float32 joints plus int32 validity per candidate), on the " + "selected device. It excludes allocator overhead and temporary/register " + "storage. These columns do not measure total peak device memory.", + "Targets are DH-generated reachable UR5 poses with limits ±2π, independent " + "random seeds, and nonuniform weights. FK metrics check one valid " + "representative per pose. Legacy parity checks run outside timings.", + ], + "performance": performance, + "quality": quality, + "leaderboard": [ + {"rank": rank, **row} for rank, row in enumerate(leaderboard, start=1) + ], + } + report = _write_report(payload, args.output_dir) + if args.json: + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + print(f"Markdown report saved: {report}") + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/tests/sim/motion/solvers/test_ur_solver.py b/tests/sim/motion/solvers/test_ur_solver.py index 2d138c8d0..338af161e 100644 --- a/tests/sim/motion/solvers/test_ur_solver.py +++ b/tests/sim/motion/solvers/test_ur_solver.py @@ -19,11 +19,12 @@ import torch import pytest import numpy as np +import warp as wp from embodichain.data import get_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.motion.solvers import URSolverCfg +from embodichain.lab.sim.motion.solvers import URSolver, URSolverCfg from embodichain.lab.sim.cfg import ( RenderCfg, JointDrivePropertiesCfg, @@ -35,6 +36,252 @@ ) +def _make_analytic_ur_solver( + device: str = "cpu", + weights: list[float] | None = None, + limits: torch.Tensor | None = None, +) -> URSolver: + """Build the analytical solver without loading robot assets or a simulator.""" + wp.init() + if limits is None: + limits = torch.tensor([[-2.0 * torch.pi, 2.0 * torch.pi]] * 6) + cfg = URSolverCfg( + ur_type="ur5", + joint_names=[f"joint_{index}" for index in range(6)], + ik_nearest_weight=weights, + user_qpos_limits=limits.tolist(), + ) + # The analytical IK path never accesses the serial chain; explicit limits + # bypass the only chain-dependent setup needed by these tests. + return cfg.init_solver(device=torch.device(device), pk_serial_chain=object()) + + +def _ur_dh_poses(solver: URSolver, qpos: torch.Tensor) -> torch.Tensor: + """Generate independent DH targets, including the configured TCP.""" + cfg = solver.cfg + dh_parameters = [ + (cfg.d1, 0.0, cfg.alpha1), + (0.0, cfg.a2, 0.0), + (0.0, cfg.a3, 0.0), + (cfg.d4, 0.0, cfg.alpha4), + (cfg.d5, 0.0, cfg.alpha5), + (cfg.d6, 0.0, 0.0), + ] + poses = [] + for joints in qpos.cpu(): + pose = torch.eye(4) + for theta, (d, a, alpha) in zip(joints, dh_parameters): + pose = pose @ URSolver.dh_matrix(theta, d, a, alpha) + poses.append(pose) + tcp = torch.as_tensor(solver.tcp_xpos, dtype=torch.float32) + return (torch.stack(poses) @ tcp).to(solver.device) + + +def _sample_ur_joints(count: int = 12) -> torch.Tensor: + generator = torch.Generator().manual_seed(712) + joints = torch.rand((count, 6), generator=generator) * 5.0 - 2.5 + # Keep the wrist away from a singularity while covering both wrist branches. + joints[:, 4] = joints[:, 4].sign() * (0.4 + joints[:, 4].abs() * 0.7) + return joints + + +def _assert_nearest_matches_all_solutions( + solver: URSolver, + poses: torch.Tensor, + seed: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + all_valid, all_qpos = solver.get_ik(poses, seed, return_all_solutions=True) + batch_size = 1 if poses.ndim == 2 else len(poses) + assert all_valid.shape == (batch_size, 512) + assert all_qpos.shape == (batch_size, 512, 6) + reference_seed = solver.get_default_qpos_seed()[None] if seed is None else seed + distances = torch.norm( + solver.ik_nearest_weight * (all_qpos - reference_seed[:, None, :]), dim=-1 + ) + distances[~all_valid] = float("inf") + indices = distances.argmin(dim=1) + rows = torch.arange(batch_size, device=solver.device) + valid, qpos = solver.get_ik(poses, seed, return_all_solutions=False) + assert valid.shape == (batch_size,) + assert qpos.shape == (batch_size, 6) + assert valid.dtype == torch.bool + assert qpos.dtype == torch.float32 + assert valid.device == solver.device + assert qpos.device == solver.device + torch.testing.assert_close(valid, all_valid[rows, indices]) + torch.testing.assert_close( + qpos, all_qpos[rows, indices], rtol=0.0, atol=2e-6, equal_nan=True + ) + return valid, qpos, all_valid, all_qpos + + +@pytest.mark.no_sim +@pytest.mark.parametrize("seed_dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize( + "weights", + [None, [0.0, 0.25, -2.0, 4.0, 0.5, 1.5]], + ids=["uniform", "nonuniform-with-zero-and-negative"], +) +def test_ur_nearest_matches_weighted_full_candidates( + weights: list[float] | None, seed_dtype: torch.dtype +) -> None: + solver = _make_analytic_ur_solver(weights=weights) + joints = _sample_ur_joints() + poses = _ur_dh_poses(solver, joints) + seed = joints.roll(1, dims=0).to(seed_dtype) + seed[::2] += 2.0 * torch.pi + seed[1::2] -= 2.0 * torch.pi + valid, _, _, _ = _assert_nearest_matches_all_solutions(solver, poses, seed) + assert valid.all() + + +@pytest.mark.no_sim +def test_ur_nearest_preserves_periodic_representatives() -> None: + solver = _make_analytic_ur_solver() + joints = torch.tensor([[0.3, -1.1, 1.4, -0.5, 0.9, 0.4]]) + shifted = joints - 2.0 * torch.pi * joints.sign() + mixed = joints.clone() + mixed[:, ::2] = shifted[:, ::2] + seed = torch.cat([joints, shifted, mixed]) + poses = _ur_dh_poses(solver, joints).repeat(3, 1, 1) + valid, qpos, _, _ = _assert_nearest_matches_all_solutions(solver, poses, seed) + assert valid.all() + torch.testing.assert_close(qpos, seed, rtol=0.0, atol=5e-5) + + +@pytest.mark.no_sim +def test_ur_nearest_respects_limits_requiring_periodic_shifts() -> None: + joints = torch.tensor([[0.4, -1.1, 1.3, -0.7, 0.8, 0.2]]) + shifted = joints.clone() + shifted[0, 0] += 2.0 * torch.pi + shifted[0, 3] -= 2.0 * torch.pi + limits = torch.stack([shifted[0] - 0.08, shifted[0] + 0.08], dim=1) + solver = _make_analytic_ur_solver(limits=limits) + poses = _ur_dh_poses(solver, joints) + valid, qpos, _, _ = _assert_nearest_matches_all_solutions(solver, poses, joints) + assert valid.all() + assert (qpos >= limits[:, 0]).all() + assert (qpos <= limits[:, 1]).all() + torch.testing.assert_close(qpos, shifted, rtol=0.0, atol=5e-5) + + +@pytest.mark.no_sim +def test_ur_nearest_zero_weights_keep_first_valid_candidate_on_ties() -> None: + solver = _make_analytic_ur_solver(weights=[0.0] * 6) + poses = _ur_dh_poses(solver, _sample_ur_joints(3)) + seed = torch.full((1, 6), 2.0) # A single seed broadcasts across targets. + valid, qpos, all_valid, all_qpos = _assert_nearest_matches_all_solutions( + solver, poses, seed + ) + assert valid.all() + assert (all_valid.sum(dim=1) > 1).all() + first_indices = all_valid.to(torch.int32).argmax(dim=1) + torch.testing.assert_close(qpos, all_qpos[torch.arange(3), first_indices]) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("seed_value", [float("nan"), float("inf"), 1e30]) +def test_ur_nearest_preserves_nonfinite_distance_selection(seed_value: float) -> None: + joints = torch.tensor([[0.4, -1.1, 1.3, -0.7, 0.8, 0.2]]) + # Exclude the unshifted first candidate to exercise validity together with + # argmin's first-NaN/first-infinity behavior, including finite overflow. + shifted = joints[0].clone() + shifted[0] += 2.0 * torch.pi + limits = torch.stack([shifted - 0.08, shifted + 0.08], dim=1) + solver = _make_analytic_ur_solver(limits=limits) + poses = _ur_dh_poses(solver, joints) + seed = torch.full_like(joints, seed_value) + _assert_nearest_matches_all_solutions(solver, poses, seed) + + +@pytest.mark.no_sim +def test_ur_nearest_broadcasts_singleton_seed_across_targets() -> None: + solver = _make_analytic_ur_solver(weights=[1.0, 0.5, 2.0, 4.0, 0.25, 1.5]) + poses = _ur_dh_poses(solver, _sample_ur_joints(3)) + seed = torch.tensor([[3.0, -4.0, 5.0, 1.0, -2.0, 4.0]], dtype=torch.float64) + valid, _, _, _ = _assert_nearest_matches_all_solutions(solver, poses, seed) + assert valid.all() + + +@pytest.mark.no_sim +def test_ur_nearest_default_seed_unbatched_pose_and_tcp() -> None: + solver = _make_analytic_ur_solver() + solver.set_tcp( + np.array( + [ + [0.0, -1.0, 0.0, 0.04], + [1.0, 0.0, 0.0, -0.02], + [0.0, 0.0, 1.0, 0.12], + [0.0, 0.0, 0.0, 1.0], + ] + ) + ) + pose = _ur_dh_poses(solver, _sample_ur_joints(1))[0] + valid, _, _, _ = _assert_nearest_matches_all_solutions(solver, pose, None) + assert valid.all() + + +@pytest.mark.no_sim +@pytest.mark.parametrize("failure", ["unreachable", "excluded-by-limits"]) +def test_ur_nearest_all_invalid_preserves_first_candidate(failure: str) -> None: + solver = _make_analytic_ur_solver() + poses = _ur_dh_poses(solver, _sample_ur_joints(2)) + if failure == "unreachable": + poses[:, :3, 3] = 10.0 + else: + solver.set_qpos_limits(torch.zeros(6), torch.full((6,), 0.01)) + valid, qpos, all_valid, all_qpos = _assert_nearest_matches_all_solutions( + solver, poses, None + ) + assert not valid.any() + assert not all_valid.any() + torch.testing.assert_close(qpos, all_qpos[:, 0], equal_nan=True) + + +@pytest.mark.no_sim +def test_ur_nearest_avoids_full_candidate_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + solver = _make_analytic_ur_solver() + joints = _sample_ur_joints(3) + poses = _ur_dh_poses(solver, joints) + # Warm up before observing the per-call allocations. + solver.get_ik(poses, joints) + allocation_sizes = [] + + def record_allocations(allocate): + def wrapped(*args, **kwargs): + array = allocate(*args, **kwargs) + allocation_sizes.append(array.size) + return array + + return wrapped + + for allocator in ("zeros", "empty"): + monkeypatch.setattr(wp, allocator, record_allocations(getattr(wp, allocator))) + valid, qpos = solver.get_ik(poses, joints, return_all_solutions=False) + assert valid.all() + assert qpos.shape == joints.shape + # Even the smallest full-candidate buffer contains N * 512 scalar values. + assert all(size < len(joints) * 512 for size in allocation_sizes) + + +@pytest.mark.no_sim +@pytest.mark.gpu +@pytest.mark.parametrize("seed_dtype", [torch.float32, torch.float64]) +def test_ur_nearest_cuda_matches_full_candidates(seed_dtype: torch.dtype) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + solver = _make_analytic_ur_solver( + device="cuda:0", weights=[0.0, 0.25, -2.0, 4.0, 0.5, 1.5] + ) + joints = _sample_ur_joints(4).to(solver.device) + poses = _ur_dh_poses(solver, joints) + seed = (joints - 2.0 * torch.pi * joints.sign()).to(seed_dtype) + valid, _, _, _ = _assert_nearest_matches_all_solutions(solver, poses, seed) + assert valid.all() + + def grid_sample_qpos_from_limits( qpos_limits: torch.Tensor, steps_per_joint: int = 4,