SafeDreamer: Safe Reinforcement Learning with World Models

Abstract: The deployment of Reinforcement Learning (RL) in real-world applications is constrained by its failure to satisfy safety criteria. Existing Safe Reinforcement Learning (SafeRL) methods, which rely on cost functions to enforce safety, often fail to achieve zero-cost performance in complex scenarios, especially vision-only tasks. These limitations are primarily due to model inaccuracies and inadequate sample efficiency. The integration of the world model has proven effective in mitigating these shortcomings. In this work, we introduce SafeDreamer, a novel algorithm incorporating Lagrangian-based methods into world model planning processes within the superior Dreamer framework. Our method achieves nearly zero-cost performance on various tasks, spanning low-dimensional and vision-only input, within the Safety-Gymnasium benchmark, showcasing its efficacy in balancing performance and safety in RL tasks. Further details can be found in the code repository: \url{https://github.com/PKU-Alignment/SafeDreamer}.

• The paper introduces SafeDreamer, an algorithm that integrates world model planning with Lagrangian optimization to balance reward maximization and strict safety constraints in RL tasks.
• It employs the Constrained Cross-Entropy Method for online action planning, ensuring near-zero-cost safety while maintaining high task performance.
• Experimental results on Safety-Gymnasium benchmarks demonstrate superior sample efficiency and adaptability in both low-dimensional and vision-based scenarios.

Introducing SafeDreamer: Bridging Safety and Performance in Reinforcement Learning through World Models

Overview

In the landscape of Reinforcement Learning (RL), deploying algorithms in real-world scenarios often comes with a paramount challenge: ensuring the safety of their actions. Recent advancements in Safe Reinforcement Learning (SafeRL) have sought to address this issue by incorporating safety criteria directly into the learning process. However, existing SafeRL methods struggle to strike a balance between achieving task objectives and adhering to safety constraints, particularly in complex, vision-based tasks. This paper introduces SafeDreamer, a novel algorithm that marries Lagrangian-based methods with the prowess of world model planning, emerging from the well-regarded Dreamer framework. SafeDreamer demonstrates a compelling capability to achieve near-zero-cost performance while maintaining task performance across various benchmarks, providing a promising avenue for enhancing both the safety and efficiency of RL tasks.

Methodology

The crux of SafeDreamer lies in its unique integration of world model planning and Lagrangian-based optimization to navigate the intricate balance between performance and safety. The algorithm encapsulates three main components:

• Online Safety-Reward Planning (OSRP): At its core, OSRP employs the Constrained Cross-Entropy Method (CCEM) for action trajectory planning within learned world models. This process selectively evaluates and optimizes action trajectories based on their anticipated rewards and safety, considerably enhancing the agent's ability to foresee and mitigate potential risks.
• Lagrangian Methods with World Models Planning: To further refine the safety-performance trade-off, SafeDreamer incorporates Lagrangian techniques into both online and background planning stages. This approach adeptly corrects for model inaccuracies by dynamically adjusting the balance between reward maximization and cost minimization, paving the way for adaptive policy optimization under constraints.
• Implementation of SafeDreamer: Built upon the architecture of DreamerV3—an advancement in model-based RL—SafeDreamer extends its utility by introducing safety-aware planning and optimization mechanisms. Through meticulous training of world models and actor-critic networks, the algorithm efficiently generates and evaluates latent state trajectories, enabling effective policy updates that honor both reward acquisition and safety considerations.

Experimental Validation

The efficacy of SafeDreamer was rigorously evaluated against a spectrum of model-based and model-free SafeRL algorithms across diverse tasks within the Safety-Gymnasium benchmark. The benchmarks encompassed a range of low-dimensional and vision-based tasks, designed to test the algorithm's versatility and robustness in varying contexts. The results unequivocally demonstrated that SafeDreamer outperforms existing SafeRL methods, achieving near-zero-cost performance while preserving or even enhancing task performance. Notably, SafeDreamer showcased superior sample efficiency and adaptability in vision-based tasks, underscoring its potential applicability in real-world scenarios where safety constraints are paramount.

Implications and Future Directions

The introduction of SafeDreamer marks a significant stride towards reconciling performance with safety in RL applications. By leveraging world model planning within a Lagrangian framework, SafeDreamer not only enhances task performance but also ensures adherence to stringent safety criteria. The algorithm's success in diverse benchmarks hints at its potential applicability across various real-world domains, from autonomous vehicles to robotic manipulation, where safety and efficiency are critical. Looking ahead, further exploration into the scalability of SafeDreamer, its applicability to more complex scenarios, and its integration with offline data for pre-training promise exciting avenues for advancing SafeRL towards practical deployment. In doing so, SafeDreamer not only elevates the state-of-the-art in SafeRL but also paves the way for safer, more reliable RL applications in the real world.