Direct Imitation Learning for Autonomous Driving: When the World-Model Latent Loses the Lane

Imitation Learning for Autonomous Driving

ALVINN introduced behavioral cloning for driving from camera images. Ross et al. formalised Dataset Aggregation (DAgger) to address distribution shift by iteratively querying the expert at states visited by the current policy. ChauffeurNet augments demonstrations with synthesised perturbations and auxiliary loss heads to improve recovery behaviour. Our work applies these ideas in a low-compute, simulation-only setting using the Intelligent Driver Model as the oracle.

World Models for Driving

DreamerV3 learns a recurrent world model (RSSM) by jointly optimising a reconstruction decoder, a reward predictor, and a discount predictor, then derives a policy by imagining rollouts in latent space. Our failure case — a world model that compresses away the precise signal the policy needs — is a known risk with RSSM-based representations when the observation space is already low-dimensional and the model is under-trained.

Distribution Shift in Behavioral Cloning

DAgger corrects distribution shift by on-policy data collection; DART injects noise to force the policy into off-distribution states during collection, providing recovery labels without requiring online expert querying. We combine both approaches and compare their effectiveness at modest compute.

MetaDrive

MetaDrive is a lightweight, procedurally generated driving simulator with a modular road block system supporting diverse geometry types (straight, curve, intersection, roundabout). Its Intelligent Driver Model (IDM) provides a rule-based expert oracle that can be queried at arbitrary states, making it ideal for DAgger and relabeling experiments.

Environment and Observation Space

All experiments use MetaDrive's MetaDriveEnv with 50 training maps and a held-out evaluation set. The observation is a 259-dimensional vector comprising ego vehicle state (speed, heading, lateral offset), LiDAR point cloud returns (120 beams), and navigation targets. Actions are 2-dimensional: steering ∈ [−1, 1] and acceleration ∈ [−1, 1]. Episode length is capped at 1,000 steps; an episode terminates on collision, going off-road, or reaching the route endpoint.

Expert Oracle

We use MetaDrive's built-in Intelligent Driver Model (IDM) as the expert oracle throughout. IDM is rule-based (not learned), models car-following and lane-keeping using closed-form kinematic equations, and achieves 99% route completion on held-out maps. We query it via IDMPolicy.act() at each environment step to obtain expert action labels.

World-Model Baseline (RSSM-BC)

We train a DreamerV3-style RSSM with a 32×32 categorical latent, GRU recurrent state of dimension 512, and decoders for observation reconstruction, reward, and discount. Training uses 8,000 IDM demonstration steps over 50,000 gradient steps. We then freeze the world-model encoder and train a two-layer MLP actor on the latent state via behavioral cloning (L1 loss against IDM actions) for 5,000 gradient steps with batch size 64.

Direct Policy (DirectBC)

The direct policy is a three-layer MLP with hidden dimension 256 and SiLU activation. Output is squashed by tanh to [−1, 1]². Trained with L1 loss against IDM action labels, Adam optimiser, learning rate 3×10⁻⁴, batch size 256.

Recovery Data via DART Perturbation

To teach recovery behaviour we inject triangular steering impulses during collection, following DART. At each step with probability p_pert = 0.02, steering is displaced by a triangular waveform with magnitude γ = 0.15 and duration τ ~ U[10, 40] steps. We execute the perturbed action in the environment but store the clean IDM action as the training label, so the policy sees the recovery state but learns the correct centering action. We collect 8,000 clean steps and 8,000 perturbed steps, training on the combined 16,000-step dataset.

Scene-Targeted Augmentation

After identifying roundabouts as the dominant failure geometry via per-scene evaluation, we inject additional demonstrations specifically on roundabout maps (map="O" in MetaDrive). We collect an additional 4,000 clean IDM steps and 4,000 recovery steps on roundabout geometry and mix them into the training dataset, giving a total of 24,000 training steps.

DAgger

At each DAgger iteration we roll out the current policy for 2,000 steps, relabel each state with the IDM action via IDMPolicy.act(), and aggregate into the training buffer. We run 3 iterations, starting from the clean+recovery+boost baseline. We evaluate two DAgger variants: one applied to the direct policy and one applied to the RSSM-conditioned actor as an ablation.

Evaluation Protocol

All evaluation uses deterministic fixed seeds on held-out maps. We report:

• Route completion %: fraction of the route navigated before episode termination, averaged across episodes.
• Off-road rate: fraction of episodes that terminate due to leaving the drivable surface.
• Crash rate: fraction of episodes terminating due to collision.
• Recovery rate: fraction of episodes in which the policy successfully re-enters the lane after being displaced off-center for the first 30 steps.

World Model Latent vs. Direct Observation

Table 1 shows the core reframing result. The RSSM-BC policy achieves route completion identical to a random agent (2%), despite 50,000 gradient steps of world-model training and 5,000 steps of actor training. The bc_loss of 2.06 — compared to 0.05 for DirectBC — indicates that the world-model latent state does not contain sufficient lane/heading information for the actor to fit the IDM actions. The direct MLP achieves 44× lower bc_loss and drives non-trivially (22% route completion) with no world-model component at all.

The 259-dim state observation already exposes heading angle, lateral lane offset, and navigation waypoints explicitly as named features. The RSSM encoder maps this structured signal through a learned compression that, at this training scale, produces a latent that is easier to reconstruct from than to plan actions from — a known failure mode of reconstruction-based world models when the prediction task does not require preserving the action-relevant signal.

Recovery Data Eliminates Crashes

Adding 8,000 recovery steps (Table 2) yields three improvements: route completion 24%→39% (+15pp), crash rate 20%→0%, off-road rate 50%→30% (−40% relative). The policy has learned to steer back toward the lane center after perturbation-induced drift. Crash rate reaches zero because collisions in MetaDrive arise from over-correction that swings the vehicle into traffic, which recovery training specifically addresses.

Doubling the dataset (32,000 steps) raises recovery rate from 40% to 70% but leaves route completion flat at 39%, confirming a capacity ceiling. Adding a shared auxiliary head predicting per-step reward alongside the action worsens off-road rate (30%→40%) and recovery rate (70%→40%), consistent with the auxiliary task competing for representation capacity when the observation already contains the lane signal explicitly.

Per-Scene Breakdown Reveals a Single Failure Geometry

The aggregate route completion of 39% masks a stark geometry-dependent split (Table 3). Three of four geometry types show near-IDM performance: the policy drives straights at 96%, curves at 82%, and intersections at 84%. Roundabouts are categorically different: route 42%, success 0%, crash 40%, off-road 60%. The aggregate 30% off-road rate is almost entirely attributable to roundabouts.

This finding reframes the next improvement target. The aggregate metric suggested a broadly underperforming policy; the per-scene breakdown reveals a policy that is near-competent on most tasks, with one failing geometry type accounting for the residual gap. The 42% roundabout route completion is consistent with the policy successfully navigating the straight approach before failing at the curved entry arc, where the required steering angle is sharp and the state distribution is qualitatively different from anything in the clean training data.

Scene-Targeted Augmentation Raises Roundabout Performance by 22pp

Adding 8,000 roundabout-specific steps (4k clean + 4k recovery) raises roundabout route completion from 42% to 64% (+22pp) and cuts off-road rate from 60% to 30% (Table 4). No other geometry regresses; curves improve by a further 5pp as a beneficial side effect of the additional curved-geometry data.

The aggregate gain is modest (+2pp) because the mixed-map evaluation samples uniformly over geometry types. The per-scene gain is the meaningful signal. Qualitatively, the boosted policy successfully navigates the roundabout entry arc in the majority of episodes, with remaining failures concentrated at tight mid-roundabout turns where the required steering exceeds the training data coverage even with augmentation.

DAgger Matches but Does Not Exceed Augmented BC

Table 5 shows direct-policy DAgger results. Iteration 1 matches the augmented BC baseline (41%). However, iterations 2 and 3 show clear regression: route completion drops to 35% and off-road rate doubles. The bc_loss rise from 0.043 to 0.316 reveals the mechanism: IDM relabels at the policy's failure states are high-variance and structurally different from the clean training distribution. As the proportion of failure-state data grows, the combined dataset becomes harder to fit, and the optimiser trades performance on well-represented clean states for a failed attempt to fit the noisy failure-state labels.

WM-based DAgger iterates from 1% to 13% route completion, confirming that the RSSM latent can be improved by on-policy data but remains hard-capped well below the direct policy. The large gap (13% vs. 41%) after equal DAgger budgets confirms that the latent representation — not the amount of on-policy data — is the binding constraint.

When World Models Compress Away the Signal

Our negative result on the RSSM latent is worth unpacking. The world model is trained to minimise reconstruction loss on a 259-dim observation. Of those 259 dimensions, the features most useful for reconstruction (LiDAR returns, which are smooth and predictable) are not the features most useful for action prediction (heading error and lateral lane offset, which are small in magnitude but carry the full lane-keeping signal). The RSSM optimises a reconstruction objective that weights all input dimensions equally, so the lane signal is under-represented in the latent relative to its policy-relevance.

This failure mode is specific to the mismatch between reconstruction-optimal and policy-optimal representations when observations are structured vectors rather than images. For image-based driving, the world model must decode every pixel including lane markings, forcing the latent to encode them. For vector observations, the latent is free to discard low-reconstruction-loss features regardless of policy relevance.

Recovery Data vs. Online Aggregation

Both DART and DAgger address distribution shift, but their stability properties differ at small compute budgets. DART injects controlled, bounded perturbations during collection: the car remains near the lane center (perturbation magnitude γ = 0.15, triangular decay), so the off-distribution states it visits form a narrow band around the nominal distribution. Expert labels at perturbed states are consistent and low-variance.

DAgger visits the policy's actual failure states, which are broader and more heterogeneous. IDM's relabeled actions at these states are well-defined but high-variance across the failure-state distribution and often conflict with training signals from clean demonstrations. With only 2,000 rollout steps per iteration — 12.5% of the 16,000-step clean buffer — the signal-to-noise ratio is low. At CPU-scale budgets where large rollout sets are prohibitive, DART-style static perturbation is a more reliable recovery data source.

Per-Scene Evaluation as a Development Tool

The per-scene evaluation methodology proved more informative than aggregate metrics throughout this project. The aggregate route completion of 39% suggested a broadly underperforming policy; the per-scene breakdown revealed that 75% of geometry types were solved at near-IDM level and identified a single failure mode. This allowed targeted intervention producing a 22pp per-scene gain at zero cost to other geometries — a gain that aggregate-metric-guided development would not have identified. Practitioners training driving policies in simulation should evaluate per-geometry from the start.

Limitations

All results are in MetaDrive simulation. Real-world transfer involves sensor noise, partial observability, and non-IDM traffic. The 259-dim state observation is provided by the simulator; a real deployment would require perception. All experiments ran on a single laptop CPU — the DAgger regression at 2,000 rollout steps/iteration might not appear at larger rollout budgets where the signal-to-noise ratio in the aggregated dataset is higher. IDM is a rule-based car-following model optimised for highway scenarios; its relabeled actions at roundabout entry arcs may not represent the globally optimal steering trajectory.