AD-R1: Closed-Loop Reinforcement Learning for End-to-End Autonomous Driving with Impartial World Models

• Optimistic bias makes world models unreliable critics because they may replace commanded collisions or off-road actions with implausibly safe futures.The paper identifies this as a breakdown in causal fidelity.Source: §1 · p. 2
• AD-R1 uses an impartial world model as an internal critic that simulates planned trajectories and returns safety-critical feedback for offline policy learning.The framework is intended to enhance the safety of driving agents.Source: §1 · p. 2
• Counterfactual Synthesis generates physically consistent, causally faithful trajectories leading to plausible failures, while model refinements keep predictions faithful to commanded trajectories.The refinements include Trajectory-Aware Gating and Ego-Trajectory Fidelity Loss.Source: §1 · p. 2
• The critic derives multi-faceted rewards from predicted 4D occupancy, penalizing collisions, off-road excursions, unsafe buffers, and vertical-clearance violations.Source: §1 · p. 2
• AD-R1 improves planning PDSM by 1.7% absolute on navsim across two baselines.Source: §1 · p. 2

2.1. World Models for Autonomous Driving

• World models have succeeded as high-fidelity, controllable driving-video generators [11] [12] [26] [38] [45] [50], while action-conditioned simulation spans future-video prediction [13] [17] [49] to explicit 3D/4D representations [19] [46].Source: §2.1 · pp. 2–3
• The authors argue that training on safe, on-distribution scenarios can leave world models biased toward safe forecasts, making them unreliable for safety-critical policy refinement.Their approach trains on synthesized failures to make the model an impartial judge of safety.Source: §2.1 · p. 3

2.2. End-to-end Autonomous Driving

• End-to-end driving policies commonly use imitation learning, which relies on expert-data distributions and can struggle with long-tail scenarios under distributional shift [3] [5] [16] [18] [25] [27].Source: §2.2 · p. 3
• The paper uses a learned world model to generate diverse future outcomes conditioned on agent actions, enabling policy learning from imagined mistakes.It presents the impartial world model as a reliable generative simulator and internal critic.Source: §2.2 · p. 3

3.1. Overview and Problem Formulation

• The world model predicts future 3D occupancy from historical observations and candidate trajectories, and the framework computes rewards from those predicted futures.This learned simulator replaces reliance on external simulators for obtaining trajectory reward signals.Source: §3.1 · p. 3

3.2.1. Backbone Architecture.

• The I²-World-inspired backbone compresses historical occupancy into latent tokens, then autoregressively forecasts future tokens for decoding into occupancy grids.Its two-stage architecture underpins the model’s failure-imagining refinements.Source: §3.2.1 · p. 4

3.2.2. Counterfactual Data Synthesis.

• Real-world datasets mainly contain safe professional trajectories, so Counterfactual Synthesis creates a structured curriculum of negative examples to address missing failure data.A unified kinematic model generates off-road, static-obstacle penetration, and dynamic-agent collision trajectories; transformed occupancy grids preserve scene geometry.Source: §3.2.2 · pp. 4–5

3.2.3. Model-Centric Refinements

• Trajectory-Aware Gating modulates scene features along the commanded path, while Ego-Trajectory Fidelity Loss penalizes predicted ego motion that deviates from that path.The fidelity loss discourages the model from predicting that the ego-vehicle autonomously swerves to avoid a collision.Source: §3.2.3 · p. 5

3.3.1. Training Paradigm

• In offline refinement, GRPO samples trajectory groups, scores their imagined outcomes, and updates the policy using relative advantages while behavior cloning anchors it to the reference policy.The weighted objective combines RL and behavior-cloning losses to regularize policy updates.Source: §3.3.1 · pp. 5–6

3.3.2. 4D Reward Modeling

• The 4D reward function converts predicted occupancy into a scalar signal with components for safety, comfort, and task progress.Its 3D and temporal predictions support reward detail beyond heuristic or 2D BEV approaches.Source: §3.3.2 · p. 6

3.3.3. Safety-Critical Penalties

• Safety penalties target volumetric collisions and ego-vehicle occupancy of non-drivable areas, with collision penalties scaling according to overlap severity.The collision penalty includes a contact penalty and an overlap-severity term.Source: §3.3.3 · p. 6

3.3.4. 3D/4D-Aware Comfort & Rule-Adherence

• 3D/4D-aware rewards penalize overhead-structure clearance risks and unstable drivable surfaces, while task rewards encourage route progress and target-speed maintenance.The paper identifies vertical-space constraints as inaccessible to 2D BEV representations.Source: §3.3.4 · pp. 6–7, §3.3.5 · p. 7

4.1. Experimental Setup

• The setup trains the impartial world model on 80% real nuScenes data and 20% synthesized counterfactual data, then evaluates the policy on NavSim.Evaluation covers both world-model fidelity and refined-policy performance.Source: §4.1 · p. 7

4.2. Quantitative Analysis

• 1.7% and 1.1% are the respective absolute PDSM score boosts reported for DiffusionDrive and Recogdrive after refinement.The paper reports stable PDMS improvements for both agents in closed-loop simulation.Source: §4.2 · p. 7, Table 2 · p. 8

4.3. Ablation Studies

• 89.8 PDMS is the best ablation result, reached by adding physically grounded 4D rewards after counterfactual synthesis and model-centric refinements.Counterfactual Synthesis raises PDMS to 89.1, and model-centric refinements raise it to 89.3.Source: §4.3 · pp. 7–8, Table 3 · p. 8

4.4. Qualitative Analysis

• Figure 3 contrasts an original agent plan ending in collision or off-road with a refined agent that safely avoids the hazard.The figure presents behavior with and without AD-R1 refinement.Source: Figure 3 · p. 8
• Given an unsafe off-road trajectory, the standard model hallucinates a safe future while IWM predicts the off-road scene and a penalty signal.The comparison uses the same unsafe trajectory as input to both models.Source: §4.4 · p. 8, Figure 4 · p. 9

• The conclusion frames world-model optimistic bias—the inability to predict failure—as a fundamental flaw for reinforcement learning in driving.Source: §5 · p. 9
• AD-R1 uses Counterfactual Synthesis to train an Impartial World Model that predicts consequences of unsafe actions and serves as an internal critic.Source: §5 · p. 9
• The paper reports that this critic enables agents to learn from imagined mistakes and substantially reduces safety violations in closed-loop simulations.Source: §5 · p. 9

• The title names closed-loop reinforcement learning and impartial world models as the paper’s subject.Source: Supplementary Material · p. 14

• Section A provides the Impartial World Model architecture, hyperparameters, and RL training configuration.Source: Overview · p. 14
• Section B elaborates on the nuScenes-CF dataset and the mathematical formulation of Counterfactual Synthesis.Source: Overview · p. 14
• Section C gives strict mathematical definitions for the Risk Foreseeing Benchmark.Source: Overview · p. 14
• Section D reports sensitivity analysis of synthetic-data ratios.Source: Overview · p. 14
• Section E contains additional qualitative analysis and a description of the attached demo.Source: Overview · p. 14

A.1. Impartial World Model (IWM) Architecture

• The IWM uses a VQ-VAE-based tokenizer and autoregressive Transformer forecaster, conditioning on 2 seconds of history to predict 3 seconds of 4D occupancy.The tokenizer uses a 200 × 200 × 16 voxel grid; the forecaster has 24 attention layers and 16 heads.Source: §A.1 · p. 14

A.2. RL Training Configuration (AD-R1)

• GRPO refines the denoising heads of DiffusionDrive and ReCogDrive, with model-specific group sizes, learning rates, and training duration.Group sizes are 64 and 8, respectively; refinement takes approximately 24 hours for 10 epochs.Source: §A.2 · p. 14

A.3. Reward Function Details

• The reward shaping is critical for guiding the agent towards safety without compromising progress, and Table 4 details the specific weights used in Eq. (18) of the main paper.Source: §A.3 · p. 14

B.1. Kinematic Trajectory Generation

• Unsafe trajectories use a blending kinematic model, with a quadratic time-varying coefficient producing a smooth deviation toward the target.Figure 5 marks trajectory starts in green and ends in red for off-road examples.Source: §B.1 · pp. 14–15, Figure 5 · p. 15

B.2. Dataset Statistics and Sampling Strategy

• Training uses 80% original nuScenes samples and 20% synthetic counterfactual clips, rather than an 80:20 split of the datasets’ absolute sizes.The counterfactual pool contains 1,800 clips derived from nuScenes training data; evaluation uses a separate 450-clip RFB set.Source: §B.2 · p. 15
• The synthetic pool balances dynamic collisions at 40%, static collisions at 30%, and off-road or sidewalk excursions at 30%.Dynamic collisions include vehicle-vehicle and vehicle-pedestrian events; static collisions include barriers, cones, and buildings.Source: §B.2 · p. 15

B.3. Counterfactual Example

• The paper visualizes synthetic counterfactual occupancy examples for penetration, off-road, and collision events.Figures 5–7 show examples of the respective unsafe trajectory categories.Source: §B.3 · p. 15

C.1. Risk Foreseeing Benchmark (RFB)

• The RFB comprises 450 challenging clips with ground-truth unsafe ego-trajectories spanning 3 seconds, testing negative-outcome prediction rather than safe expert driving.It is designed to quantify the ability to foresee failures.Source: §C.1 · pp. 15–16
• RFB metrics assess whole-scene fidelity with G-IoU, occupancy in critical failure regions with f-IoU, and dynamic-agent preservation with DAF.A high f-IoU indicates that occupied space in the ego-vehicle’s path is predicted rather than hallucinated as empty.Source: §C.1 · pp. 16–17

C.2. Driving Policy Evaluation

• On NavSim, the final agent is evaluated with safety and planning metrics, including PDSM for progress, comfort, and safety.PDSM is the stated holistic planning-quality score.Source: §C.2 · p. 17

C.3. Baselines and Comparison Schemes

• World-model comparisons use I2-World and DOME, while policy refinement is tested on DiffusionDrive and ReCogDrive.The world-model baselines are trained only on safe data; the two driving agents are pre-trained end-to-end systems.Source: §C.3 · p. 17

D.1. Ablation: Synthetic Data Ratio

• The ablation varies the proportion of synthetic counterfactual data mixed into training and reports that a balance is crucial.The reported optimal balance is the 80:20 mix.Source: §D.1 · p. 17
• At 50% synthetic data, f-IoU improves slightly while G-IoU degrades because the synthetic physics are simplified approximations.The ablation describes a safety-awareness versus general-scene-generation trade-off.Source: §D.1 · p. 17
• The 80:20 synthetic-to-real training mix offers the best trade-off between scene fidelity (G-IoU) and safety awareness (f-IoU).Table 5 summarizes the effect of synthetic data ratio.Source: §D.1 · p. 17, Table 5 · p. 18

D.2. Ablation: Reward Function Components

• Removing safety penalties returns Collision Rate to the 1.8% baseline, a 0.2% absolute difference described as meaningful.The text notes this improvement is meaningful given that the baseline policy is already highly optimized.Source: §D.2 · p. 18
• Removing the 3D stability reward yields a PDMS of 87.6, below the original baseline of 88.1.The text attributes the lower ride quality to jerky trajectories when comfort constraints are absent.Source: §D.2 · p. 18
• Removing task-oriented rewards drops PDMS to 76.8 and slightly raises collision rate to 2.0%.The text suggests that lacking goal-directed behavior may make the agent unpredictable or less efficient at exiting dangerous zones.Source: §D.2 · p. 18

E. Qualitative Visualization

• In the sudden cut-in scenario, the baseline plan collides, while the refined agent predicts the risk and brakes safely.Figure 8 contrasts the original and refined agents; its caption also describes baseline collision or off-road outcomes and safe hazard avoidance after refinement.Source: E. Qualitative Visualization · p. 18, Figure 8 · p. 19