MRIC: Model-Based Reinforcement-Imitation Learning With Mixture-of-Codebooks for Autonomous Driving Simulation

Accurately simulating the diverse behaviors of heterogeneous agents in varied scenarios is fundamental to autonomous driving simulation. Our first insight is to leverage state-matching through differentiable simulation to provide informative learning signals and enable efficient credit assignment for the policy. We demonstrate this by identifying gradient highways and inter-agent gradient pathways. However, challenges such as gradient explosion and weak supervision in low-density regions emerge. Our second insight addresses these issues by applying dual policy regularizations to constrain the function space. To further enhance diversity, our third insight involves compressing the behaviors of heterogeneous agents into a series of prototype vectors for retrieval. These insights culminate in our model-based reinforcement-imitation learning framework with a temporally abstracted mixture-of-codebooks (MRIC). MRIC integrates open-loop model-based imitation learning to stabilize training and model-based reinforcement learning (RL) to inject domain knowledge. The RL component introduces differentiable rewards for collision avoidance, road adherence, and traffic rule compliance. We further propose a dynamic multiplier mechanism to maintain the effectiveness of the regularizations while avoiding interference. Experiments on the large-scale Waymo Open Motion dataset show that MRIC significantly outperforms strong baselines in diversity, behavioral realism, and distributional fidelity, achieving notable improvements in metrics such as collision rate, minSADE, and time-to-collision JSD.