SafeDreamer: Safe Reinforcement Learning with World Models

Authors: Weidong Huang, Jiaming Ji, Chunhe Xia, Borong Zhang, Yaodong Yang

Venue: ICLR 2024

arXiv: 2307.07176

Abstract

SafeDreamer integrates Lagrangian-based optimization into world model planning within the Dreamer framework to achieve safe reinforcement learning. Existing SafeRL methods using cost functions often fail to achieve zero-cost performance in complex scenarios, especially vision-only tasks. SafeDreamer demonstrates nearly zero-cost performance on Safety-Gymnasium benchmark tasks spanning both low-dimensional and vision-only inputs.

Key Method

• Combines Lagrangian-based constraint satisfaction with DreamerV3-style world model planning
• Imagines future rollouts in latent space, applies Lagrangian penalty to cost-violating transitions
• Trains actor to maximize reward while satisfying safety constraints via dual-variable optimization
• Works on both state-based and image-based observations

Results

• Nearly zero constraint violations on Safety-Gymnasium benchmark
• Outperforms prior SafeRL methods (CPO, PCPO, PPO-Lag) in both safety and reward
• Effective for low-dimensional control and vision-only environments

Code

GitHub: PKU-Alignment/SafeDreamer