Add differentiable Warp TOPPRA planning on CPU and CUDA

[yuecideng]

Description

TOPPRA planning currently detaches joint waypoints into NumPy and dispatches CPU process workers. This PR replaces that path with an in-tree float64 Warp implementation and connects first-order gradients through ToppraPlanner, typed option copies, and MotionGenerator resampling. CPU and CUDA use Warp as their only runtime implementation.

• Differentiate positions, velocities, accelerations and arrival intervals with respect to waypoint tensors and scalar, per-joint or signed velocity/acceleration bounds, with TIME and QUANTITY sampling.
• Keep CUDA numerical work on the input device. CUDA inference assigns a 64-lane block to each path for cooperative constraint reductions, with fast polygon projection for large batches. CPU inference uses a bounded shared thread pool and fast projection with an exhaustive fallback; dense sampling shares time searches across joints and narrows only on final output stores. Small CPU requests run inline. Recorded calls retain the serial kernels used by autograd.
• Remove the NumPy planner backend, process dispatch, fork/spawn selection, worker cleanup and broken-pool recovery. Simplify ToppraPlanner to a single Warp adapter. The offline NumPy oracle remains for numerical tests; runtime validation is shared without importing that oracle.
• Preserve current main behavior for robot speed caps, control_part, TIME defaults, frozen joints and float32 planner results. Reject raw motion in a frozen joint before waypoint filtering.
• Expose retime_joint_paths() for existing joint paths, preserving valid prefixes, float32/float64 precision, isolated failures and endpoint padding. Update the tutorial, Task Program fixtures, benchmark adapter, planner/API documentation and motion-planning context.

Migration: remove backend, max_workers and mp_context from ToppraPlannerCfg constructors and configuration dictionaries. These fields are no longer accepted. The robot selects the device; the CPU thread executor is shared by the Warp runtime. close() remains an idempotent lifecycle hook.

Timing optimizes conservative interval constraints on the existing not-a-knot spline, so durations can differ from the former external solver. Gradients are first-order and local to fixed filtering, sampling topology and active constraints; higher-order derivatives are unsupported. Trainable boundary states use PlanState.qpos; the generic MotionGenOptions.start_qpos deepcopy behavior is outside this change.

This is standalone on current main (82563fd5), extracted from the TOPPRA work in #712 with the CPU implementation improvements. It does not depend on #712 or include its Atomic Skill benchmark framework, evidence runtime, or NMG retiming feature.

Dependencies: declare SciPy directly and remove the external toppra==0.6.3 runtime dependency. Benchmark availability no longer depends on that external package.

Type of change

• New feature (differentiable CPU/CUDA planning)
• Breaking change (remove legacy planner backend/process configuration)
• Documentation update

Validation

196 focused checks passed:

• Compute/CPU/CUDA/autograd and dependency-isolation tests: 96 passed; three native simulator import-order cases excluded.
• Planner/facade tests: 48 passed; two real-simulator cases excluded. Benchmark availability: 1 passed. Atomic-action quantity retiming: 2 passed. These groups used import-only isolation for unavailable native simulator/cuRobo/Open3D dependencies; the TOPPRA, joint-limit resolver, MotionGenerator and interpolation implementations ran unchanged.
• Agent-context tests and API-doc checker tests: 49 passed.
• black . with Black 26.3.1; git diff --check; changed Python modules compile.
• API documentation coverage: 2595/2595 exports. Context check and affected-topic review completed; planner context updated for the single Warp route and shared thread ownership.

Coverage includes finite differences, independent HiGHS/analytic projection oracles on CPU/CUDA, cooperative blocks with more than 64 planes, finite infeasible rows, extreme scales, both output precisions, runtime-cap gradients before/after saturation, frozen-joint sub-tolerance motion, failed-row gradients, variable-length output, CUDA streams, repeated backward, CPU concurrency, foreign-tape isolation, failure cleanup and external-fork recovery. New regressions forbid CPU process creation and runtime imports of the offline NumPy solver; removed configuration fields are rejected.

Full native simulation, selected Task Program integration tests and the Sphinx dummy build remain blocked locally by the existing DexSim import error: missing reconstruct_cable_rings_kernel in dexsim.engine.newton_physics.runtime_kernels. They are not reported as passing.

Inference performance

Paired A/B against this PR’s previous c316cab1 revision on EPYC with an eight-physical-core affinity budget and RTX PRO 5000 Blackwell. Same 7-DoF paths, 8 waypoints, limits and grid; three warmups and 21 synchronized observations per case. CPU variants share one bounded executor to hold worker placement and allocator state constant. Times are median real .plan() latency, excluding setup/JIT.

Device / Torch threads	Batch	Output samples	Before (ms)	After (ms)	Speedup
cuda	1	128	5.96	2.10	2.84x
cuda	128	128	8.17	2.59	3.15x
cuda	2048	128	19.36	10.55	1.84x
cpu-t1	128	512	11.40	7.71	1.48x
cpu-t8	2048	512	113.97	80.71	1.41x

Across the 2/8/32-waypoint fixed-128-output GPU matrix, speedup is 1.84–5.76x. CPU dense-output cases improve up to 1.52x; small outputs retain the original sampler and show timing variation around baseline. All 162 measured version/device/shape combinations passed sampled endpoint/finite/limit checks, and the first eight rows’ q/v/a/dt outputs match exactly before/after in float32. Internal computation stays float64; training/backward throughput is not part of these claims.

Context was revised for CUDA block-uniform exits and the split inference/recorded paths. Existing public planner documentation remains accurate: signatures, configuration, units, sampling and gradient contracts are unchanged by this optimization.

Screenshots

N/A — numerical planning and API changes.

Checklist

• I have run the black . command to format the code base.
• I reviewed affected documentation and agent context, updated it where needed, or explained why no update was needed.
• Public API changes are reflected in the API docs (python docs/scripts/check_api_docs.py), if applicable.
• I have added tests that prove my fix is effective or that my feature works.
• Dependencies have been updated, if applicable.

[greptile-apps]

Confidence Score: 5/5

[High impact] The PR appears safe to merge, with the earlier non-blocking path-length validation issue still open.

Findings

1.  Invalid lengths silently shorten paths ▶

Fix with agent prompt

### Issue 1
embodichain/lab/sim/motion/planners/toppra_planner.py:141-143
Both `retime_joint_paths()` and `_tensor_path_lengths()` convert `path_lengths` to `torch.int32` before checking its range. An invalid `torch.int64` length such as `4294967298` wraps to `2` and passes the check for a four-waypoint path. The planner then uses only the first two waypoints instead of rejecting the input.

Check the original values against `2..count` before converting them, and add an overflow case to the length tests.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

Replaces external TOPPRA and CPU process workers with an in-tree Warp solver on CPU and CUDA.

• Warp plans joint paths and returns differentiable timed motion.
• Existing joint paths can be retimed with gradients intact.
• TOPPRA planning no longer accepts external backend or worker settings.

No new findings were accepted.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[ToppraPlanner] --> B[Resolve robot limits]
  B --> C[Warp adapter]
  D[retime_joint_paths] --> C
  C --> E{Recording gradients?}
  E -->|Yes| F[Recorded solve and sampling]
  F --> G[Torch backward through Warp]
  E -->|No| H{Input device}
  H -->|CPU| I[Inline or shared threads]
  H -->|CUDA| J[Cooperative solve]
  I --> K[Timed positions and derivatives]
  J --> K
  F --> K

Loading

_{Reviews (3) · Last reviewed commit: "chore: sync TOPPRA branch with latest ma..." · Reviewed by Greptile}

[greptile-apps]

Invalid lengths silently shorten paths

Both retime_joint_paths() and _tensor_path_lengths() convert path_lengths to torch.int32 before checking its range. An invalid torch.int64 length such as 4294967298 wraps to 2 and passes the check for a four-waypoint path. The planner then uses only the first two waypoints instead of rejecting the input.

Check the original values against 2..count before converting them, and add an overflow case to the length tests.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/planners/toppra_planner.py
Line: 318-320

Comment:
**Invalid lengths silently shorten paths**

Both `retime_joint_paths()` and `_tensor_path_lengths()` convert `path_lengths` to `torch.int32` before checking its range. An invalid `torch.int64` length such as `4294967298` wraps to `2` and passes the check for a four-waypoint path. The planner then uses only the first two waypoints instead of rejecting the input.

Check the original values against `2..count` before converting them, and add an overflow case to the length tests.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.