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, with the same backend selection Replace external TOPPRA with differentiable CPU/CUDA retiming #724 introduces: Warp for CUDA input or when gradients are requested, NumPy workers on CPU. 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, 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.21–5.29	52–122	N/A	43–295
retimed	100%	100%	0%	0%	0%	0.20–0.61	0.998–1.000	0.57–1.73 s	79–424

Per joint on one W=3 case, against each joint's own velocity limit and the uniform 10 rad/s² acceleration limit (retimed values are the planner's analytic derivatives):

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̇|	6.74	3.03	3.09	5.22	1.75	2.27	7.33	5 of 7
retimed |q̇|	0.92	0.35	0.32	0.70	0.15	0.40	0.88	0 of 7
nominal |q̈|	329	109	84	269	94	176	692	7 of 7
retimed |q̈|	9.97	3.30	2.99	7.70	1.87	4.47	10.00	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 1.000 on every case, never over. 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.

Latency is noisy on this machine; the nominal arm itself ranges 43–295 ms per case. Retiming adds roughly 30–130 ms per plan at B=1.

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 97ebc8c0 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/ — 1277 passed, 111 skipped, 7 deselected.
• 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. 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 risk] Adds optional retiming to neural planner rollouts.

The PR appears safe to merge; no outstanding finding or new actionable defect was established.

Findings

1.  Long rollouts slow retiming ▶

Fix with agent prompt

### Issue 1
embodichain/lab/sim/motion/planners/neural_planner.py:696-699
With constraints enabled, a rollout that uses the default 240-step budget requests at least 24,100 TOPPRA gridpoints, compared with the kernel’s roughly 2,410-point default for that path. The shared entry point solves batch rows serially, so slow-to-converge or failed rollouts add planning latency that grows with batch size. A bounded or adaptive density would avoid paying this cost for every rollout sample.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

---

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

Summary

The PR adds optional velocity- and acceleration-constrained timing for neural-planner rollouts through a shared TOPPRA entry point.

• Retimed results recompute poses and check the returned samples against joint bounds and waypoints.
• The shared solver now selects NumPy or Warp according to device and gradient needs.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Neural rollout samples] --> B{Constraints set?}
  B -- No --> C[Nominal timing]
  B -- Yes --> D[Shared TOPPRA retiming]
  D --> E[Resampled joints and timing]
  E --> F[Batched FK and result checks]
  F --> G[PlanResult]

Loading

_{Reviews (9) · Last reviewed commit: "perf(motion): run post-retime forward ki..."}

[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.

[greptile-apps]

Long rollouts slow retiming

With constraints enabled, a rollout that uses the default 240-step budget requests at least 24,100 TOPPRA gridpoints, compared with the kernel’s roughly 2,410-point default for that path. The shared entry point solves batch rows serially, so slow-to-converge or failed rollouts add planning latency that grows with batch size. A bounded or adaptive density would avoid paying this cost for every rollout sample.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/planners/neural_planner.py
Line: 696-699

Comment:
**Long rollouts slow retiming**

With constraints enabled, a rollout that uses the default 240-step budget requests at least 24,100 TOPPRA gridpoints, compared with the kernel’s roughly 2,410-point default for that path. The shared entry point solves batch rows serially, so slow-to-converge or failed rollouts add planning latency that grows with batch size. A bounded or adaptive density would avoid paying this cost for every rollout sample.

---

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

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[yangchen73]

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

[yangchen73]

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