Add differentiable Warp TOPPRA planning on CPU and CUDA

agent_context/topics/motion-planning/planner-details.md

@@ -77,17 +77,61 @@ retain their existing timing contracts and do not consume these backend limits.
 
 ### TOPPRA
 
-`toppra_planner.py` fans batched solves into independent CPU work. The module-level
-NumPy worker must remain picklable and must not touch CUDA/Warp/simulation state.
-Single-worker/single-row execution can run inline. The automatic multiprocessing
-context uses spawn with GPU simulation to avoid forking initialized CUDA state.
-Inspect the config for overrides; do not force fork merely to reduce overhead.
+`toppra_planner.py` uses Warp on CPU and CUDA. Numerical ownership is in
+`compute/trajectory/_toppra_warp.py` and `_warp/toppra.py`, with cooperative
+CUDA inference in `_warp/toppra_cuda.py`; shared validation
+lives in `_toppra_common.py`. `_toppra.py` is an offline NumPy oracle and must
+not be imported by the runtime planner or Warp adapter.
+Robot-backed planning retains shared speed-cap and control-part resolution,
+including gradients of uncapped requested limits.
+The runtime and reference preserve uniform not-a-knot cubic geometry, the existing waypoint
+deduplication tolerance and final knot, and signed velocity/acceleration bounds
+containing rest. Planner outputs retain float32 and quantity sampling retains
+the existing integer coercion, independently of backend or input precision.
+Interval bounds are conservative; timing need not match the former external
+TOPPRA package. Tiny real moves must retain positive duration.
+
+Warp constructs splines, bounds, controllable sets and samples on the input
+device. It honors the current Torch stream and reads a scalar to size output.
+CUDA inference reduces constraint planes across one block per path; all
+early exits must remain uniform across the block. Inference may use fast
+polygon projection, while recorded calls retain the selected-expression
+projection and serial solver needed by the adjoint. Dense CPU inference sampling
+shares scalar timing and interval searches across the joints of each sample;
+recorded sampling retains its per-joint adjoint path.
+CPU inference amortizes sufficiently large batches over a shared bounded thread
+pool; small requests and adjoint-recording calls stay inline. Row blocks own
+their numerical buffers, inherit Torch grad/inference modes and merge using
+the global endpoint/zero-derivative/zero-dt padding contract. Ordinary inference
+launches must not record into Warp's process-global tape. CPU pool state is
+discarded after fork; it is not owned by an individual planner's `close()`.
+Planning has no process backend, per-planner worker state or process cleanup.
+The removed `backend`, `max_workers` and `mp_context` configuration fields must
+be removed from callers; do not restore them as ignored compatibility options.
+
+Warp records first-order adjoints when Torch gradient tracking is enabled and
+waypoints or tensor velocity/acceleration bounds require gradients. Both time
+and quantity sampling propagate position, velocity, acceleration and duration
+losses through the timing solve. Scalar, per-joint and signed tensor bounds keep
+their graphs on either device; TOPPRA option copies preserve these tensor
+references while copying mutable containers. `torch.no_grad()` retains
+ordinary forward planning.
+Filtering, grid topology, sample counts and active-constraint choices are
+discrete: derivatives apply locally away from their switches. Time sampling
+keeps `ceil(duration / sample_dt) + 1` points uniformly spread across the full
+duration, with the integer count fixed during backward. Failed rows contribute
+zero gradients; stationary rows differentiate only their retained position.
+Higher-order gradients are unsupported. Test finite-difference agreement,
+generator option copying/resampling and CUDA backward stream behavior in the
+TOPPRA tests.
 
 Variable-length outputs pad with the terminal position without adding duration.
-Per-row failures preserve other rows; a broken worker pool is discarded and
-rebuilt on a later call. Test these lifecycle boundaries in
-`tests/sim/motion/planners/test_toppra_batched.py` and keep dependency versions
-in the package declarations rather than this context.
+Per-row failures preserve other rows. Test planner contracts and the absence
+of process workers in
+`tests/sim/motion/planners/test_toppra_batched.py`; numerical limits, geometry,
+CPU/Warp agreement and CUDA stream behavior are covered in
+`tests/compute/test_toppra.py`; CPU projection, thread isolation, variable-length
+assembly and pool inheritance are covered in `tests/compute/test_toppra_cpu.py`.
 
 ### Trapezoidal and Double-S
 
@@ -133,6 +177,17 @@ semantics are owned by [collision worlds](collision-worlds.md).
 
 ## Retiming and playback
 
+`toppra_planner.py::retime_joint_paths()` is the shared entry point for
+re-parameterizing an already-planned joint path under velocity and acceleration
+limits, used by planners that emit geometry without executable timing. It is
+the planner's own time parameterization. All CPU/CUDA rows run on Warp,
+including single-row calls, so batch size does not change waypoint cleanup
+precision or the helper's float32/float64 output dtype. It resamples
+along the fitted spline, so recompute derived poses from its output rather
+than reusing the input samples. Optional per-row path lengths exclude
+rectangular tail holds from spline fitting; without them, an explicit repeated
+final waypoint retains the historical final-knot geometry.
+
 `compute/trajectory/timing.py` owns time-domain differentiation/resampling and
 `retime_to_control_grid()`. Zero-time position changes are invalid; repeated-time
 unchanged samples may represent padding/junctions. Time resampling preserves

docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst

@@ -144,3 +144,22 @@ Utilities
     :members:
     :undoc-members:
     :show-inheritance:
+
+Joint path retiming
+------------------
+
+TOPPRA uses a single Warp calculation chain on CPU and CUDA. Its configuration
+accepts ``grid_size`` and the shared robot/manager fields; the removed
+``backend``, ``max_workers`` and ``mp_context`` fields are no longer accepted.
+CPU inference uses shared threads without creating planner-owned processes.
+
+Re-parameterize existing joint paths with the same differentiable timing
+solve. The helper preserves float32/float64 input precision, supports valid
+per-row waypoint prefixes, and returns explicit timing with isolated failures.
+
+.. autosummary::
+   :nosignatures:
+
+   embodichain.lab.sim.motion.planners.toppra_planner.retime_joint_paths
+
+.. autofunction:: embodichain.lab.sim.motion.planners.toppra_planner.retime_joint_paths

docs/source/overview/sim/motion/planners/index.rst

@@ -17,7 +17,7 @@ Choose a planner
      - Native targets and requirements
    * - :doc:`toppra_planner`
      - Time-parameterizing joint paths under velocity and acceleration limits.
-     - Joint waypoints; requires TOPPRA.
+     - Joint waypoints; in-tree Warp on CPU or CUDA.
    * - :doc:`trapezoidal_planner`
      - Batched joint paths with trapezoidal or jerk-limited Double-S timing.
      - Joint waypoints; Torch or Warp.

docs/source/overview/sim/motion/planners/toppra_planner.md

@@ -1,8 +1,9 @@
 # ToppraPlanner
 
-`ToppraPlanner` time-parameterizes a joint path with
-[TOPPRA](https://github.com/hungpham2511/toppra) under velocity and acceleration
-limits. Install `toppra==0.6.3` before use. For Cartesian goals, use
+`ToppraPlanner` time-parameterizes joint paths under velocity and acceleration
+limits with the in-tree Warp implementation of
+[TOPPRA](https://github.com/hungpham2511/toppra). CPU and CUDA use the same
+calculation chain; the external `toppra` package is not required. For Cartesian goals, use
 [MotionGenerator](../motion_generator.md) to generate joint waypoints through IK.
 
 ## Direct planning
@@ -30,6 +31,7 @@ goal_qpos[:, 0] += 0.1  # Choose a goal within this robot's limits.
 result = planner.plan(
     [PlanState.from_qpos(start_qpos), PlanState.from_qpos(goal_qpos)],
     ToppraPlanOptions(
+        control_part="arm",
         constraints={"velocity": 0.2, "acceleration": 0.5},
         sample_method=TrajectorySampleMethod.TIME,
         sample_interval=0.01,
@@ -41,17 +43,89 @@ planner.close()
 
 ## Constraints and sampling
 
-- Set `constraints` on `ToppraPlanOptions`; a scalar applies to all controlled
-  joints and a list specifies per-joint limits.
+- Set `constraints` on `ToppraPlanOptions`; positive scalars apply to all
+  controlled joints, vectors specify per-joint magnitudes, and `(DOF, 2)`
+  arrays specify signed bounds containing zero. Explicit velocity requests
+  are capped by the robot's runtime speed limits. `None` uses robot speed
+  limits and the default planning acceleration budget.
+- `control_part` selects the trajectory's joint order. A partial-joint path
+  needs an explicit group when the mapping is ambiguous. Runtime zero-speed
+  joints must remain constant at every waypoint, including points that
+  would otherwise be removed by duplicate filtering.
 - `TIME` uses `sample_interval` as the requested time spacing in seconds;
   samples include both endpoints, so actual spacing can be smaller.
 - `QUANTITY` uses `sample_interval` as the output count, with at least two
   samples.
 
 Results follow the shared
-[PlanResult contract](../motion_generator.md#inputs-and-results). Planning
-runs on CPU; batched workloads can use the configured worker pool. Call
-`close()` when a directly owned planner is no longer needed.
+[PlanResult contract](../motion_generator.md#inputs-and-results). Planner
+results use float32 on the robot's device. Internal spline fitting, interval
+bounds, timing and sampling use float64. `grid_size` sets the minimum timing
+grid size; segment alignment and path length can increase the actual size.
+The interval bounds are conservative, so durations can differ from the
+former external solver. Timing alone does not certify collision freedom.
+
+## CPU execution and configuration migration
+
+CPU inference automatically splits sufficiently large batches across a shared
+pool of at most eight threads. Small requests and calls recording derivatives
+run inline. CUDA numerical work remains on the input device. Planning does
+not create a CPU process pool, and individual planners own no worker processes.
+`close()` is an idempotent lifecycle hook; it does not close the shared CPU
+thread executor.
+
+Remove `backend`, `max_workers`, and `mp_context` from existing
+`ToppraPlannerCfg` constructors and configuration dictionaries. These options
+are no longer accepted. The robot selects the device, and Warp is the only
+runtime implementation:
+
+```python
+cfg = ToppraPlannerCfg(
+    robot_uid=robot.uid,
+    sim_instance_id=sim.instance_id,
+    grid_size=100,
+)
+```
+
+The NumPy implementation is an offline numerical reference used by tests;
+it cannot be selected as a planner backend.
+
+## First-order gradients
+
+Waypoint tensors and tensor velocity/acceleration requests retain their graphs
+through `ToppraPlanOptions` copying and the timing solve on CPU or CUDA. For
+example, with a valid joint goal in the direct-planning setup above:
+
+```python
+goal = goal_qpos.detach().requires_grad_()
+result = planner.plan(
+    [PlanState.from_qpos(start_qpos), PlanState.from_qpos(goal)],
+    ToppraPlanOptions(
+        control_part="arm",
+        constraints={"velocity": 0.2, "acceleration": 0.5},
+        sample_method=TrajectorySampleMethod.QUANTITY,
+        sample_interval=101,
+    ),
+)
+assert result.is_all_success()
+result.duration.sum().backward()
+assert goal.grad is not None
+```
+
+Gradients are local to fixed waypoint filtering, timing-grid topology, output
+counts and active constraints. Higher-order derivatives are unsupported.
+Failed rows contribute zero gradients; stationary rows differentiate only
+their retained position. Robot speed-cap saturation stops gradients through
+the capped request. Supply trainable boundary states through `PlanState.qpos`.
+
+## Retiming an existing path
+
+`retime_joint_paths()` exposes the same Warp timing solve for a `(B, N, DOF)`
+tensor without constructing a planner. Provide explicit velocity and
+acceleration limits; optional `path_lengths` exclude per-row tail padding
+from spline fitting. Outputs preserve float32/float64 input precision, while
+other input dtypes are promoted to float32. Recompute poses from the returned
+positions because the helper resamples along the fitted spline.
 
 ## References