Source-linked AI summary
Aligning Multi-Trajectory Supervision with Policy Optimization for VLA Driving
Tian Zhang, Zhuo Huang, Hongrui Ye, Yu Wu, Zengmao Wang, Kaixuan Zhou
TL;DR
VLA driving needs trajectory supervision that remains compatible with the policy being optimized, because high-scoring imitation targets can degrade later GRPO. AMPT aligns candidate selection, feasibility-first credit assignment, and dynamic teacher distillation, achieving strong NAVSIM results and recovering many failed scenes.
Problem
Trajectory sets that improve imitation can reduce feasible, high-quality rollouts during GRPO because aggregate quality does not ensure policy compatibility or preserve safety structure.
Method
AMPT filters candidates by Pareto quality, frozen-policy compatibility, and local feasibility, then uses feasibility-first Pareto credit and repeatedly validated teacher distillation.
Results
91.4 PDMS and 89.1 EPDMS are achieved on NAVSIM v1 and v2, respectively, while 440 of 658 initially failed scenes are recovered.
Takeaways & Limitations
Expanded supervision should be judged by the policy distribution it induces for downstream optimization rather than by trajectory scores alone.
Takeaways & Limitations
The compatibility test uses finite policy samples and teacher validation depends on an offline planning evaluator; future work targets learned estimates and interactive closed-loop environments.
Abstract
from arXiv · showhide
Vision-language-action (VLA) driving methods increasingly combine multi-trajectory imitation learning with group-relative policy optimization (GRPO), making trajectory selection critical to final performance. However, some high-scoring trajectories that improve imitation can degrade subsequent GRPO by inducing advantage estimates misaligned with the current policy's feasible behavior distribution, driving updates away from safe and compliant behaviors. To address this, we propose a novel framework that aligns multi-trajectory supervision with policy optimization. To address the policy gradient bias induced by infeasible noisy trajectories outside the feasible region, augmented trajectories are constrained to a neighboring manifold of the ground-truth feasible region, and a Pareto-optimality criterion is adopted in place of the conventional aggregate score, retaining only non-dominated candidates and thereby filtering out conflicting samples at the source. To ensure that expanded trajectory supervision is effectively absorbed during policy optimization, we introduce two complementary mechanisms: feasibility-first advantage assignment and dynamic distillation. The former adapts Pareto credit to the feasibility composition of each rollout group and guides fully infeasible groups toward safe references. The latter updates teacher trajectories across refinement rounds to continually transfer useful supervision. Together, they progressively translate the benefits of expanded supervision into policy improvement. On NAVSIM v1 and v2, our method achieves 91.4 PDMS and 89.1 EPDMS, respectively, under single-trajectory inference, and recovers 440 of 658 initially failed scenes, 11.1\% higher than the original GRPO baseline.
Introduction
The paper identifies a supervision–optimization mismatch: trajectories that improve imitation may be incompatible with the policy distribution needed for safe and effective GRPO. AMPT addresses this by aligning trajectory selection, credit assignment, and policy refinement.
- Multi-trajectory selection based on aggregate quality, physical validity, or geometric diversity does not establish downstream usefulness for GRPO.
- 87.2 and 87.5 PDMS imitation checkpoints deteriorate to 85.8 and 86.7 after GRPO under score-based and Pareto-based supervision.
- 90.3 and 91.1 PDMS are reached after GRPO by GT-only supervision and policy-compatible selection, respectively.
- AMPT combines policy-compatible multi-trajectory supervision, feasibility-first Pareto GRPO, and adaptive Pareto-guided policy refinement.The framework screens candidates for compatibility and feasibility, adapts rollout credit, and repeatedly updates teacher trajectories.
- 91.4 PDMS and 89.1 EPDMS are achieved on NAVSIM v1 and v2, while 440 of 658 initially failed scenes are recovered.The recovery count is 73 higher than the original GRPO baseline.
Related Work
Related work broadens driving-policy behavior distributions and applies GRPO or multi-objective feedback, while AMPT focuses on selecting and refining trajectories that remain useful to the evolving policy. It produces a single trajectory without test-time scoring or reranking.
- Multi-Trajectory Supervision for Driving Policies: Multi-trajectory driving methods represent plausible futures using motion anchors, proposal sets, learned scoring, or simulator-generated alternatives.
- Multi-Trajectory Supervision for Driving Policies: Curious-VLA broadens the policy distribution, whereas AMPT studies which additional trajectories provide useful supervision for downstream optimization.
- Reinforcement Post-Training and Policy Refinement: Prior reinforcement post-training methods apply GRPO and expose metric-specific or multi-objective feedback during policy optimization.
- Reinforcement Post-Training and Policy Refinement: AMPT reassesses teachers after every round, converts them into locally validated targets, and distills them with behavior retention.
- Reinforcement Post-Training and Policy Refinement: The resulting policy predicts a single trajectory without a test-time scorer or candidate reranking.
Preliminaries
The preliminaries define a diffusion-based VLA trajectory policy trained first by imitation and then improved with group-relative policy optimization. Standard GRPO derives advantages from aggregate rewards, which can obscure safety and compliance trade-offs.
- Driving VLA policy: A driving VLA maps multimodal scene context to a continuous future trajectory through a conditional diffusion policy πθ(τ | o).
- Imitation learning: The policy is initialized with logged demonstrations before reinforcement-learning refinement.
- Group-relative policy optimization: GRPO samples G trajectories for one scene and evaluates each with an aggregate planning reward Ri.
- Group-relative policy optimization: Standard GRPO estimates advantages from rewards within the sampled group without an additional value network.
- Group-relative policy optimization: Aggregate-reward normalization does not distinguish safety violations from meaningful trade-offs among feasible trajectories.
Method
AMPT constructs policy-compatible, locally feasible supervision, then applies feasibility-first Pareto credit assignment and iterative local distillation. The method keeps safety and reference-level performance ahead of efficiency improvements during optimization.
- Overall pipeline: AMPT trains and freezes πGT, uses PC-MTS to produce πinitial, applies FF-PGRPO, and repeatedly mines improvements with APR.
- Metric hierarchy: NAVSIM feasibility and compliance checks determine which metrics can be compared within the feasible region.
- Policy-Compatible Multi-Trajectory Supervision: PC-MTS retains candidates on the scene-wise Pareto front after quality, safety, compliance, compatibility, and local-feasibility screening.
- Policy-Compatible Multi-Trajectory Supervision: Compatibility is estimated by average waypoint distance to the K nearest frozen-policy rollouts, with command-specific thresholds calibrated from held-out samples.
- Feasibility-First Pareto GRPO: FF-PGRPO selects a single executable scene reference from logged, retained, and predicted trajectories, then constructs a Pareto front among admissible rollouts.
- Feasibility-First Pareto GRPO: Feasible Pareto rollouts may receive positive credit, non-Pareto feasible rollouts receive non-positive credit, and infeasible rollouts receive stronger penalties.
- Feasibility-First Pareto GRPO: Fully infeasible groups receive no positive examples, safe-reference supervision, and downweighted updates while retaining regularization toward πinitial.
- Adaptive Pareto-guided Policy Refinement: APR reassesses multi-source teachers against the current policy, interpolates them into verified local targets, and distills accepted gains with policy retention.
Experiments
Experiments show that AMPT improves single-trajectory planning across NAVSIM versions, while its components address trajectory-set construction, feasibility-first optimization, and iterative refinement. The method also improves recovery of initially failed scenes.
- NAVSIM v1: 91.4 PDMS on NAVSIM v1 is state of the art among compared methods, 0.2 points above DiffusionDriveV2.AMPT improves NC, DAC, and TTC, including a 1.2-point TTC gain, while retaining competitive EP.
- NAVSIM v2: 89.1 EPDMS on NAVSIM v2 exceeds Drive-JEPA by 1.3 points under the expanded evaluation protocol.The gain is primarily associated with higher EP and EC, while several other components remain near saturated ranges.
- Trajectory-set construction: Score and Pareto supervision improve imitation checkpoints to 87.2 and 87.5 but decline to 85.8 and 86.7 after identical scalar GRPO.PC-MTS reaches only 86.9 after imitation yet improves to 91.1 after GRPO, demonstrating that imitation quality alone does not identify the best initialization.
- Component ablation: FF-PGRPO raises DAC from 94.7 to 98.0, TTC from 94.2 to 95.8, and EP from 80.9 to 85.9, reaching 91.0 PDMS.The joint gains in safety-sensitive components and progress indicate that efficiency is improved without relaxing feasibility.
- Iterative refinement: APR raises PDMS from 91.10 to 91.45 across three refinement rounds, with gains remaining monotonic.APR increases EP from 86.0 to 86.7 while NC, DAC, TTC, and comfort remain unchanged.
- Failure recovery: AMPT recovers 440 of 658 initially failed scenes versus 367 for original scalar GRPO, increasing recovery from 55.8% to 66.9%.This corresponds to 73 additional recovered scenes and an absolute gain of 11.1 percentage points.
Conclusion
AMPT addresses the supervision–optimization mismatch in VLA driving and reports strong single-trajectory performance across NAVSIM protocols. Its compatibility test relies on finite policy samples, while teacher validation depends on an offline planning evaluator.
- AMPT combines PC-MTS, FF-PGRPO, and APR to align multi-trajectory supervision with downstream policy optimization.
- 91.4 PDMS on NAVSIM v1 and 89.1 EPDMS on NAVSIM v2 are achieved under single-trajectory inference.
- 440 of 658 initially failed scenes are recovered.
- The compatibility test uses finite policy samples, and teacher validation depends on an offline planning evaluator.
Optimization for VLA Driving”
AMPT focuses on improving VLA driving through training-data construction and policy optimization, while leaving several representation-learning components outside its scope.
- AMPT improves VLA driving through multi-trajectory candidate selection and trajectory-policy optimization.
- The framework adds no external driving data, does not separately fine-tune the VLM backbone, and does not increase deployed policy parameters.
- The study does not address feature extraction, representation learning, feature fusion, or multimodal semantic alignment.
A Additional Method Details
AMPT selects policy-compatible, locally feasible trajectories, assigns GRPO credit with feasibility-first Pareto rules, and dynamically refines teachers around the evolving policy.
- PC-MTS: PC-MTS filters candidates by safety, quality, Pareto non-dominance, compatibility with a frozen GT-only policy, and local feasibility.
- FF-PGRPO: Feasibility and protected metrics are checked before Pareto comparison, so efficiency gains cannot offset safety or compliance regressions.
- PC-MTS: Each scene contributes exactly one supervision target, preventing scenes with more candidates from receiving greater loss weight.
- Reference construction: Scene references use one highest-scoring feasible trajectory, preserving coherent aggregate and component metrics from the same executable plan.
- FF-PGRPO: Fully infeasible groups receive no positive rollout credit and instead use violation-aware advantages plus safe-reference recovery supervision.
- APR: APR interpolates distant teachers into local targets, re-evaluates them exactly, weights accepted targets by verified gains, and refreshes teachers after each refinement round.
B Implementation and Evaluation Protocol
The implementation preserves the underlying VLA backbone and single-trajectory inference while applying AMPT through training-only candidate construction, rollout credit assignment, and policy refinement.
- Implementation: AMPT reuses the InternVL3-2B backbone and ReCogDrive diffusion planner without adding deployed trainable modules or external driving data.
- Implementation: The framework changes trajectory targets, rollout credit, and policy updates while preserving the perception and language representation stack.
- Evaluation protocol: NAVSIM v1 evaluates 12,138 valid scored scenes and NAVSIM v2 evaluates 12,146 scenes, with exactly one predicted trajectory per scene.
- Results: The final v1 score is supported jointly by NC, DAC, and TTC, while v2 retains near-saturated DDC, TLC, and TTC with EP at 89.45.
- Recovery evaluation: The hard subset contains 658 initially failed scenes, and recovery counts scenes whose evaluated trajectory obtains non-zero PDMS.
- Results: 91.4 PDMS is reported on NAVSIM v1 and 89.1 EPDMS on NAVSIM v2 for the final single-trajectory policy.
- Training pipeline: The training algorithms use staged policy checkpoints, including PC-MTS initialization, FF-PGRPO optimization, and APR refinement.
D Analysis of Policy-Compatible Supervision
Trajectory quality at imitation time does not reliably predict post-training utility. Policy-compatible selection instead preserves feasible behavior support and enables stronger subsequent optimization.
- Higher imitation performance does not predict a better optimization outcome.
- Score and Pareto supervision improve imitation but reduce scores after scalar GRPO, revealing a ranking reversal.The imitation gains are 0.8 and 1.1 points, followed by GRPO declines of 1.4 and 0.8 points, respectively.
- PC-MTS changes the learned sampling distribution before GRPO rather than merely maximizing isolated trajectory quality.Its one-target-per-scene rule prevents larger candidate sets from mechanically increasing supervision.
- Candidate-to-GT distance and policy compatibility measure different properties of a trajectory.Rollout-bank references assess whether a candidate lies near behavior modes the learner can absorb.
- PC-MTS favors local extensions, while APR later recalibrates compatibility around the improved policy.Together they form a curriculum whose boundary moves with policy capability.
E Analysis of Feasibility-First Pareto GRPO
Feasibility-first credit assignment protects safety and compliance before rewarding efficiency, while APR iteratively refines teacher targets as the policy improves. The resulting pipeline improves multiple driving metrics and substantially increases recovery on failed scenes, subject to evaluator and candidate-pool limits.
- FF-PGRPO raises PDMS from 86.5 to 91.0 while improving DAC, TTC, and EP over GT-only training.The gains are 3.3 DAC points, 1.6 TTC points, and 5.03 EP points.
- FF-PGRPO enforces zero positive credit for rollouts violating encoded feasibility and reference conditions.Protected violations are excluded from positive reinforcement rather than traded against progress on a common scale.
- Safe-reference supervision supplies a feasible target when every rollout in a group is infeasible.Group downweighting prevents hard scenes from dominating broader updates.
- APR increases PDMS monotonically from 91.10 to 91.45 across three rounds, with the final gain concentrated in EP.Displayed NC, DAC, TTC, and comfort values remain unchanged.
- The largest APR increment occurs in Round 2, indicating that teacher reconstruction differs from repeatedly fitting a fixed teacher set.Later audits expose improvements that were previously too distant or insufficiently advantageous.
- 26.4% of teacher candidates are rejected after interpolation and exact re-evaluation.Endpoint satisfaction alone does not ensure that an interpolated target is feasible and improving.
- AMPT raises recovery from 55.8% to 66.9% on the 658-scene hard subset, recovering 440 positive outcomes.This is an absolute gain of 11.1 percentage points over scalar GRPO.
- AMPT achieves 73 net repairs after accounting for 20 regressions, with 78.5% of gross repairs retained.DAC-only failures contribute 45 net repairs and collision-only scenes contribute 26.