- The paper introduces a propose–evaluate–correct planner that uses a collision critic and trace-conditioned reinforcement learning to revise unsafe motion tokens before execution.
- CorrectionPlanner reduces Waymax collision rates to 1.68 in reactive and 2.43 in non-reactive settings, exceeding 20% improvement over the strongest baselines while maintaining competitive progress.
- The method’s gains come from conditioning on the full correction trace rather than resampling, with a five-step correction budget and 0.75 critic threshold balancing safety, progression, and latency.
Overview
CorrectionPlanner introduces an explicit self-correction mechanism into autoregressive motion planning for autonomous driving (2603.15771). The core observation is that most learning-based planners, whether continuous regression, diffusion-based, or autoregressive next-token predictors, commit to an action without any internal evaluation or revision step. Once an unsafe trajectory is proposed, nothing in the pipeline prevents its execution. The paper addresses this gap by modeling planning as motion-token generation within a propose–evaluate–correct loop: a policy proposes an ego motion token, a learned collision critic predicts whether it will cause a collision within a short horizon (e.g., 2.5 seconds), and if the token is flagged unsafe, it is retained—not executed—and used to condition subsequent proposals until a safe token is produced or a correction budget is exhausted.
The sequence of rejected unsafe tokens is termed the correction trace, and it plays a role analogous to chain-of-thought reasoning traces in LLMs, except that reasoning occurs directly in motion-token space rather than in natural language. The authors motivate this design by noting that language-based self-reflection may not faithfully represent physical dynamics and may induce redundant reasoning steps.
Method
Planning is cast as a finite-horizon MDP where the action space consists of discretized motion tokens—trajectories segmented into 0.5-second intervals and clustered via K-disk clustering into a 1024-token vocabulary, following SMART (2603.15771). The policy is conditioned on ego history, agent histories, predicted agent motions, map/navigation information, and, during correction, the accumulated trace of rejected unsafe tokens encoded through a self-attention block.
A central design decision concerns simulation fidelity. Naive RL setups either replay logged agent trajectories or pre-predict all agent futures before planning; both are non-reactive and ignore how the ego's interventions—including corrections themselves—alter other agents' behavior. CorrectionPlanner therefore adopts model-based RL: a frozen pretrained world model built on SMART's architecture (a transformer decoder with temporal self-attention, agent-map cross-attention, and agent-agent cross-attention, using relative spatio-temporal embeddings for rotation/translation invariance) serves as a reactive simulator during rollout collection.
Two-stage training
Training proceeds in two stages:
- Imitation learning with corrective supervision. Beyond standard next-token prediction on expert trajectories, the IL phase exposes the policy to collision scenarios: when a proposed token would collide at the current timestep, the policy iteratively generates revised tokens conditioned on the correction trace, with cross-entropy loss against expert trajectories applied at every proposal. Importantly, IL only corrects current-step collisions; correcting future collisions requires counterfactual rollouts unavailable from expert data alone, which motivates the second stage.
- Model-based RL. Rollouts are generated in the frozen world model, including both executed tokens and intermediate correction tokens. The reward is deliberately sparse and two-term: progression rate relative to the expert, multiplied by a collision indicator, minus a penalty when a collision occurs. Optimization uses REINFORCE with KL regularization against the imitation policy, and the policy-gradient update is applied only to the executed token—the trace's intermediate tokens receive no gradient credit. This choice reflects that the trajectory-level reward is attributable to the executed action rather than the discarded proposals.
Collision critic
A binary classifier with temporal self-attention and ego-agent interaction layers predicts collision within k future planning steps, trained on policy/world-model rollout data with balanced safe/collision sampling. The critic gates execution at both training and inference time.
Experimental results
Closed-loop evaluations were conducted on WOMD via the Waymax simulator and on nuPlan (Val14, Test14-random, Test14-hard), under reactive (IDM agents) and non-reactive (log-replay) settings.
On Waymax, CorrectionPlanner achieves the lowest collision rates in both modes: 1.68 (reactive) and 2.43 (non-reactive), versus 2.36 and 3.02 for the best baselines—a reduction exceeding 20% relative to the strongest baseline—while maintaining competitive off-road rates and near-baseline progression. On nuPlan, the method attains state-of-the-art planning scores across most settings (e.g., 91.22 Val14 NR, 85.19 Val14 R, 90.41 Test-Random R), with SMART remaining best only on Test-Hard NR (76.39 vs. 75.37). Gains over SMART are consistently larger in reactive settings, which the authors interpret as evidence that reactive agent modeling better reflects interactive traffic.
| Method |
Collision (reactive) ↓ |
Off-road (reactive) ↓ |
Progression ↑ |
| PlanT |
2.94 |
1.65 |
95.85 |
| LatentDrive |
3.27 |
2.33 |
98.70 |
| SMART |
2.36 |
0.87 |
91.33 |
| CorrectionPlanner |
1.68 |
0.94 |
94.23 |
Ablations
Several ablations isolate the source of the improvement:
- Self-correction is not achievable by resampling. Rejection sampling under the same critic yields 2.09 collision on Waymo versus 1.68 for CorrectionPlanner; candidate selection over 10 sampled trajectories performs similarly to pure RL (2.14). Because rejection sampling redraws from the same conditional distribution, proposals cluster near the original unsafe token, whereas conditioning on the full trace shifts the proposal distribution away from collision regions. A variant conditioning only on the most recent unsafe token (2.01 collision) also underperforms full-trace conditioning, indicating the entire trace carries useful information.
- IL alone does not teach correction. IL with self-correction barely improves over vanilla IL (2.31 vs. 2.34 collision on Waymo); effective self-correction emerges only after RL training.
- Gains are not explained by extra compute. Running the same trained policy without self-correction at inference matches pure RL performance (2.20 collision), confirming that the collision reduction comes specifically from the correction loop.
- Threshold and budget trade-offs. At classification threshold 0.7, the critic achieves recall 0.93 but precision 0.21, triggering frequent early corrections and dropping progression to ~77.5; at threshold 0.8, precision rises to 0.87 but recall falls to 0.31 with first detection only ~1 second before impact. Threshold 0.75 with maximum correction length 5 gives the best safety-efficiency balance. Increasing the correction length monotonically reduces collisions within this regime, though lengths of 8–10 do not improve further due to unavoidable not-at-fault collisions and local trapping in unsafe regions. Setting the threshold to zero (correcting every step) drastically degrades progression, showing that gains come from selective, not constant, correction.
Latency measurements show the deployed configuration (threshold 0.75, length 5) incurs 0.434s per trajectory on an H800 GPU, comparable to the base SFT policy's 0.357s and SMART's 0.329s, because self-correction triggers rarely. Zero-shot transfer of a nuPlan-trained policy to Waymax also shows that enabling self-correction reduces collision from 2.93 to 2.62, suggesting the mechanism generalizes across datasets.
Qualitative examples exhibit diverse corrective behaviors: widening left turns preemptively to avoid not-at-fault collisions, aborting aggressive lane changes, and modulating turn speed rather than braking.
Limitations and open questions
The paper concedes several constraints. First, the collision critic's precision-recall trade-off directly bounds system behavior: at high precision, roughly 70% of potential collisions go undetected, and the authors identify critic calibration as future work. Second, the reward uses only progression and collision; the observed slight degradation in progression stems from yielding-and-re-accelerating behaviors, which the authors accept as a safety-for-efficiency trade-off but do not optimize explicitly. Third, longer correction budgets plateau due to not-at-fault collisions and local trapping in unsafe regions, indicating the correction trace cannot always escape genuinely infeasible situations. Fourth, the method depends on the fidelity of the frozen world model for counterfactual rollouts—an assumption whose accuracy bounds what RL can teach about correction. Finally, the framework is demonstrated only for collision avoidance; extending self-correction to other metrics (comfort, off-road, traffic rules) remains open, as does the question of whether gradient attribution solely to executed tokens is optimal versus crediting intermediate trace tokens.
Conclusion
CorrectionPlanner demonstrates that autoregressive driving planners can acquire an explicit, learned self-correction capability operating in motion-token space, analogous to reasoning traces in LLMs. Through a two-stage scheme combining corrective imitation learning with model-based RL over a reactive world model, the planner reduces Waymax collision rates by more than 20% over the best baseline while achieving state-of-the-art planning scores on nuPlan, with negligible added latency. The ablations establish that the benefit derives specifically from trace-conditioned distribution shifting rather than resampling, candidate selection, or additional forward passes, positioning trace-based self-correction as a concrete mechanism for safety improvement in learned planners.
