# Closed-Loop Supervised Fine-Tuning of Tokenized Traffic Models

arXiv: https://arxiv.org/abs/2412.05334 · HTML: https://mlsanity.com/paper/2412.05334/

## Key facts

- Cited by 40 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 49 references that are in the map.
- First posted on arXiv 5 Dec 2024.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Large Scale Interactive Motion Forecasting for Autonomous Driving : The Waymo Open Motion Dataset](https://mlsanity.com/paper/2104.10133/) (2021)
- [NuPlan: A closed-loop ML-based planning benchmark for autonomous vehicles](https://mlsanity.com/paper/2106.11810/) (2021)
- [ChauffeurNet: Learning to Drive by Imitating the Best and Synthesizing the Worst](https://mlsanity.com/paper/1812.03079/) (2018)
- [Motion Transformer with Global Intention Localization and Local Movement Refinement](https://mlsanity.com/paper/2209.13508/) (2022)
- [MotionDiffuser: Controllable Multi-Agent Motion Prediction using Diffusion](https://mlsanity.com/paper/2306.03083/) (2023)
- [End-to-End Urban Driving by Imitating a Reinforcement Learning Coach](https://mlsanity.com/paper/2108.08265/) (2021)
- [TrafficSim: Learning to Simulate Realistic Multi-Agent Behaviors](https://mlsanity.com/paper/2101.06557/) (2021)
- [Guided Conditional Diffusion for Controllable Traffic Simulation](https://mlsanity.com/paper/2210.17366/) (2022)
- [Scene Transformer: A unified architecture for predicting multiple agent trajectories](https://mlsanity.com/paper/2106.08417/) (2021)
- [MotionLM: Multi-Agent Motion Forecasting as Language Modeling](https://mlsanity.com/paper/2309.16534/) (2023)
- [MTR++: Multi-Agent Motion Prediction with Symmetric Scene Modeling and Guided Intention Querying](https://mlsanity.com/paper/2306.17770/) (2023)
- [Latent Variable Sequential Set Transformers For Joint Multi-Agent Motion Prediction](https://mlsanity.com/paper/2104.00563/) (2021)
- [Imitation Is Not Enough: Robustifying Imitation with Reinforcement Learning for Challenging Driving Scenarios](https://mlsanity.com/paper/2212.11419/) (2022)
- [Language-Guided Traffic Simulation via Scene-Level Diffusion](https://mlsanity.com/paper/2306.06344/) (2023)
- [Overcoming Limitations of Mixture Density Networks: A Sampling and Fitting Framework for Multimodal Future Prediction](https://mlsanity.com/paper/1906.03631/) (2019)
- [Query-Efficient Imitation Learning for End-to-End Autonomous Driving](https://mlsanity.com/paper/1605.06450/) (2016)
- [BITS: Bi-level Imitation for Traffic Simulation](https://mlsanity.com/paper/2208.12403/) (2022)
- [Symphony: Learning Realistic and Diverse Agents for Autonomous Driving Simulation](https://mlsanity.com/paper/2205.03195/) (2022)
- [Trajeglish: Traffic Modeling as Next-Token Prediction](https://mlsanity.com/paper/2312.04535/) (2023)
- [TrafficBots: Towards World Models for Autonomous Driving Simulation and Motion Prediction](https://mlsanity.com/paper/2303.04116/) (2023)
- [BehaviorGPT: Smart Agent Simulation for Autonomous Driving with Next-Patch Prediction](https://mlsanity.com/paper/2405.17372/) (2024)
- [Scenario Diffusion: Controllable Driving Scenario Generation With Diffusion](https://mlsanity.com/paper/2311.02738/) (2023)
- [Gen-Drive: Enhancing Diffusion Generative Driving Policies with Reward Modeling and Reinforcement Learning Fine-tuning](https://mlsanity.com/paper/2410.05582/) (2024)
- [Solving Motion Planning Tasks with a Scalable Generative Model](https://mlsanity.com/paper/2407.02797/) (2024)
- [Real-Time Motion Prediction via Heterogeneous Polyline Transformer with Relative Pose Encoding](https://mlsanity.com/paper/2310.12970/) (2023)

## Built on by (most cited first, 30 of 40)

- [AutoVLA: A Vision-Language-Action Model for End-to-End Autonomous Driving with Adaptive Reasoning and Reinforcement Fine-Tuning](https://mlsanity.com/paper/2506.13757/) (2025)
- [Pseudo-Simulation for Autonomous Driving](https://mlsanity.com/paper/2506.04218/) (2025)
- [Building reliable sim driving agents by scaling self-play](https://mlsanity.com/paper/2502.14706/) (2025)
- [Advancing Multi-agent Traffic Simulation via R1-Style Reinforcement Fine-Tuning](https://mlsanity.com/paper/2509.23993/) (2025)
- [LANGTRAJ: Diffusion Model and Dataset for Language-Conditioned Trajectory Simulation](https://mlsanity.com/paper/2504.11521/) (2025)
- [DRoPE: Directional Rotary Position Embedding for Efficient Agent Interaction Modeling](https://mlsanity.com/paper/2503.15029/) (2025)
- [TrajTok: Technical Report for 2025 Waymo Open Sim Agents Challenge](https://mlsanity.com/paper/2506.21618/) (2025)
- [SPACeR: Self-Play Anchoring with Centralized Reference Models](https://mlsanity.com/paper/2510.18060/) (2025)
- [Learning to Drive in New Cities Without Human Demonstrations](https://mlsanity.com/paper/2602.15891/) (2026)
- [TeraSim-World: Worldwide Safety-Critical Data Synthesis for End-to-End Autonomous Driving](https://mlsanity.com/paper/2509.13164/) (2025)
- [RoaD: Rollouts as Demonstrations for Closed-Loop Supervised Fine-Tuning of Autonomous Driving Policies](https://mlsanity.com/paper/2512.01993/) (2025)
- [DecompGAIL: Learning Realistic Traffic Behaviors with Decomposed Multi-Agent Generative Adversarial Imitation Learning](https://mlsanity.com/paper/2510.06913/) (2025)
- [Surprise Potential as a Measure of Interactivity in Driving Scenarios](https://mlsanity.com/paper/2502.05677/) (2025)
- [Learning Through Retrospection: Improving Trajectory Prediction for Automated Driving with Error Feedback](https://mlsanity.com/paper/2504.13785/) (2025)
- [RIFT: Group-Relative RL Fine-Tuning for Realistic and Controllable Traffic Simulation](https://mlsanity.com/paper/2505.03344/) (2025)
- [Do LLM Modules Generalize? A Study on Motion Generation for Autonomous Driving](https://mlsanity.com/paper/2509.02754/) (2025)
- [nuPlan-R: A Closed-Loop Planning Benchmark for Autonomous Driving via Reactive Multi-Agent Simulation](https://mlsanity.com/paper/2511.10403/) (2025)
- [Human-like autonomy emerges from self-play and a pinch of human data](https://mlsanity.com/paper/2606.19370/) (2026)
- [Scaling Self-Play for End-to-End Driving](https://mlsanity.com/paper/2606.19641/) (2026)
- [Autoregressive Meta-Actions for Unified Controllable Trajectory Generation](https://mlsanity.com/paper/2505.23612/) (2025)
- [Long-term Traffic Simulation with Interleaved Autoregressive Motion and Scenario Generation](https://mlsanity.com/paper/2506.17213/) (2025)
- [SceneStreamer: Continuous Scenario Generation as Next Token Group Prediction](https://mlsanity.com/paper/2506.23316/) (2025)
- [Post-Training and Test-Time Scaling of Generative Agent Behavior Models for Interactive Autonomous Driving](https://mlsanity.com/paper/2512.13262/) (2025)
- [AutoWorld: Learning Multi-Agent Traffic Simulation with Self-Supervised World Models](https://mlsanity.com/paper/2603.28963/) (2026)
- [SCORP: Scene-Consistent Multi-agent Diffusion Planning with Stable Online Reinforcement Post-Training for Cooperative Driving](https://mlsanity.com/paper/2604.11734/) (2026)
- [RLFTSim: Realistic and Controllable Multi-Agent Traffic Simulation via Reinforcement Learning Fine-Tuning](https://mlsanity.com/paper/2605.19033/) (2026)
- [Learning Rollout from Sampling:An R1-Style Tokenized Traffic Simulation Model](https://mlsanity.com/paper/2603.24989/) (2026)
- [Efficient Equivariant Transformer for Self-Driving Agent Modeling](https://mlsanity.com/paper/2604.01466/) (2026)
- [Learning from Mistakes: Rollout-Retrieval Lifelong Policy Learning for Autonomous Driving](https://mlsanity.com/paper/2606.30537/) (2026)
- [World Models as Adversaries: Multi-Agent Self-Play Fine-Tuning for Robust Motion Planning](https://mlsanity.com/paper/2607.10630/) (2026)

## Abstract (arXiv)

Traffic simulation aims to learn a policy for traffic agents that, when unrolled in closed-loop, faithfully recovers the joint distribution of trajectories observed in the real world. Inspired by large language models, tokenized multi-agent policies have recently become the state-of-the-art in traffic simulation. However, they are typically trained through open-loop behavior cloning, and thus suffer from covariate shift when executed in closed-loop during simulation. In this work, we present Closest Among Top-K (CAT-K) rollouts, a simple yet effective closed-loop fine-tuning strategy to mitigate covariate shift. CAT-K fine-tuning only requires existing trajectory data, without reinforcement learning or generative adversarial imitation. Concretely, CAT-K fine-tuning enables a small 7M-parameter tokenized traffic simulation policy to outperform a 102M-parameter model from the same model family, achieving the top spot on the Waymo Sim Agent Challenge leaderboard at the time of submission. The code is available at https://github.com/NVlabs/catk.