Research map / Multi-agent systems and games / Human-AI and goal-conditioned
Closed-Loop Supervised Fine-Tuning of Tokenized Traffic Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes CAT-K rollouts, a closed-loop supervised fine-tuning method that picks the closest-to-ground-truth action among the policy's top-K tokens, letting a 7M tokenized traffic model (SMART-tiny) beat the 102M SMART-large on the WOSAC leaderboard.
Problem
Tokenized next-token-prediction traffic simulation policies are trained open-loop with behavior cloning and suffer covariate shift when rolled out in closed loop, degrading realism.
Approach
During fine-tuning, roll out the policy by choosing, at each step and per agent, the token among the policy's top-K most likely tokens that lands closest to the ground-truth next state, then apply DAD-style recovery targets computed from GT; this stays on-policy while keeping GT a valid supervision target.
Key results
- SMART-tiny fine-tuned with CAT-K reaches 0.7702 RMM on the WOSAC test leaderboard, above SMART-large (102M) at 0.7614 and top of the leaderboard at submission.
- On the local validation split, CAT-32 achieves 0.7616 RMM versus 0.7581 for BC pre-training, while top-K rollouts without the closest-selection degrade performance (e.g., Top-5 + distance filter 0.6860).
- On ego-motion planning with a GMM policy, CAT-3 fine-tuning reduces collision rate from 0.0568 to 0.0422 and off-road rate from 0.0053 to 0.0035 relative to BC pre-training (the paper's headline reductions of 25.7% and 33.9% are for the planning task).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fine-tune SMART-tiny with CAT-K | SMART-tiny · WOSAC 2024 leaderboard test split | RMM: 0.7591 -> 0.7702 | Helped |
| Fine-tune SMART-tiny with CAT-K instead of scaling to SMART-large | SMART-tiny · WOSAC 2024 leaderboard test split | RMM: 0.7614 -> 0.7702 | Helped |
| Continue fine-tuning with BC (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7581 -> 0.7590 | No clear effect |
| Replace CAT-K rollouts with Trajeglish noisy tokenization (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7581 -> 0.7562 | Hurt |
| Replace CAT-K rollouts with SMART trajectory perturbation (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7581 -> 0.7556 | No clear effect |
| Use policy top-K rollouts instead of CAT-K rollouts (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7581 -> 0.6478 | Hurt |
| Filter top-5 rollouts by distance to the GT (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7581 -> 0.6860 | Hurt |
| Sample top-5 rollouts by distance to the GT (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7581 -> 0.7058 | Hurt |
| Use deterministic maximum-probability rollouts instead of CAT-K (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7581 -> 0.6361 | Hurt |
| Increase CAT-K K from 5 to 16 (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7423 -> 0.7604 | Helped |
| Increase CAT-K K from 16 to 32 (ablation) | BC-pretrained SMART-tiny · WOSAC 2% validation split | RMM: 0.7604 -> 0.7616 | Helped |
| Apply CAT-K fine-tuning with K=3 (ablation) | GMM ego-policy · WOSAC 2% validation split, ego-motion planning | Collision rate: 0.0568 -> 0.0422 | Helped |
Surprising findings
- Fine-tuning with top-K sampling or trajectory noising (Trajeglish, SMART perturbation) fails to improve RMM and often hurts it, while CAT-K is the only method that significantly helps.
- Additional BC fine-tuning of the pre-trained model does not help, so gains come from the closed-loop rollouts rather than extra training.
Limitations
- Requires sampling rollouts from the policy during training, which raises data generation cost (mitigated by BC pre-training).
- Winner-takes-all selection of the closest token is greedy and only approximates the minADE-optimal rollout.
- Evaluated only on WOSAC/WOMD with the NTP SMART architecture and a GMM ego-policy; not tested on diffusion or VAE policies.
Experimental setup
Datasets: Waymo Open Motion Dataset (WOMD)
Benchmarks: WOSAC 2024 leaderboard (test split) WOSAC local validation split (2%)
Metrics: Realism Meta-Metric (RMM) kinematic metrics interactive metrics map-based metrics minADE collision rate off-road rate
Baselines: SMART-large SMART-tiny KiGRAS FDriver-tiny BehaviorGPT GUMP Trajeglish noisy tokenization SMART trajectory perturbation
Builds on: SMART Data As Demonstrator (DAD) Trajeglish DAgger
Closest papers by meaning
- Advancing Multi-agent Traffic Simulation via R1-Style Reinforcement Fine-Tuning (ICLR 2026 · similarity 0.89)
- RLFTSim: Realistic and Controllable Multi-Agent Traffic Simulation via Reinforcement Learning Fine-Tuning (CVPR 2026 · similarity 0.88)
- LASIL: Learner-Aware Supervised Imitation Learning For Long-term Microscopic Traffic Simulation (CVPR 2024 · similarity 0.86)
- Trajeglish: Traffic Modeling as Next-Token Prediction (ICLR 2024 · similarity 0.85)
- Prompt to Transfer: Sim-to-Real Transfer for Traffic Signal Control with Prompt Learning (AAAI 2024 · similarity 0.85)
- Causal Composition Diffusion Model for Closed-loop Traffic Generation (CVPR 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.