Add differentiable time-optimal path parameterization

[Yuan-Xinyi]

Stack

• Layer: 4/5
• Base: xinyi/nmg-retiming-comparison
• Depends on: Pair a nominal and a retimed NMG arm in the benchmark #715 (and 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: Let NeuralPlanner retime with the differentiable TOPP #725

Description

Stacked on #715 only so that #725 can depend on both this solver and NeuralPlanner's retiming with a clean diff; nothing in this PR uses anything from the layers below.

This PR adds a differentiable time-optimal path parameterization to embodichain.compute.trajectory: the fastest rest-to-rest timing of a joint path under per-joint velocity and acceleration limits — the problem TOPP-RA solves — written so gradients reach the path.

The toppra library solves a small linear program at every grid point, so nothing downstream can differentiate through it. That matters for NMG: a policy that emits only path geometry cannot currently be trained against the timing the planner will actually give it, which is why #684 proposes teaching the network about velocity at all. With a differentiable parameterization the limits are met by construction and the policy only has to produce paths that are cheap to time.

Why it can be differentiated. With squared path speed x = sdot² and path acceleration u = sddot, every joint velocity and acceleration limit is linear in (x, u) at each grid point, so each stage is a two-variable linear program. Its feasible set projects onto x in closed form. The backward controllable-set sweep and the forward greedy sweep then reduce to min, max, division and square roots, which autograd and Warp's tape both differentiate almost everywhere.

What is added:

• parameterize_time_optimal() — speed profile on a sampled path. A Warp backend runs grid-parallel kernels for the constraint rows and one thread per environment for the two sequential sweeps; a torch implementation is the reference.
• retime_time_optimal() — end to end from waypoints: merges repeated waypoints, fits the not-a-knot cubic spline toppra.SplineInterpolator uses, parameterizes, and samples on a time grid. Every output, including the sample times, is differentiable with respect to the waypoints.
• Both toppra discretizations; interpolation, its default, is the default here.

Dependencies: none new. toppra and warp are already core dependencies; toppra and scipy are used only as references in tests.

Type of change

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

Measured evidence

It is TOPP-RA, not an approximation of it

Same path and same grid, against the toppra library, worst case over random paths, a real NMG rollout, and that rollout with repeated samples:

discretization	case	duration, relative	speed profile, max |Δx| / max x
interpolation	all cases	≤ 2.3e-7	≤ 7.5e-6
collocation	smooth paths	≤ 2.1e-7	≤ 4.5e-6
collocation	paths that must stop mid-way	≤ 3.4e-5	≤ 7.5e-6

The collocation residual is toppra's LP tolerance, located: 98% of it comes from points where the exact speed is zero, where toppra returns 2.8e-7 or 4e-8 instead of 0 and the square root amplifies it. The spline matches scipy.interpolate.CubicSpline(bc_type="not-a-knot") to 2e-11 for 2 to 20 knots, including scipy's two- and three-knot special cases.

An analytic check: two waypoints 0.5 rad apart on every joint, acceleration limit 10 — the optimum is a triangular speed profile lasting 2·√(0.5/10) = 0.4472 s. Measured: 0.4472 s.

The gradients are correct and usable

check	result
parameterization, autograd vs central finite differences (fp64)	160/160 agree, median relative error 1e-9
end to end retime_time_optimal, loss on duration and sampled q, q̇, q̈ (fp64)	40/40 within 1e-4, median 3e-8, both backends
Warp tape vs torch autograd (fp64, CPU and CUDA, both discretizations)	agree to 1e-12
fp32 gradient w.r.t. waypoints vs the fp64 reference	cosine 1.00000, relative norm error ≤ 2.5e-3, 99.4–99.7% of entries within 1%

Individual fp32 entries can differ more: where two joints nearly tie for the binding constraint, operation order decides which receives the gradient. That is the piecewise nature of the problem, not a defect of either backend, and it averages out at the waypoint level as the last row shows.

Usability, using only this API: a seven-joint path with its start, one via-point and its goal held fixed, optimized by Adam on d(duration)/d(waypoints):

	duration	peak |q̇| / limit	peak |q̈| (limit 10)
initial path	2.049 s	0.769	10.001
after 200 steps	0.887 s	0.741	10.000

Joint position, velocity and acceleration against time, before and after. The optimized path loses most of its direction reversals, joint 7 now reaches its 2.62 rad/s velocity limit, and the acceleration switches far less often while staying within ±10.

A near-stationary stretch was the hard case

The first formulation omitted one bound of the stage projection — that maximal acceleration must not force the speed below zero. It only binds where the path nearly stops, so smooth random paths all matched; on a real NMG rollout with a held tail, the shape a batched rollout produces when an environment converges early, it was 6% off. With the bound, that case matches to 8e-8.

The same case exposed a pre-existing issue in ToppraPlanner._toppra_solve_one_env: its deduplication keeps one trailing duplicate knot, which changes the whole spline and makes the retimed path reverse at the goal and take 6.6% longer. This module drops a held tail entirely, so holding at the goal leaves the motion unchanged, which a test asserts. The planner fix is left for a separate change.

Limits hold at grid points; between them, density decides

Peak acceleration utilization over the resampled output of ten random twelve-waypoint paths:

grid points	200	500	1000	1200 (default here)	2000	4000
peak |q̈| / limit	1.052	1.019	1.011	1.000	1.002	1.000

The default is 100 grid points per retained waypoint with a floor of 1000, keeping the overshoot to about one percent.

Speed

End to end, 12 waypoints × 7 joints, default grid, CUDA fp32, against the production path (_toppra_solve_one_env, CPU, forward only):

batch	this, forward	this, forward + backward	production
1	8.5 ms	39 ms	46 ms
64	10.7 ms	32 ms	1.3 s
1024	48 ms	107 ms	19 s (extrapolated)

The first call for a new joint count, discretization and dtype compiles its Warp kernels, 2–25 s depending on shape; Warp caches them on disk afterwards.

Review follow-up

Automated review found four real defects, fixed in 9b51ae1e; every fix below has a test that fails on the implementation before it:

• a move smaller than duplicate_tolerance was returned as a hold at the start, never reaching the goal;
• the velocity bound skipped any path derivative below the degeneracy threshold, which with a tight limit let a joint exceed it at the grid point — it now divides by a floored square instead;
• the square-root floor reported a small nonzero speed at rest, so the first sample was not at rest;
• a non-finite or negative duplicate_tolerance classified every path as stationary — it is now rejected.

A second round, in 45ab0cf4, caught that two of those first fixes were themselves wrong:

• answering the sub-tolerance move by starting at the first waypoint and holding the last moved the pose across a zero-length interval. Now only a row whose waypoints are all identical is stationary; any difference between first and last, however small, is timed as a real move;
• the square-root floor also treated a slow but real speed as rest — a tight velocity limit can put the squared speed below 1e-12 — and then timed those segments from the floor. Rest now means below the dtype's smallest normal number.

Scope and limitations

• Rest to rest only. Nonzero start or end speeds, or interior speed constraints, would need lower bounds on the controllable sets; not implemented.
• Limits are enforced at grid points, as in toppra; see the density table above.
• Warp kernels support up to 16 joints. Reverse mode needs every loop that updates a local to be unrolled, so the joint count is a compile-time constant and each loop is bounded by it; auto falls back to torch beyond.
• Not yet wired in. Nothing in lab calls this; NeuralPlanner retiming and NMG training integration are follow-ups.

Validation

• pytest tests/compute/test_trajectory_topp.py --run-gpu — 35 passed, CPU and CUDA.
• pytest tests/compute/ — 118 passed.
• black . — clean.
• python docs/scripts/check_api_docs.py — 2322/2322; the four new exports are documented on the curated embodichain.compute page. sphinx -b dummy — no warnings for the touched pages.
• context.py affected --base origin/main — motion-planning, through compute/trajectory/. Updated: the owner table gains a row for this module, and the retiming section states what it solves, its scope limits, the Warp joint cap, and the deduplication difference from ToppraPlanner. context.py check — ok.

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 time-optimal trajectory timing computation.

The PR is not yet safe to merge because tiny moves can still bypass joint acceleration limits.

Findings

1.  Tiny moves bypass acceleration limits ▶

Fix with agent prompt

### Issue 1
embodichain/compute/trajectory/topp.py:593-596
A float32 move smaller than `duplicate_tolerance` is now retained and timed. If its displacement is at most `1e-7`, the parameterizer treats its nonzero path derivative as inactive when checking acceleration. It can then return a successful trajectory whose joint acceleration exceeds the caller’s limit. For example, a `5e-8`-rad move with a 10-rad/s² limit can exceed that limit.

---

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

Summary

Adds differentiable, rest-to-rest time-optimal parameterization for joint paths, with Torch and Warp backends, waypoint retiming, tests, and API documentation.

_{Reviews (4) · Last reviewed commit: "fix(compute): time tiny moves and keep s..."}

[Yuan-Xinyi]

All four findings were valid; fixed in 9b51ae1e. Each fix has a test that fails on the previous implementation — 8 new test cases fail before the change and pass after.

1. Sub-tolerance move ended at the start — fixed. Such a row now starts at its first waypoint and holds its last, so it ends exactly at the requested goal. test_a_move_below_the_duplicate_tolerance_still_ends_at_the_goal.

2. Small nonzero dq/ds dropped from the velocity bound — fixed. The bound is now v² / max((dq/ds)², 1e-30) with no threshold, in both backends. The acceleration rows keep their threshold: treating a near-zero-coefficient row as the pure speed bound |b x| ≤ A is the conservative side, and the remaining effect is bounded by the other rows' acceleration range. test_a_tiny_path_derivative_still_bounds_the_speed runs your exact case (fp32, dq/ds = 1e-8, limit 1e-5) on both backends.

3. Nonzero speed at rest — fixed. Zero squared speed now gives exactly zero speed; the floor survives only inside the untaken branch, to keep the gradient finite, and in the segment-time denominator. test_samples_start_and_end_exactly_at_rest checks the first and last real samples on both backends. Agreement with toppra improved slightly as a result, since its profiles are also exactly zero there.

4. Non-finite tolerance — fixed. NaN, infinity and negative values are rejected.

[greptile-apps]

Tiny moves bypass acceleration limits A float32 move smaller than duplicate_tolerance is now retained and timed. If its displacement is at most 1e-7, the parameterizer treats its nonzero path derivative as inactive when checking acceleration. It can then return a successful trajectory whose joint acceleration exceeds the caller’s limit. For example, a 5e-8-rad move with a 10-rad/s² limit can exceed that limit.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/compute/trajectory/topp.py
Line: 593-596

Comment:
**Tiny moves bypass acceleration limits** A float32 move smaller than `duplicate_tolerance` is now retained and timed. If its displacement is at most `1e-7`, the parameterizer treats its nonzero path derivative as inactive when checking acceleration. It can then return a successful trajectory whose joint acceleration exceeds the caller’s limit. For example, a `5e-8`-rad move with a 10-rad/s² limit can exceed that limit.

---

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

[Yuan-Xinyi]

Closing in favor of #724, which solves the same problem with a stronger formulation.

Both implement differentiable TOPP-RA with the same closed-form reachability sweeps. The difference is the constraint set: this PR enforces toppra's grid-point constraints, which reproduces the toppra library to 1e-7 but can exceed the limits slightly between grid points; #724 bounds the whole interval through Bernstein coefficients and guarantees the limits everywhere. #724 also carries signed limits, both sampling modes, limit gradients and the ToppraPlanner integration, and removes the external dependency.

Two things from this PR are worth carrying over, and will follow as separate PRs against #724's branch:

• dropping a held tail instead of keeping a trailing duplicate knot, which otherwise lengthens a batched environment's motion by about 6% and adds a reversal at the goal;
• the kernel structure (grid-parallel constraint rows, unrolled joint loops, an fp32 path), which was about 10x faster at equal grid density.

The last automated finding here (tiny float32 moves bypassing the acceleration rows through a degeneracy threshold) does not apply to #724, whose rows have no such threshold.