WAM-Flow: Parallel Coarse-to-Fine Motion Planning via Discrete Flow Matching for Autonomous Driving

AI summary

Overview

Research area: Autonomous driving / robotics — specifically end-to-end motion planning with vision-language-action (VLA) models, drawing on discrete flow matching and reinforcement learning.

Technical level: Advanced. The paper assumes familiarity with generative modeling (diffusion, flow matching, continuous-time Markov chains), autoregressive vision-language models, GRPO-style policy optimization, and closed-loop driving benchmarks such as NAVSIM.

One-sentence scope: The paper proposes WAM-Flow, a 1.5B-parameter VLA planner that reframes ego-trajectory prediction as discrete flow matching over a numerically structured token space, and evaluates it on NAVSIM v1, NAVSIM v2, and nuScenes.

What This Paper Is About

End-to-end autonomous driving models must simultaneously reason about scenes, produce precise continuous control, and perform well in closed-loop simulation — but existing VLA approaches generally split into autoregressive decoders (accurate and interpretable, but slow and sequential) and diffusion planners (parallel, but harder to reason through and lighter on interpretability). The paper asks whether discrete flow matching (DFM), a continuous-time Markov chain over discrete tokens that supports fully parallel, bidirectional denoising, can serve as a better policy representation for driving. The goal is a single model that can produce a coarse trajectory in one denoising step and refine it with more steps, giving a tunable speed-versus-quality trade-off.

Key Contributions

1. Discrete flow matching for VLA motion planning. WAM-Flow casts ego-trajectory planning as discrete flow matching over a structured token space, enabling fully parallel and bidirectional denoising instead of sequential autoregressive decoding. The authors state this is a largely unexplored direction for VLA policies in end-to-end driving.

2. Metric-aligned numerical tokenizer. Continuous scalars (position, heading, velocity, acceleration) are discretized into a uniform codebook over [-100, 100] at 0.01 resolution (N = 20,001). A triplet-margin ranking loss forces Euclidean embedding distances to be monotonic in the underlying scalar differences, so latent geometry reflects numerical proximity.

3. Geometry-aware flow objective. The probability path is defined as a Gibbs distribution induced by a distance metric, with coordinate-wise dissimilarities tailored to data type (tokenizer-induced distances for numbers, circular metrics for angles, semantic distances for text). This replaces the standard masked/mixture path.

4. Simulator-guided GRPO alignment. A Group Relative Policy Optimization objective with a composite reward decomposes NAVSIM PDMS into multiplicative safety penalties (No-Collision, Drivable Area Compliance) and a weighted average of performance objectives (Ego Progress, Time-to-Collision, Comfort). The authors describe this as the first exploration of GRPO within discrete flow matching for autonomous driving VLA, and note it goes beyond conventional likelihood-based training by explicitly incorporating safety alignment.

Main Findings

• NAVSIM v1 state of the art on PDMS. The full model reaches 90.3 PDMS on NAVSIM v1, with sub-metrics NC 99.2, DAC 98.3, TTC 97.0, Comfort 99.7, and EP 82.3. The paper notes this is achieved using only a single front-view camera, outperforming methods that use multi-view camera setups or LiDAR.

• Coarse-to-fine trade-off is real and monotonic up to 5 steps. PDMS rises from 89.1 at 1 step, to 89.7 at 2 steps, 90.0 at 3 steps, 90.3 at 5 steps, and then 90.2 at 10 steps. Inference time grows roughly linearly with steps (0.09s, 0.19s, 0.29s, 0.48s, 0.94s respectively).

• Latency advantage. With 1-step denoising, WAM-Flow is reported as 4.67× faster than RecogDrive; 5-step processing is stated to match RecogDrive's latency. The 1.5B model is also reported as 3× faster than the Janus autoregressive baseline (Janus-1.5B baseline latency is listed as 0.27s; WAM-Flow at 1 step is 0.09s).

• NAVSIM v2 results. The full model achieves 84.7 EPDMS, leading on NC (98.5), DDC (99.5), and LK (97.4) among compared methods. Other reported sub-scores include DAC 94.5, TLC 99.8, EP 86.9, TTC 96.8, HC 97.6, and EC 73.9.

• Component ablation on NAVSIM v1 (Table 5). Starting from the Janus-1.5B text tokenizer for numbers gives 76.2 PDMS. Replacing it with a dedicated numerical tokenizer raises this to 81.1 (+4.9). Adding metric-aligned embeddings reaches 83.4 (+2.3). Adding large-scale VQA pretraining reaches 86.7 (+3.3). Adding simulator-guided GRPO reaches 90.3. The paper separately notes that comparing rows with and without VQA pretraining gives a +3.4 PDMS improvement.

• GRPO group size matters (Table 2). Without GRPO the PDMS is 86.7; group size 2 gives 89.2, group size 3 gives 90.3, and group size 4 gives 89.6. Group size 3 is reported as the best balance between exploration diversity and training stability.

• Reward weighting matters (Table 3). The default EP:TTC:C = 5:5:2 yields the best PDMS of 90.3. Skewing toward safety (5:20:2) yields 89.7, skewing toward progress (20:5:2) yields 90.0, and shifting weight to comfort (5:5:8) yields 90.1.

• Pretraining helps, with diminishing returns after 3 epochs. PDMS peaks at 3 pretraining epochs on 6.5M VQA data, an improvement of +3.3 versus 0 epochs. On data scaling, 0.65M pretraining data yields +1.9 PDMS for the numerical tokenizer and +5.2 for the text tokenizer; further scaling to 6.5M adds +1.4 (numerical) and +2.6 (text) over the 0.65M setting.

• nuScenes collision rates. Under ST-P3 metrics, WAM-Flow reports an average collision rate of 0.12% (1s: 0.04%, 2s: 0.10%, 3s: 0.23%). Under UniAD metrics, it reports the lowest average collision rate among evaluated VLA methods at 0.23% (1s: 0.00%, 2s: 0.10%, 3s: 0.60%).

Methodology in Plain English

The researchers start from an existing general-purpose multimodal backbone (Janus-1.5B) that was built for autoregressive text generation, and convert it into a non-causal flow model. The input is a front-view image, a natural-language navigation command plus system prompt, and the current ego-vehicle state (position, heading, velocity, acceleration). The output is 8 waypoints covering the next 4 seconds.

Two things have to be fixed before this works. First, ordinary text tokens are a poor way to represent numbers like positions and velocities, so the authors build a numerical codebook over [-100, 100] at 0.01 resolution (20,001 entries) and train the embeddings with a triplet-margin ranking loss so that numerically close values end up geometrically close in embedding space. Second, the model needs domain knowledge, so they run a four-stage curriculum:

1. Train the new numerical embeddings and the language head on 668K nuPlan samples for 4 epochs, using the flow-matching cross-entropy loss plus the triplet-margin ranking loss.
2. Continue multimodal pretraining for 3 epochs on 6.5M VQA items — 3.4M general multimodal VQA from LLaVA-v1.5 and 3.1M driving-specific VQA from RecogDrive — to build road-scene competence.
3. Supervised fine-tune the full VLA backbone for 2 epochs on the 668K nuPlan data.
4. Run simulator-guided GRPO for 0.5 epoch on 103K NAVSIM samples, sampling candidate trajectories in parallel and rewarding them with a composite safety-plus-performance signal (EP:TTC:Comfort weighted 5:5:2).

At inference, the model discretizes the time interval [0, 1] into n steps, samples a starting token per coordinate, predicts a target token from the posterior, and applies a jump rule based on the outgoing transition rate against a uniform random draw. Because every coordinate can take a jump independently, generation is parallel rather than left-to-right.

Why This Matters

Impact on research. The paper positions discrete flow matching as a third paradigm alongside autoregressive VLA decoders and continuous diffusion planners. If the reported gains hold, it suggests that parallel, bidirectional decoding can be combined with a safety-aligned reinforcement learning stage without losing the parallel structure — a combination that is not straightforward with purely autoregressive policies. The metric-aligned tokenizer also offers a reusable idea anywhere continuous values must be embedded into a discrete token vocabulary.

Real-world applications:

• Real-time urban driving. The 1-step setting (89.1 PDMS at 0.09s) is aimed at scenarios with tight latency budgets, where an approximate plan is good enough.
• Complex intersection and interaction handling. The 5-step setting (90.3 PDMS at 0.48s) allows more compute when the scene demands it.
• Efficient fleet deployment. A 1.5B-parameter backbone that runs 4.67× faster than RecogDrive at 1 step lowers hardware requirements for edge inference.
• Safety-critical scenario collection and simulation. The reward decomposition into No-Collision, Drivable Area Compliance, Ego Progress, Time-to-Collision, and Comfort could inform how training data is curated and how simulators are used for policy alignment.

Industry relevance. The training used 4 × 8 Ascend 910B NPUs across four sequential stages, and the model is derived from an openly described backbone (Janus-1.5B). The efficiency table places WAM-Flow at 0.09s to 0.48s against FSDrive (10.58s), Epona (1.24s), and ReCogDrive (0.42s), a range that matters for automotive compute budgets.

Future Directions

• Closing the NAVSIM v2 gap. EPDMS is 84.7 while v1 PDMS is 90.3; the paper does not report its own EC (Extended Comfort) score above 73.9, and the abstract-level claims focus on v1. Extending the gains to the nine-metric v2 suite is an obvious next step.
• Scaling the pretraining data further. The data-scaling results show diminishing but still positive returns (0.65M → +1.9/+5.2; 6.5M → additional +1.4/+2.6), leaving open how far the curve continues.
• Compute-adaptive step selection. The paper presents a tunable step count but does not report a policy that automatically chooses the number of denoising steps per scene. Learning when to stop refining is a natural extension.
• Generalization beyond the evaluated benchmarks. Evaluation is on NAVSIM v1 (codebase v1.1), NAVSIM v2 (codebase v2.2), and nuScenes; the paper does not report cross-city or cross-sensor transfer results. The authors also state in the appendix that limitations and future work are discussed there, though the truncated content does not include that section.

Target Audience

This paper is most valuable to autonomous driving researchers working on end-to-end planning and VLA architectures, particularly those already familiar with NAVSIM benchmarks and diffusion- or autoregressive-based planners. It will also interest generative modeling researchers exploring discrete flow matching and non-autoregressive generation outside of language, and reinforcement learning practitioners studying how GRPO transfers from reasoning tasks to continuous-control-adjacent planning problems. Readers seeking an introductory treatment of autonomous driving will find the mathematical formulation (continuous-time Markov chains, probability velocity, Kolmogorov forward equation) demanding.

Authors’ abstract

We introduce WAM-Flow, a vision-language-action (VLA) model that casts ego-trajectory planning as discrete flow matching over a structured token space. In contrast to autoregressive decoders, WAM-Flow performs fully parallel, bidirectional denoising, enabling coarse-to-fine refinement with a tunable compute-accuracy trade-off. Specifically, the approach combines a metric-aligned numerical tokenizer that preserves scalar geometry via triplet-margin learning, a geometry-aware flow objective and a simulator-guided GRPO alignment that integrates safety, ego progress, and comfort rewards while retaining parallel generation. A multi-stage adaptation converts a pre-trained auto-regressive backbone (Janus-1.5B) from causal decoding to non-causal flow model and strengthens road-scene competence through continued multimodal pretraining. Thanks to the inherent nature of consistency model training and parallel decoding inference, WAM-Flow achieves superior closed-loop performance against autoregressive and diffusion-based VLA baselines, with 1-step inference attaining 89.1 PDMS and 5-step inference reaching 90.3 PDMS on NAVSIM v1 benchmark. These results establish discrete flow matching as a new promising paradigm for end-to-end autonomous driving. The code will be publicly available soon.