Describe the bug
URSolver.get_ik() always allocates and computes 512 candidate joint configurations per target pose, even when return_all_solutions=False and the caller only needs the nearest valid solution. The 512 candidates are the 8 analytical branches expanded across 2^6 periodic joint-shift combinations. This is functionally useful when all solutions are requested, but it creates substantial avoidable memory traffic and compute overhead for the normal single-solution path.
This is especially visible in workspace Cartesian/plane analysis, where compute_batch_ik() can submit thousands of targets in one batch. The workspace analyzer then performs an additional nearest-solution selection in PyTorch after the Warp kernel has already materialized the full candidate tensor.
Steps to reproduce
- Run a Cartesian workspace analysis with a UR preset and a sufficiently large sample count:
embodichain analyze-workspace \
--robot ur \
--robot-params '{"robot_type":"ur5"}' \
--control-part arm \
--mode cartesian_space \
--bounds -0.8 0.8 -0.8 0.8 0.0 1.2 \
--num-samples 10000 \
--batch-size 1000 \
--headless
- Profile the call to
URSolver.get_ik() / Robot.compute_batch_ik().
- Compare it with
OPWSolver, which materializes 8 candidates per target for the equivalent analytic IK path.
The relevant implementation currently contains:
# embodichain/lab/sim/motion/solvers/ur_solver.py
N_SOL = 512
all_qpos_wp = wp.zeros(n_sample * N_SOL * DOF, ...)
and then computes distances for all 512 candidates before selecting one result when return_all_solutions is false.
Expected behavior
When return_all_solutions=False, the UR analytic IK path should avoid materializing the complete candidate tensor where possible. It should either:
- select the nearest valid candidate inside the Warp kernel and return only one solution per target; or
- use a compact, configurable candidate representation and only expand periodic variants when they are required by the joint limits or when
return_all_solutions=True.
The return_all_solutions=True behavior should remain available and preserve all valid periodic representatives.
System Info
- Commit:
797fe56e (main)
- OS: Ubuntu Linux (kernel 5.15)
- GPU: NVIDIA RTX PRO 5000 72GB Blackwell
- CUDA: 13.0 (reported by
nvidia-smi)
- GPU Driver: 580.126.20
Additional context
The corresponding OPW implementation uses N_SOL = 8 and performs candidate selection with a second kernel. The current UR implementation therefore has a much larger per-target intermediate footprint despite both solvers being analytic 6-DoF IK implementations. A fused “generate, validate, and select” kernel would also remove the PyTorch-side distance tensor and argmin for the common single-solution path.
A benchmark should report warm-up-separated throughput, peak device memory, and total time for 1k/10k/100k targets with return_all_solutions both disabled and enabled. Any optimization should preserve joint-limit handling, periodic representatives, seed-weighted nearest-solution selection, and validity flags.
Checklist
Describe the bug
URSolver.get_ik()always allocates and computes 512 candidate joint configurations per target pose, even whenreturn_all_solutions=Falseand the caller only needs the nearest valid solution. The 512 candidates are the 8 analytical branches expanded across 2^6 periodic joint-shift combinations. This is functionally useful when all solutions are requested, but it creates substantial avoidable memory traffic and compute overhead for the normal single-solution path.This is especially visible in workspace Cartesian/plane analysis, where
compute_batch_ik()can submit thousands of targets in one batch. The workspace analyzer then performs an additional nearest-solution selection in PyTorch after the Warp kernel has already materialized the full candidate tensor.Steps to reproduce
embodichain analyze-workspace \ --robot ur \ --robot-params '{"robot_type":"ur5"}' \ --control-part arm \ --mode cartesian_space \ --bounds -0.8 0.8 -0.8 0.8 0.0 1.2 \ --num-samples 10000 \ --batch-size 1000 \ --headlessURSolver.get_ik()/Robot.compute_batch_ik().OPWSolver, which materializes 8 candidates per target for the equivalent analytic IK path.The relevant implementation currently contains:
and then computes distances for all 512 candidates before selecting one result when
return_all_solutionsis false.Expected behavior
When
return_all_solutions=False, the UR analytic IK path should avoid materializing the complete candidate tensor where possible. It should either:return_all_solutions=True.The
return_all_solutions=Truebehavior should remain available and preserve all valid periodic representatives.System Info
797fe56e(main)nvidia-smi)Additional context
The corresponding OPW implementation uses
N_SOL = 8and performs candidate selection with a second kernel. The current UR implementation therefore has a much larger per-target intermediate footprint despite both solvers being analytic 6-DoF IK implementations. A fused “generate, validate, and select” kernel would also remove the PyTorch-side distance tensor andargminfor the common single-solution path.A benchmark should report warm-up-separated throughput, peak device memory, and total time for 1k/10k/100k targets with
return_all_solutionsboth disabled and enabled. Any optimization should preserve joint-limit handling, periodic representatives, seed-weighted nearest-solution selection, and validity flags.Checklist