Dreamer-SAC: Off-Policy Learning in Latent World Models for Sample-Efficient Autonomous Driving

The paper addresses the challenge of sample-efficient reinforcement learning for autonomous driving, which is often limited by the trade-off between data efficiency and model bias. The authors note that while world models reduce the reliance on costly environment interactions, policy optimization over learned dynamics remains sensitive to prediction errors. The primary bottleneck in autonomous driving deployment lies in "the prohibitive cost of safety validation, as ensuring reliable performance under long-tail and safety-critical scenarios, such as collisions, unexpected cut-ins, and complex interactions, requires extensive driving experience."

The paper identifies a key limitation in existing latent world model approaches: Most of these methods employ on-policy actor-critic optimization strategies, limiting the agent to training within the 'imagined' world and lacking the ability to learn directly from real-world data. Furthermore, "unlike real transitions, predicted trajectories inevitably suffer from accumulated prediction errors, causing model bias that increases rapidly with rollout horizon, which will affect the policy's assessment of safety-critical driving scenarios (such as collisions)."

The paper proposes the Dreamer-SAC framework, which integrates a recurrent state-space world model with an off-policy soft actor-critic algorithm trained directly in latent space. The framework uses a combination of real interactions and short-horizon generated trajectories with n-step target estimation and multi-objective supervision.

The key architectural components include:

1. Multi-modal Encoding: The observation at each timestep consists of a front-facing camera image It ∈ R H×W×3 providing rich visual context of the road environment, and a vector vt ∈ R 125 comprising LiDAR range measurements and ego-vehicle state information. The image is processed by a CNN, while the vector is processed by an MLP, and both embeddings are concatenated.

2. Latent Dynamics Modeling: The framework uses a Recurrent State-Space Model (RSSM) that maintains a deterministic hidden state ht that summarizes past information, and a stochastic state zt that captures uncertainty. The transition is defined with a prior distribution for latent rollouts and a posterior distribution when real observations are available.

3. Multi-objective Reward: The reward function consists of three components: driving efficiency reward Reff (encouraging forward progress and appropriate speed), lane center offset penalty Rlane (penalizing deviation from lane center), and termination reward/punishment Rterm (sparse terminal reward for completing route or penalty for violations). The framework employs three independent prediction heads where efficiency and lane-center rewards use MSE loss, while terminal reward uses a symlog-discretized categorical distribution.

4. Hybrid Real and Model-Generated Replay: The framework samples trajectory segments from the real replay buffer, uses them to optimize the RSSM, and generates latent rollouts from posterior states. "Since rollout is performed from every sampled real transition, the number of predicted transitions is directly determined by the rollout horizon H. Consequently, the ratio between real and predicted samples automatically becomes 1: H. Predicted trajectories are discarded immediately after policy optimization" and regenerated with the latest models.

5. Multi-step Bootstrapping: Real transitions use the standard one-step SAC target, while predicted trajectories use an n-step target that propagates long-horizon reward information to the critic while reducing bootstrap variance inside predicted trajectories.

The paper lists three main contributions:

1. "A model-based off-policy reinforcement learning framework is proposed by integrating Soft Actor-Critic into the Recurrent State-Space Model, enabling autonomous driving policies to efficiently exploit model-generated rollouts."

2. "A hybrid real and model-generated experience learning strategy is developed to improve policy reliability by combining real transitions and latent rollouts with different value estimation schemes, thereby enhancing data utilization while mitigating the impact of model prediction errors."

3. "Comprehensive experiments demonstrate the superiority of the proposed framework over representative reinforcement learning baselines and reveal an inverted-U relationship between rollout horizon and performance, indicating that short-horizon rollouts provide the optimal balance between learning benefits and model bias."

The framework is evaluated on the MetaDrive simulator, specifically the BIG BLOCK SEQUENCE-CC scenario, which consists of a multi-lane highway segment with consecutive curved sections. Key parameters include traffic density of 0.1 (Approximately 30 vehicles per kilometer), 10 distinct training scenarios and 50 distinct testing scenarios. Each episode runs for up to 300 decision steps, with each decision step corresponding to 20 simulation frames.

The world model uses an RSSM with a deterministic recurrent state of 256 dimensions and a discrete stochastic state with 32 categorical variables and 32 classes per variable. The resulting latent feature dimension is 1280. The SAC policy module operates directly on RSSM latent features with two fully connected layers with 256 hidden units. Training is performed for 40,000 environment interaction steps with a rollout horizon of H = 5.

The proposed Dreamer-SAC is compared against three baselines: DreamerV3 (latent world-model-based on-policy learning), SAC (off-policy maximum-entropy RL), and PPO (on-policy policy-gradient method).

Training Return Results: The proposed Dreamer-SAC achieves the highest final performance, reaching an average return of 371.4. In comparison, DreamerV3 achieves an average return of 189.2 and exhibits relatively high variance, SAC converges at 134.5, and PPO reaches a score of 65.5.

Generalization Evaluation: On extended road networks approximately three times longer than training environments, evaluated over 50 independent trials with unseen random seeds, Dreamer-SAC achieves the lowest collision frequency (1.56/km) and out-of-road frequency (1.37/km), while also obtaining the longest average travel distance (320.8 m) and maximum travel distance (951.6 m).

An interesting finding is that DreamerV3 achieves the highest average speed (39.7 km/h) but suffers from higher collision frequency and shorter travel distance. The paper explains this by noting that "training exclusively with predicted experiences may amplify model bias, where safety-critical transitions are inaccurately represented in the learned dynamics. Consequently, the policy may exploit inaccuracies in the predicted environment and adopt overly aggressive behaviors."

In contrast, "Dreamer-SAC employs a dual-network architecture and learns a more conservative and reliable driving strategy by continuously incorporating real experiences during policy optimization. When encountering slower vehicles, the agent tends to follow the preceding vehicle and reduce speed when necessary rather than aggressively performing risky maneuvers."

The paper reveals a clear inverted-U relationship between rollout horizon and performance. "Without rollouts (H = 0), the agent fails to learn meaningful driving behavior and achieves an average return of only −2.6. Increasing the horizon significantly improves performance, reaching the best result at H = 5 with an average return of 371.4. However, further increasing the horizon leads to gradual degradation, with returns decreasing to 285.6 at H = 20."

The paper explains that the poor performance without rollouts indicates that directly optimizing SAC in the learned latent space using only real samples is insufficient during early training. With rollouts, short-horizon rollouts provide reliable synthetic experiences and improve exploration, whereas excessive rollouts amplify prediction errors and lead to overestimated value functions.

Additional analysis shows that the average predicted Q-value increases with larger horizons, and the Q-value variance also becomes larger, suggesting that the critic obtains stronger action discrimination but also becomes increasingly optimistic. Notably, although H = 20 achieves the highest Q-value estimates, its final return is lower than that of H = 5, revealing a mismatch between value estimation and actual driving performance.

The ablation studies investigate the contribution of each major component:

1. Effect of real and predicted data combination: Removing real samples decreases performance for both short and long horizon rollout settings. With H = 1, training solely on predicted data achieves 167.3 return, while incorporating real transitions improves to 226.2. With H = 5, the gap is smaller (347.5 vs. 371.4) because longer rollouts generate more synthetic samples.

2. Effect of multi-objective world model learning: "Detaching the observation reconstruction gradient causes a severe performance degradation (A5, return of −5.6), indicating that reconstruction provides an essential regularization signal for learning meaningful latent representations. Detaching reward and continuation gradients while retaining only reconstruction results in moderate degradation (280.8), suggesting task-oriented prediction gradients provide additional guidance by encouraging the latent dynamics to preserve decision-relevant information."

3. Effect of temporal target estimation: The results show that n-step targets achieve better performance (371.4 vs. 331.6), demonstrating that latent rollouts contain useful future information beyond immediate transitions.

4. Effect of stochastic latent representation: "Replacing the discrete latent state with a continuous Gaussian distribution reduces the final performance (303.4 vs. 371.4), suggesting that the discrete stochastic representation provides a more expressive and stable latent space for modeling complex driving dynamics."

The paper performs a trajectory-level comparison between real and predicted transitions by randomly sampling 512 trajectories from the replay buffer and initializing rollouts from the same real states with identical action sequences. The results show:

• The cosine similarity between real and predicted hidden states gradually decreases from 0.438 at the first predicted step to 0.371 after five steps, indicating that latent dynamics remain consistent with the real environment in short horizons but gradually accumulate prediction errors during longer rollouts.

• The reward prediction error increases with rollout depth. The mean absolute error grows from 0.855 at the first step to 3.698 at the fifth step, representing more than a four-fold increase.

• "The predicted distribution assigns higher probability to extreme negative rewards, indicating that the world model tends to overestimate the occurrence of safety-critical failures such as collisions or off-road events. This conservative prediction behavior results from the difficulty of accurately modeling rare transition events with limited real interactions."

The paper concludes that Dreamer-SAC improves policy learning while reducing the adverse effects of accumulated model errors by generating short-horizon rollouts from real trajectory states and incorporating n-step target estimation. The experiments demonstrate that the proposed approach consistently outperforms representative reinforcement learning baselines, including DreamerV3, SAC, and PPO.

Future work will focus on developing adaptive rollout strategies that dynamically adjust rollout length according to model reliability and uncertainty and evaluating the proposed framework with larger-scale real-world driving datasets and more diverse traffic scenarios.

Improvements to AI systems based on this paper:

1. Hybrid Experience Replay with Model-Error Awareness: Implement a dual-replay system that combines real transitions (1-step targets) with short-horizon latent rollouts (n-step targets) in a fixed 1:H ratio. The improved system can automatically balance exploitation of real data with exploration of model-generated scenarios, reducing model bias while maintaining sample efficiency.

2. Adaptive Rollout Horizon Control: Replace fixed rollout lengths with a dynamic mechanism that monitors prediction error (e.g., cosine similarity between real and predicted hidden states, or reward prediction MAE) and adjusts rollout depth accordingly. The improved system can maintain the optimal inverted-U performance point by shortening rollouts when model uncertainty increases and lengthening them when the model is reliable.

3. Multi-Objective Latent Representation Learning: Incorporate three parallel prediction heads (efficiency reward, lane-center penalty, and symlog-discretized terminal reward) with separate loss functions (MSE for continuous rewards, categorical cross-entropy for terminal events). The improved system can learn richer latent dynamics that preserve both continuous driving quality and discrete safety-critical event information.

4. Conservative Value Estimation for Safety-Critical Domains: Use n-step bootstrapping specifically for predicted trajectories while keeping 1-step targets for real data, combined with the observed tendency of world models to overestimate negative outcomes. The improved system can produce more reliable value estimates that avoid over-optimism in long-horizon planning, particularly for collision and off-road risk assessment.

5. Discrete Stochastic State Representation with Reconstruction Regularization: Maintain the discrete latent space (32 categorical variables × 32 classes) and retain observation reconstruction gradients during world model training. The improved system can achieve more stable and expressive latent dynamics, preventing representation collapse and preserving decision-relevant information.

6. Conservative Driving Policy via Real-Experience Anchoring: Continuously inject real transitions during policy optimization to counteract the aggressive behaviors that emerge from pure model-based training (as observed with DreamerV3). The improved system can learn safer driving strategies that follow slower vehicles and reduce speed appropriately rather than exploiting model inaccuracies.

What the improved AI system can do:

• Learn autonomous driving policies with 2-3× higher final returns than on-policy world model methods (371 vs. 189) while maintaining lower collision rates (1.56/km vs. higher) in unseen environments.

• Dynamically balance sample efficiency and model bias by adjusting rollout horizons based on real-time prediction confidence, avoiding both undertraining (H=0) and over-optimistic value estimates (H=20).

• Generalize to longer, unseen road networks (3× longer than training) with reliable safety performance, maintaining low collision and out-of-road frequencies while achieving longer travel distances.

• Provide calibrated uncertainty about safety-critical events by leveraging the world model's conservative bias toward overestimating negative outcomes, enabling safer decision-making in long-tail scenarios.

• Operate with only 40,000 environment steps (vs. millions typically required), making it feasible for real-world deployment where interaction costs are prohibitive.