Retime NMG rollouts under joint velocity and acceleration limits

[Yuan-Xinyi]

Stack

• Stack: main ← Replace external TOPPRA with differentiable CPU/CUDA retiming #724 ← Measure dynamic limits and mark nominal timing in the NMG benchmark #713 ← Retime NMG rollouts under joint velocity and acceleration limits #714 ← Pair a nominal and a retimed NMG arm in the benchmark #715
• Base: xinyi/nmg-dynamic-limit-metrics
• Depends on: Measure dynamic limits and mark nominal timing in the NMG benchmark #713, and through it Replace external TOPPRA with differentiable CPU/CUDA retiming #724 (the TOPPRA backend this PR retimes with)
• Followed by: Pair a nominal and a retimed NMG arm in the benchmark #715
• Related: Drop a held tail instead of appending a duplicate end knot #726 (held-tail fix on Replace external TOPPRA with differentiable CPU/CUDA retiming #724, see below)

Description

This PR implements item 2 of #684: give NeuralPlanner an opt-in time parameterization so its output can respect joint velocity and acceleration limits.

NeuralPlanner integrates action * action_scale per step and reports NeuralPlannerCfg.dt as a constant. No dynamic limit informs that value, so the reported timing is bookkeeping: at the default action_scale = 0.2 and dt = 0.01, a saturated step implies 20 rad/s, against 2.62 rad/s for the slowest FR3 joint. The rollout is a usable geometric path, but nothing turned it into an executable trajectory.

What changed:

• NeuralPlannerCfg.constraints, default None. The default preserves current behavior exactly. When set, the rollout samples are re-parameterized under the given velocity and acceleration limits, and the returned timing, derivatives and poses all come from the solved trajectory.
• toppra_planner.retime_joint_paths() as the shared entry point. It applies the TOPPRA planner's own parameterization to a path another planner already produced, on Replace external TOPPRA with differentiable CPU/CUDA retiming #724's backends: Warp for CUDA input, for gradient requests, and for a CPU batch of more than one row (a stateless function has no worker pool, and batched Warp is 8–9× faster than solving rows one after another in NumPy); a single CPU row stays on NumPy. The rest-to-rest scalar law in _scalar_time_law.py was the wrong tool: it would brake at every one of the hundreds of rollout samples. ([Proposal] Make joint velocity and acceleration first-class in the NMG waypoint contract #684 originally named that module; this corrects it.)
• Success is derived from the trajectory returned. Retiming resamples along a spline, so the rollout's success flag does not carry over. Arrival is re-checked on the returned grid with the rollout's own test, run once per waypoint over all samples rather than once per sample, so nothing synchronizes with the host per sample. Retimed positions are checked against joint limits, and a row that moves in zero time fails. A row whose path cannot be parameterized is reported as failed.
• The output is sampled on NeuralPlannerCfg.dt, not on the rollout's step count (see below).
• Poses are recomputed from the retimed positions with one batched FK call (Robot.compute_batch_fk), reused by the waypoint re-check. The per-step version cost 618 ms of an 1188 ms plan for a batch of eight; the batched result is identical (max difference 0.0) and takes about 100 ms.
• One result assembly shared between ToppraPlanner and the new entry point, instead of duplicating the tail-padding of unequal-length rows.

What this does not do: retiming bounds the derivatives but cannot smooth the path. Jerk still comes from the policy, and the cost of respecting the limits shows up as duration.

Dependencies: none beyond #724.

Refs #684

Type of change

• New feature (non-breaking change which adds functionality)

Measured evidence

Real NMG checkpoint (swa3.onnx, K=5) on Franka FR3, smoke suite, 9 free-space cases at W=1/3/5, 3 measured trials each, on this branch with #724's backend. The same rollout is planned twice, once with constraints=None and once with the FR3 limits, so the path is identical and only the timing differs.

arm	motion_valid	waypoints	v viol	a viol	j viol	v util	a util	duration	plan ms (per-case median)
constraints=None	100%	100%	100%	100%	100%	1.77–5.42	117–245	N/A	67–256
retimed	100%	100%	0%	0%	0%	0.21–0.63	0.9986–0.99997	0.57–1.73 s	100–266

Per joint on one W=3 case, against each joint's own velocity limit and the uniform 10 rad/s² acceleration limit, measured with the benchmark's operator from #713:

joint	1	2	3	4	5	6	7	over limit
velocity limit	2.62	2.62	2.62	2.62	5.26	4.18	5.26
nominal |q̇|	8.57	3.18	3.49	6.50	1.75	3.55	7.73	5 of 7
retimed |q̇|	0.92	0.35	0.32	0.70	0.15	0.39	0.88	0 of 7
nominal |q̈|	1056	279	153	816	203	479	1384	7 of 7
retimed |q̈|	9.97	3.31	2.89	7.55	1.79	4.16	9.99	0 of 7

Duration 0.060 s → 0.993 s; 7 samples → 101, one per NeuralPlannerCfg.dt. Start and end positions are identical.

Acceleration sits just under the limit on every case (at most 0.99997), never over, even under #713's non-smoothing operator. A time-optimal parameterization rides its binding constraint, and #724's interval bounds are conservative, so the limit holds between grid points rather than only at them. An earlier revision of this PR used the external toppra package, whose accelerations exceeded the limit by up to 4.5% at the default grid, and compensated with a denser grid. That workaround is gone: #724 makes it unnecessary, and planning latency falls from 581–840 ms to the range above.

Per case, retiming changes the median planning time by −14 to +39 ms at B=1, within the run-to-run spread of the rollout itself.

The same rollout before and after time parameterization. Position is unchanged: retiming moves the clock, not the path. The nominal acceleration flips sign between adjacent steps; the retimed one is bang-bang, with joints 1 and 7 alternating at ±10 rad/s².

Output sample count, not the parameterization, was losing waypoints

Worst waypoint error on a five-waypoint case against the 10 mm tolerance (measured with the previous backend; the effect depends on output sample density, not on the solver):

output samples	worst error
rollout itself (13)	8.77 mm	within
retimed, N = 13 (inherited from the rollout)	10.55 mm	outside
retimed, N = 64	8.64 mm	within
retimed, N = 1024	8.64 mm	unchanged

Inheriting the rollout's step count left a trajectory stretched 16× sampled on 13 points; the policy stops as soon as it is barely inside tolerance, leaving about 1 mm of slack for resampling to consume. Fixed in e423555e by sampling on NeuralPlannerCfg.dt.

Batched rollouts need #726

In a batch, an environment that converges early holds its pose while the others keep rolling out, so its path ends in repeated samples. #724's waypoint cleanup keeps one of those duplicates as a knot, which reshapes the spline and adds a reversal before the goal. On the rollout above, four held samples lengthen the retimed duration from 0.993 s to 1.055 s (+6.2%), so a row's timing depends on what else is in its batch. #726 fixes this on #724's branch; with it, the held and unheld durations are equal. Single-environment plans are unaffected.

Validation

• python -m pytest tests/sim/motion/ tests/compute/ tests/benchmark/motion_generation/ — 1296 passed, 111 skipped, 7 deselected (whole stack).
• black . — clean.
• python docs/scripts/check_api_docs.py — 2315/2315 documented. retime_joint_paths is a new public export and is documented in public_api.rst; the result assembly stayed private since nothing outside the module needs it.
• context.py affected --base xinyi/nmg-dynamic-limit-metrics --explain — motion-planning and simulation-system. motion-planning/planner-details.md is updated here: the neural-adapter section states that dt is nominal and what constraints changes, and the retiming section names the shared entry point and its backend selection, including Warp for CPU batches. simulation-system matched only through the tests/sim/ watch path and describes no behavior this PR changes, so it needs no edit. context.py check — ok.

Two existing tests in test_neural_batched.py stub planner.cfg with a hand-built object; they gained the new field. No assertion changed.

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. No dependency changes were required.

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge, though the previously noted long-rollout retiming cost remains.

Summary

This PR adds opt-in dynamic-limit retiming to NeuralPlanner rollouts and exposes a shared joint-path retiming entry point.

• Retimed results use solved timing, recomputed poses, and renewed waypoint and joint-limit checks.
• The latest changes route multi-row CPU retiming through Warp and vectorize waypoint checks.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Neural rollout] --> B{Constraints set?}
  B -- No --> C[Nominal-timing result]
  B -- Yes --> D[Retime joint path]
  D --> E[Recompute poses and check returned samples]
  E --> F[Timed plan result]

Loading

_{Reviews (10) · Last reviewed commit: "perf(motion): drop per-sample host syncs..." · Reviewed by Greptile}

[Yuan-Xinyi]

All five findings were valid; fixed in ff8a48ea.

Findings 1-3 share one root cause worth naming: retiming fits a spline through the rollout samples and resamples it, so the returned trajectory is not the one the rollout verified. Carrying the rollout's success flag across that boundary was the mistake, and it let all three failures through. Success is now re-derived from the samples actually returned.

1. Waypoint arrival after resampling — fixed. The resampled grid can step past a waypoint the rollout stopped on. Arrival is re-checked on the returned grid using the rollout's own _is_active_reached, so a trajectory that passes a waypoint only between output samples is no longer reported as reaching it.

2. Zero-duration result claimed as success — fixed. The shared kernel's near-coincident-endpoint shortcut returns dt = 0 with success=True. A row that moves in zero total time is now a failure. Fixed locally in the neural planner rather than in _toppra_solve_one_env, since changing that shortcut would alter ToppraPlanner behavior beyond this change.

3. Spline overshoot past joint limits — fixed. The fit sees only the sampled values, so a path clamped at a limit can overshoot between knots. Retimed positions are now checked against the joint limits. Not clamped: clamping would reintroduce a kink and break the acceleration bound that is the point of the feature, so the row is reported as failed instead.

4. sample_interval=None with TIME sampling — fixed. The default is a sample count, and passing it as seconds silently collapsed a dense path to two samples whenever its duration was shorter. That combination now raises.

5. Bare dict annotation — fixed. constraints is now dict[str, float | list[float]] | None. Note this is deliberately more precise than the sibling ToppraPlanOptions.constraints and TrapezoidalPlanOptions.constraints, which are bare dict; AGENTS.md requires full annotation for public APIs and a new field should meet it.

[yangchen73]

Check completion on the CPU at every sample spend a lot of time.

[Yuan-Xinyi]

Agreed. The re-check walked the output grid one sample at a time and read a device flag back to the host on every step, which is about a hundred synchronizations per second of trajectory at the default period.

In 45e5a67 it runs the arrival test once per waypoint over all samples at once and finds each first arrival with a reduction. The order rule is kept: waypoint k+1 counts only at a sample after the one where waypoint k was first reached. The loop is now over waypoints (at most five), with no host read inside. test_retimed_waypoint_check_keeps_order_and_runs_once_per_waypoint checks both the order semantics and that the arrival test runs once per waypoint.

[yangchen73]

It solves CPU batch rows one at a time. Any better solution?

[Yuan-Xinyi]

Yes. ToppraPlanner fans CPU rows out to a worker pool it owns, but this is a stateless function without one, so rows ran back to back in NumPy. In 45e5a67, a CPU batch of more than one row goes through the batched Warp solve on the same CPU. A single row stays on NumPy, matching ToppraPlanner's CPU default. Measured on a real NMG rollout:

B	NumPy, one row at a time	Warp, CPU batch	speedup
8	28.7 ms	3.4 ms	8.4×
64	225 ms	24.7 ms	9.1×
256	900 ms	95.8 ms	9.4×

Durations agree to within 4e-7 s. Covered by test_retime_joint_paths_solves_a_cpu_batch_in_one_call, which also checks that each batched row matches its single-row solve.