diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index 342556550..2b90c69f8 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -138,6 +138,19 @@ 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. The kernel flags +nearby competing distances for legacy `torch.norm` / `argmin` selection in +chunks of at most 128 targets, preserving backend rounding and candidate order +at branch bisectors. Ordinary targets avoid the full candidate tensor; even an +all-ambiguous batch uses bounded candidate buffers. The ambiguity margin routes +the fallback and never decides a tie. `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. @@ -149,11 +162,13 @@ mask with the leading shape and qpos with a final `dof` axis; existing concrete `get_fk`/`get_ik` signatures remain available. Robot owns local-arena/root frame conversion and broadcasts root transforms without materializing per-target copies. -UR/OPW reuse internal candidate buffers through `solvers/_buffers.py`. +UR's all-solutions path and OPW reuse internal candidate buffers through +`solvers/_buffers.py`. UR's single-solution path keeps its own compact outputs +and bounded ambiguity-fallback buffers, without allocating the full-batch cache. `prepare_buffers(max_batch)` reserves a high-water capacity; allocation is lazy and later larger calls grow it. Public results remain independent of subsequent calls, including `return_all_solutions=True`. Calls are serialized by a lock and CUDA events; Warp kernels use the current Torch stream. Buffer storage is released -with the solver. UR retains all 512 periodic candidates and the existing nearest -selection; OPW retains eight candidates. OPW packs live joint limits in one host +with the solver. UR's all-solutions path retains all 512 periodic candidates; +OPW retains eight candidates. OPW packs live joint limits in one host transfer per call, so limit updates do not require cache invalidation. 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 c2b03245b..f5c8d7232 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 @@ -139,6 +139,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..73919d3cc 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,109 @@ 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), + selection_ambiguous: 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,)``. + selection_ambiguous: Output flags requesting legacy distance selection, + shape ``(N,)``. + + Candidates follow the all-solutions kernel's branch and periodic order. + Scalar Warp norms can round differently from PyTorch reductions. Nearby + competitors are flagged for bounded legacy selection by the caller; the + comparison margin must not itself be used to break ties. When all candidates + are invalid, the first candidate is returned with a false flag. + """ + 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) + second_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 + distance = wp.sqrt(squared_distance) + if valid != 0 and ( + distance < best_distance + or (wp.isnan(distance) and not wp.isnan(best_distance)) + ): + second_distance = best_distance + best_distance = distance + best_q = candidate + best_valid = valid + elif valid != 0 and distance < second_distance: + # Repeated periodic representatives have identical distances + # under either reduction and need no compatibility fallback. + different = int(0) + for t in range(6): + if candidate[t] != best_q[t]: + different = int(1) + if different != 0: + second_distance = distance + + for t in range(6): + qpos[i, t] = best_q[t] + ik_valid[i] = best_valid + # Both kernels generate float32 candidates, including for float64 seeds. + # Allow 32 float32 epsilons for the six weighted terms, their accumulation, + # square root and candidate rounding. This only routes ambiguous targets; + # the original PyTorch reduction determines their actual ordering. + margin = qpos_seed.dtype(3.814697265625e-6) * wp.max( + qpos_seed.dtype(1.0), wp.abs(best_distance) + ) + selection_ambiguous[i] = int( + wp.isfinite(best_distance) + and wp.isfinite(second_distance) + and second_distance - best_distance <= margin + ) diff --git a/embodichain/lab/sim/motion/solvers/ur_solver.py b/embodichain/lab/sim/motion/solvers/ur_solver.py index 629b2d96a..670a16768 100644 --- a/embodichain/lab/sim/motion/solvers/ur_solver.py +++ b/embodichain/lab/sim/motion/solvers/ur_solver.py @@ -25,11 +25,14 @@ from embodichain.compute.kinematics._warp.ur import ( URParam, ur_ik_kernel, + ur_ik_nearest_kernel, ) import math from ._buffers import _with_ik_buffers from embodichain.utils.device_utils import standardize_device_string +__all__ = ["URSolverCfg", "URSolver"] + @configclass class URSolverCfg(SolverCfg): @@ -145,19 +148,25 @@ 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 uses a fused Warp selection, with bounded + legacy selection for distances close enough to be affected by + floating-point rounding. **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 @@ -169,7 +178,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 = ( @@ -183,12 +192,87 @@ 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)) + 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) + ambiguous_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, ambiguous_wp], + device=wp_device, + ) + best_valid = wp.to_torch(best_valid_wp).bool() + best_qpos = wp.to_torch(best_qpos_wp) + ambiguous = wp.to_torch(ambiguous_wp).nonzero(as_tuple=True)[0] + legacy_seed = qpos_seed.to(device=device).expand(n_sample, DOF) + # Do not approximate ties with an epsilon: even a one-ULP difference + # in the legacy norm can choose another analytical branch. Reuse + # the original candidates and reduction for ambiguous targets only. + # Cap each candidate buffer at 128 * 512 * 6 floats (1.5 MiB), even + # when every target is on a branch bisector. + for rows in ambiguous.split(128): + count = rows.numel() + if count == 0: + continue + candidates_wp = wp.empty( + count * N_SOL * DOF, dtype=float, device=wp_device + ) + validity_wp = wp.empty(count * N_SOL, dtype=int, device=wp_device) + wp.launch( + kernel=ur_ik_kernel, + dim=count, + inputs=[ + wp.from_torch(target_xpos_batch[rows].reshape(-1)), + self._ur_params, + lower_qpos_limits_wp, + upper_qpos_limits_wp, + ], + outputs=[candidates_wp, validity_wp], + device=wp_device, + ) + candidates = wp.to_torch(candidates_wp).view(count, N_SOL, DOF) + validity = wp.to_torch(validity_wp).view(count, N_SOL).bool() + distances = torch.norm( + self.ik_nearest_weight * (candidates - legacy_seed[rows, None, :]), + dim=-1, + ) + distances[~validity] = float("inf") + nearest = distances.argmin(dim=1) + local_rows = torch.arange(count, device=device) + best_qpos[rows] = candidates[local_rows, nearest] + best_valid[rows] = validity[local_rows, nearest] + return best_valid, best_qpos + all_qpos_wp = self._ik_buffers.zeros( "qpos", n_sample, N_SOL * DOF, torch.float32 ) all_ik_valid_wp = self._ik_buffers.zeros("valid", n_sample, N_SOL, torch.int32) - lower_qpos_limits_wp = wp.from_torch(self.lower_qpos_limits) - upper_qpos_limits_wp = wp.from_torch(self.upper_qpos_limits) wp.launch( kernel=ur_ik_kernel, dim=(n_sample,), @@ -214,22 +298,8 @@ def get_ik( .to(device=device) ) - if return_all_solutions: - return all_solutions_validity.clone(), all_solutions.clone() - # 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] - # ``ik_qpos`` indexes the reusable candidate scratch buffer. Copy both - # outputs before releasing the borrow so a subsequent solve cannot - # mutate tensors returned to the caller. - return ik_validity.clone(), ik_qpos.clone() + # Copy reusable scratch storage before releasing the buffer borrow. + return all_solutions_validity.clone(), all_solutions.clone() @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..637064fe8 --- /dev/null +++ b/scripts/benchmark/robotics/kinematic_solver/ur_solver.py @@ -0,0 +1,363 @@ +# ---------------------------------------------------------------------------- +# 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) + # Primary output buffers only. The single kernel also emits an int32 + # ambiguity flag; bounded fallback candidate buffers are excluded here. + # Torch conversion shares joints but allocates boolean validity. + warp_bytes = count * (1 if mode == "single" else 512) * (6 * 4 + 4) + if mode == "single": + warp_bytes += count * 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 primary Warp output " + "buffers (6 float32 joints plus int32 validity per candidate, plus an " + "int32 ambiguity flag per target in single mode). It excludes bounded " + "legacy-fallback buffers, 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 8b073413c..05a349014 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, @@ -34,6 +35,396 @@ ) +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 +def test_ur_nearest_nonzero_distance_rounding_regression() -> None: + """The reported bisector selected branch 384 instead of legacy branch 256.""" + solver = _make_analytic_ur_solver() + pose = torch.tensor( + [ + [ + [ + 0.11713490635156631, + 0.8637253642082214, + -0.49016112089157104, + 0.061536163091659546, + ], + [ + -0.34352806210517883, + -0.4278513193130493, + -0.8360213041305542, + -0.12725721299648285, + ], + [ + -0.9318088889122009, + 0.2663114070892334, + 0.24659760296344757, + 0.7414405345916748, + ], + [0.0, 0.0, 0.0, 1.0], + ] + ], + dtype=torch.float32, + ) + seed = torch.tensor( + [ + [ + 0.6716036796569824, + -2.2640819549560547, + 1.3243775367736816, + -0.8974207639694214, + 0.0, + 0.3734279274940491, + ] + ], + dtype=torch.float32, + ) + # Include both sides of the boundary, where treating near-equality as a tie + # would incorrectly retain the earlier branch on one side. + seeds = seed.repeat(3, 1) + seeds[:, 4] += torch.tensor([-1e-6, 0.0, 1e-6]) + _, selected, all_valid, all_qpos = _assert_nearest_matches_all_solutions( + solver, pose.repeat(3, 1, 1), seeds + ) + distances = torch.norm(all_qpos - seeds[:, None], dim=-1) + distances[~all_valid] = float("inf") + assert (distances[:, [256, 384]] > 2.0).all() + assert distances[0, 256] != distances[0, 384] + assert distances[2, 256] != distances[2, 384] + assert (selected[0] - selected[2]).abs().max() > 3.0 + # Compare against the actual backend reduction, without assuming all CPU + # architectures or PyTorch builds round these two norms identically. + expected = all_qpos[torch.arange(3), distances.argmin(dim=1)] + torch.testing.assert_close(selected, expected, rtol=0.0, atol=0.0) + + +@pytest.mark.no_sim +@pytest.mark.parametrize( + "device,seed_dtype,weights,count", + [ + ("cpu", torch.float32, None, 128), + ("cpu", torch.float64, [0.75, 0.25, -2.0, 4.0, 0.5, 1.5], 16), + pytest.param("cuda:0", torch.float32, None, 8, marks=pytest.mark.gpu), + pytest.param( + "cuda:0", + torch.float64, + [0.75, 0.25, -2.0, 4.0, 0.5, 1.5], + 8, + marks=pytest.mark.gpu, + ), + ], +) +def test_ur_nearest_branch_bisectors_match_legacy( + device: str, + seed_dtype: torch.dtype, + weights: list[float] | None, + count: int, +) -> None: + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + solver = _make_analytic_ur_solver(device=device, weights=weights) + poses = _ur_dh_poses(solver, _sample_ur_joints(count)) + _, candidates = solver.get_ik(poses, return_all_solutions=True) + branches = candidates[:, ::64].to(seed_dtype) + pairs = torch.combinations(torch.arange(8, device=solver.device)) + left, right = branches[:, pairs[:, 0]], branches[:, pairs[:, 1]] + midpoint = (left + right) * 0.5 + # 128 * 28 = 3,584 branch midpoints, plus the adjacent representable seeds + # on either side. Nonuniform weights retain midpoint equidistance. + seeds = torch.stack( + [torch.nextafter(midpoint, left), midpoint, torch.nextafter(midpoint, right)], + dim=2, + ).reshape(-1, 6) + targets = poses[:, None, None].expand(-1, len(pairs), 3, -1, -1).reshape(-1, 4, 4) + valid, _, all_valid, all_qpos = _assert_nearest_matches_all_solutions( + solver, targets, seeds + ) + assert valid.all() + distances = torch.norm( + solver.ik_nearest_weight * (all_qpos - seeds[:, None]), dim=-1 + ) + distances[~all_valid] = float("inf") + nearest = distances.sort(dim=1).values[:, :2] + gap = nearest[:, 1] - nearest[:, 0] + nonzero = nearest[:, 0] > 0.0 + assert (nonzero & (gap == 0.0)).any(), "Exercise nonzero exact ties" + assert (nonzero & (gap > 0.0) & (gap < 1e-5)).any(), "Exercise near ties" + + +@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) + # Main's reusable buffers allocate through Torch instead of wp.empty. + # Single-solution calls must not lazily allocate that full-candidate cache. + assert not solver._ik_buffers.arrays + + +@pytest.mark.no_sim +def test_ur_nearest_ambiguous_selection_bounds_candidate_allocations( + monkeypatch: pytest.MonkeyPatch, +) -> None: + solver = _make_analytic_ur_solver(weights=[0.0] * 6) + # Every target needs tie selection; exceed the 128-target fallback chunk. + joints = _sample_ur_joints(257) + poses = _ur_dh_poses(solver, joints) + all_valid, all_qpos = solver.get_ik(poses, return_all_solutions=True) + allocate = wp.empty + allocation_sizes = [] + + def record_allocation(*args, **kwargs): + array = allocate(*args, **kwargs) + allocation_sizes.append(array.size) + return array + + monkeypatch.setattr(wp, "empty", record_allocation) + valid, qpos = solver.get_ik(poses, joints[:1]) + first = all_valid.to(torch.int32).argmax(dim=1) + torch.testing.assert_close(valid, all_valid[torch.arange(len(joints)), first]) + torch.testing.assert_close(qpos, all_qpos[torch.arange(len(joints)), first]) + assert allocation_sizes.count(128 * 512 * 6) == 2 + assert max(allocation_sizes) <= 128 * 512 * 6 + + +@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,