[Proposal] Make joint velocity and acceleration first-class in the NMG waypoint contract

[Yuan-Xinyi]

Proposal

Make joint velocity and acceleration first-class in the NMG waypoint contract, both as limits the rollout respects and as targets a waypoint can ask for.

The main decisions are:

• Land the dynamic-limit metrics and correct dt semantics before any mechanism, since nothing else can be accepted without them.
• Bound velocity and acceleration through output-side time parameterization, reusing the existing scalar time law rather than adding a second one.
• Add an in-rollout clamp as a differentiable operator shared by training and inference, first as a diagnostic on existing checkpoints.
• Consume the already-declared PlanState.qvel and PlanState.qacc as waypoint targets, mirroring the existing joint-position modality.
• Group every observation-layout change into one retraining and re-export rather than three.
• Keep current defaults unchanged; every new mechanism is opt-in.

Related work: #437, #576.

Current behavior

neural_planner.py integrates the policy as qpos += action * action_scale and clamps joint position only. Velocity and acceleration are never bounded; they are reported by _compute_vel_acc_via_finite_diff over a constant NeuralPlannerCfg.dt = 0.01.

With action in [-1, 1] and action_scale = 0.2, the rollout already carries an implicit velocity bound, but a constant, isotropic, robot-independent one: 20 rad/s, against 2.62 rad/s for the slowest joint of the FR3 asset the benchmark uses (Franka/Panda/PandaWithHand.urdf, joints 1-4 at 2.62, joints 5-7 at 5.26 / 4.18 / 5.26). The action can swing from +1 to -1 between steps, so acceleration is unbounded within a step.

BENCHMARK_DESIGN.md section 2.1 already records this as a known capability boundary, and section 4.2 already names the dynamic-limit metrics and the dynamic_limits_satisfied_when_applicable term of motion_valid. Neither is implemented.

PlanState.qvel and PlanState.qacc exist and are documented as target fields. No planner reads them; the only reference under embodichain/lab/sim/motion/ is base_planner.py:112, for batch-size inference.

Motivation

The reported derivatives are consistent bookkeeping on a time axis no arm can execute. NMG native timing therefore cannot be compared fairly against cuRobo native timing, motion_valid has a term that can never evaluate, and the trajectory NMG hands to an unspecified controller carries no feasibility guarantee — which matters precisely because the NMG contract ends at the joint trajectory, with nothing downstream to enforce limits.

On the target side, every waypoint is a pose-or-configuration constraint with no say over arrival speed. That excludes via-points passed through at speed rather than braked into, a specified approach velocity before contact or insertion, and an explicit rest condition at the final waypoint, which today is an emergent side effect of the hold mechanism rather than a stated goal.

Velocity targets are not symmetric with joint-position targets. Since qd_t = action_t * action_scale / dt, the policy does control velocity directly and last_action is already a velocity proxy. But a target in rad/s needs dt to mean something, the arrival check needs a causal velocity definition rather than the central difference used for reporting, and position-plus-velocity is a coupled terminal condition that costs arrival rate unless trained for. This is what fixes the ordering below.

Items 1-4 run on existing checkpoints. Items 5-7 change _WaypointObservationLayout.block_widths, which alters the fingerprint that _validate_observation_metadata() enforces, so they require retraining and re-export and should share one such event.

1. Fix dt semantics and implement the dynamic-limit metrics

Report PlanResult.dt as unavailable rather than 0.01 when no time parameterization has been applied. Implement the metrics section 4.2 already names: velocity, acceleration and jerk violation rates, maxima and means, and dynamic_limits_satisfied.

This provides no guarantee; it is the measuring instrument and a prerequisite for both halves of the proposal. Metrics are recomputed from the output trajectory, consistent with match_ordered_joint_waypoints already refusing to trust PlanResult.success.

2. Output-side time parameterization

Add NeuralPlannerCfg.constraints: dict | None = None, defaulting to current behavior. When set, treat the rollout output as a geometric path and recompute its timing under velocity, acceleration and jerk limits before deriving velocities, reusing the existing ToppraPlanner (embodichain/lab/sim/motion/planners/toppra_planner.py, toppra==0.6.3), which already accepts an (N, DOF) joint path with per-joint velocity and acceleration limits. The rest-to-rest scalar time law in _scalar_time_law.py is the wrong tool here: it would brake at every one of the hundreds of rollout samples.

Hard guarantee, and dt becomes a per-sample array so duration becomes meaningful. Path geometry is unchanged, so jerk remains the policy's and retiming can only absorb it as time; expect duration to grow by close to an order of magnitude.

3. State-dependent clamping in the rollout

Bound the applied joint delta before it is integrated:

dq_max = torch.minimum(v_max * dt, qd_prev.abs() * dt + a_max * dt**2)
dq = torch.clamp(action * action_scale, -dq_max, dq_max)

last_action must record the clamped value, or the closed-loop observation reports an action that was never applied.

Hard guarantee on a uniform time axis. The policy never saw this clamp, so arrival degrades; measuring that degradation is the purpose, since it quantifies how much the current policy depends on out-of-limit actions and therefore decides whether items 5-7 justify their training budget. Land it as a diagnostic mode, not a default.

4. Self-consistent action_scale, as a control experiment

Set action_scale <= min(v_max) * dt, about 0.026 for the FR3 asset at dt = 0.01, adding no mechanism. Velocity is satisfied; acceleration is not, since the action can still reverse in one step. Per-step displacement drops roughly tenfold, so steps_per_waypoint: 30 no longer reaches its waypoints and max_steps must grow with it.

Proposed as an experiment reported alongside item 2, not as a shipped configuration change. It separates "the policy needs this speed" from "dt is mislabeled".

5. Limits and current velocity in the observation

Add qvel, qvel_limit and qacc_limit blocks to the observation layout.

No guarantee on its own: encoding a quantity lets the policy see it, not obey it, which is exactly the property the existing waypoint joint constraint has. Train it together with item 3's clamp, which is closed-form and differentiable, so the policy learns inside the feasible set, inference applies the same operator, and item 3's arrival penalty disappears. Conditioning on the limit values is what keeps the result from being Franka-specific.

6. Waypoint joint-velocity targets

Give velocity the treatment joint_mask gives joint position, consuming PlanState.qvel:

• a waypoint_qvel block of width K * 7, a vel_mask of width K, and a waypoint_qvel_err block under use_relative_obs, mirroring waypoint_joint / joint_mask / waypoint_joint_err, with zero as the masked value;
• a vel_eps term conjoined into _is_active_reached, compared against a causal velocity definition rather than the reporting central difference;
• parse-time rejection of a target exceeding the configured velocity limit, rather than a silent non-arrival;
• relaxation of the if / elif in _parse_waypoints that forces a slot to be EEF xor joint. The observation layout already supports independent per-modality masks; only the parser forbids it.

Soft guarantee, the same as every other waypoint target.

7. Waypoint joint-acceleration targets

The same structure with waypoint_qacc, acc_mask and acc_eps, consuming PlanState.qacc. Kept separable and lower priority: acceleration is controlled through the action difference rather than the action, so it is one order harder to hit, and its realistic use is continuity and rest conditions rather than free targets. Skip it if item 6's arrival rate is poor.

Cartesian twist targets are the more natural specification for contact tasks but need the Jacobian inside the rollout loop; out of scope until joint-space targets are shown to work.

Suggested implementation sites

• embodichain/lab/sim/motion/planners/neural_planner.py — NeuralPlannerCfg, _WaypointObservationLayout, _parse_waypoints, _build_obs, _is_active_reached, and the rollout loop.
• embodichain/lab/sim/motion/planners/_scalar_time_law.py and embodichain/compute/trajectory/timing.py — reused, not modified.
• scripts/benchmark/motion_generation/metrics/trajectory.py, config.py and the suite YAMLs — metrics, thresholds and opt-in configuration.
• embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py — _set_joint_targets_kernel applies the same clamp(current + action * scale, lo, hi) as the planner, so items 3, 5, 6 and 7 must change both sites or training and inference semantics diverge.
• tests/sim/motion/planners/test_neural_planner.py and tests/benchmark/motion_generation/ — focused coverage per item.

Velocity limits resolve correctly: Articulation._qvel_limits is populated from entity.get_joint_velocity_limit(), and the benchmark's Franka asset is an FR3 whose URDF declares 2.62 / 2.62 / 2.62 / 2.62 / 5.26 / 4.18 / 5.26 rad/s. The max_velocity=1e10 default at embodichain/lab/sim/cfg/robot.py:70 is drive configuration and does not feed these limits.

Acceleration limits have no in-repository source: there is no qacc_limits anywhere and URDF carries no acceleration field, so they must be supplied by configuration, following the explicit-value convention of TrapezoidalPlanOptions.constraints. Franka publishes them for this arm -- FR3 joint-space limits are 10 rad/s^2 and 5000 rad/s^3 uniformly across all seven joints (robot specifications), and the same page's velocity row matches the asset URDF exactly, which confirms the asset is a genuine FR3. Suites should state these values with that citation, and metrics must report acceleration results as N/A rather than as satisfied when limits are unset.

Naming reuses existing vocabulary: the velocity / acceleration / jerk keys from TrapezoidalPlanOptions, the waypoint_* / *_mask / *_eps pattern from the position, rotation and joint modalities, and section 4.2's metric terms.

Acceptance criteria

• PlanResult.dt is reported as unavailable when no time parameterization has been applied, rather than as a nominal constant.
• Section 4.2's velocity, acceleration and jerk metrics and dynamic_limits_satisfied are computed from output trajectories and reported by the existing suites.
• NeuralPlannerCfg.constraints defaults to None and reproduces current behavior; previously published NMG numbers remain reproducible.
• With constraints set, dynamic_limits_satisfied is 100%, and arrival rate, duration and jerk are reported against the unconstrained baseline.
• The in-rollout clamp is a single operator shared by the planner rollout and the APG training kernel, and last_action records the clamped value.
• Velocity and acceleration limits resolve from robot.get_qvel_limits() for the benchmark assets, with an explicit error rather than a silent no-op when limits are unpopulated.
• PlanState.qvel and PlanState.qacc are consumed as waypoint targets, and a target exceeding the configured limit is rejected at parse time.
• A waypoint slot can carry EEF, joint-position and velocity constraints independently; the observation layout and the parser agree on this.
• Velocity-target arrival rate and velocity error at arrival are reported beside position arrival rate on the same cases, not averaged together.
• Any observation-layout change updates the fingerprint and the exported ONNX metadata together, and a mismatched export is still rejected.
• Focused tests cover retiming, the clamp, velocity-target parsing and arrival, and the limits-unpopulated case.

Checklist

• I have checked that there is no similar issue in the repo (required)