Let NeuralPlanner retime with the differentiable TOPP

[Yuan-Xinyi]

Stack

• Layer: 5/5
• Base: xinyi/differentiable-topp
• Depends on: Add differentiable time-optimal path parameterization #723 (and Pair a nominal and a retimed NMG arm in the benchmark #715, Retime NMG rollouts under joint velocity and acceleration limits #714, Measure dynamic limits and mark nominal timing in the NMG benchmark #713 below it)
• Followed by: none

Description

NeuralPlannerCfg.constraints (#714) retimes the NMG rollout with the toppra library, one environment at a time on the CPU. #723 adds compute.trajectory.retime_time_optimal, which solves the same time-optimal problem batched on the planner's device and differentiably. This PR lets the planner choose between them.

• NeuralPlannerCfg.retime_backend: toppra (default, so existing behavior is unchanged) or differentiable. An unknown value is rejected at construction rather than at the first constrained plan.
• Scalar limits apply to every joint the policy drives; a per-joint limit fixes how many joints are retimed, and any further control-part joints are held.
• The benchmark's NMG adapter forwards retime_backend.

Nothing is trained through the solver yet. This only replaces the forward computation; using its gradients in NMG training is a separate change.

Dependencies: none.

Refs #684

Type of change

• Enhancement (non-breaking change which improves an existing functionality)

Measured evidence

Real NMG checkpoint (swa3.onnx, K=5) on Franka FR3.

Same trajectories, less time

smoke suite, 9 free-space cases at W=1/3/5, 3 measured trials, both batch sizes:

arm	motion_valid	v / a / j violations	duration	plan ms, median (B=1 / B=8)
no retiming	100%	100% / 100% / 100%	N/A	213 / 161
retime_backend: toppra	100%	0 / 0 / 0	0.697–1.541 s	283 / 613
retime_backend: differentiable	100%	0 / 0 / 0	0.697–1.542 s	202 / 283

Validity, violations and limit utilization are identical and durations agree within a millisecond. The differentiable backend adds almost nothing over the unretimed rollout; toppra grows with the batch because it solves each environment separately.

The benchmark's B=8 cases are one problem replicated eight times — a throughput measurement — so every environment converges on the same step. That hides the one place the backends differ, which needs a batch of different problems.

The backends differ when an environment finishes early

Eight environments holding three distinct problems, each planned alone and then inside the batch:

environment	toppra alone → in batch	differentiable alone → in batch
converges first	0.978 → 1.038 s (+6.2%)	0.979 → 0.978 s (0.0%)
converges second	1.116 → 1.173 s (+5.1%)	1.117 → 1.117 s (0.0%)
converges last	1.711 → 1.711 s	1.712 → 1.712 s

Solid: planned alone. Dashed: the same problem inside the batch. With toppra the batched trajectory departs after about 0.75 s, joint 7 gains an extra reversal before the goal, and the motion ends 60 ms later. With the differentiable backend the two coincide.

An environment that converges first holds its pose while the others finish, so its rollout ends in repeated samples. toppra's deduplication in _toppra_solve_one_env keeps one trailing duplicate knot, which changes the spline. With toppra, an environment's trajectory depends on which other environments share its batch. The differentiable solver drops the held tail. The toppra deduplication itself is left for a separate fix.

Time for the batch of eight: toppra 837 ms, of which retiming 506 ms; differentiable 349 ms, of which retiming 9 ms.

Validation

• pytest tests/sim/motion/planners/test_neural_planner.py tests/sim/motion/planners/test_neural_batched.py tests/sim/motion/test_motion_generator.py tests/benchmark/motion_generation/ tests/compute/test_trajectory_topp.py — 208 passed, 3 skipped.
• New tests run both backends on a saturating rollout: limits respected, every joint's start and end match the rollout, durations agree within 2%, and an unknown backend is rejected.
• black . — clean. python docs/scripts/check_api_docs.py — aligned; no new exports, the new config field is documented by autodoc.
• context.py affected — motion-planning and simulation-system. planner-details.md's neural-adapter section now describes retime_backend and where the backends differ. simulation-system matches only through the tests/sim/ watch path. context.py check — ok.

A bug the unit tests missed. An earlier version inferred the joint count from the length of the limits, so scalar limits retimed only joint 0 and the other joints jumped to their final values at t = 0. The first unit test passed because it checked only limits and shapes. Running the real benchmark exposed it. The test now checks each joint's start and end against the rollout, and fails on that version.

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: 4/5

[Medium risk] Adds a retiming backend option to the motion planner.

The PR should not merge until short per-joint limit lists are rejected or handled without changing a controlled joint’s starting position.

Findings

1.  Short limits cause joint jumps ▶
2.  Key backend cases untested ▶

Fix with agent prompt

### Issue 1
embodichain/lab/sim/motion/planners/neural_planner.py:765-770
When a per-joint limit list has fewer than seven entries, the policy can still move all seven joints during the rollout, but this code retimes only the joints covered by the list. It then holds each remaining joint at its final rollout position from the first output sample and reports zero velocity for it. With six limits, joint 7 therefore jumps to its endpoint at time zero instead of following the planned path. Reject limit lists that do not cover every policy-driven joint.

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!

### Issue 2
tests/sim/motion/planners/test_neural_planner.py:510-517
These backend tests use only scalar limits and a single-environment rollout. Add a per-joint-limit case and a constrained batch where one environment finishes early. Without those cases, tests will not catch regressions in the new joint-counting logic or the held-tail behavior that distinguishes this backend.

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

Adds a configurable differentiable TOPP retiming backend to NeuralPlanner, forwards the choice through the NMG benchmark adapter, and documents and tests basic backend parity.

• Short per-joint limit lists can cause a controlled joint to jump to its endpoint.
• Integration tests do not cover per-joint limits or early-held rows in a constrained batch.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[NeuralPlanner rollout] --> B{retime_backend}
  B -->|toppra| C[Per-environment CPU retiming]
  B -->|differentiable| D[Batched TOPP retiming]
  C --> E[FK and trajectory checks]
  D --> E
  E --> F[Timed PlanResult]

Loading

_{Reviews (1) · Last reviewed commit: "feat(motion): let NeuralPlanner retime w..."}

[greptile-apps]

Short limits cause joint jumps When a per-joint limit list has fewer than seven entries, the policy can still move all seven joints during the rollout, but this code retimes only the joints covered by the list. It then holds each remaining joint at its final rollout position from the first output sample and reports zero velocity for it. With six limits, joint 7 therefore jumps to its endpoint at time zero instead of following the planned path. Reject limit lists that do not cover every policy-driven joint.

Knowledge Base Used:

• Motion planning and kinematics
• Trajectory timing and velocity tracking

Prompt To Fix With AI

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

Comment:
**Short limits cause joint jumps** When a per-joint limit list has fewer than seven entries, the policy can still move all seven joints during the rollout, but this code retimes only the joints covered by the list. It then holds each remaining joint at its final rollout position from the first output sample and reports zero velocity for it. With six limits, joint 7 therefore jumps to its endpoint at time zero instead of following the planned path. Reject limit lists that do not cover every policy-driven joint.

**Knowledge Base Used:**
- [Motion planning and kinematics](https://app.greptile.com/dexforce/-/custom-context/knowledge-base/dexforce/embodichain/-/docs/motion-planning-and-kinematics.md)
- [Trajectory timing and velocity tracking](https://app.greptile.com/dexforce/-/custom-context/knowledge-base/dexforce/embodichain/-/docs/trajectory-timing-and-tracking.md)

---

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!

[greptile-apps]

Key backend cases untested These backend tests use only scalar limits and a single-environment rollout. Add a per-joint-limit case and a constrained batch where one environment finishes early. Without those cases, tests will not catch regressions in the new joint-counting logic or the held-tail behavior that distinguishes this backend.

Prompt To Fix With AI

This is a comment left during a code review.
Path: tests/sim/motion/planners/test_neural_planner.py
Line: 510-517

Comment:
**Key backend cases untested** These backend tests use only scalar limits and a single-environment rollout. Add a per-joint-limit case and a constrained batch where one environment finishes early. Without those cases, tests will not catch regressions in the new joint-counting logic or the held-tail behavior that distinguishes this backend.

---

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!

[Yuan-Xinyi]

Closing: this switch selected between toppra and #723's solver, and #723 is closed in favor of #724. Once #724 lands, ToppraPlanner's own retiming is batched and differentiable, so NeuralPlanner needs no separate backend option. The held-tail comparison from this PR's description carries over to the deduplication follow-up on #724.