PaperlayerView exact revision on arXiv

Source-linked AI summary

SafeDreamer: Safe Reinforcement Learning with World Models

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

arXiv:2307.07176v3cs.LGcs.AIcs.CV

TL;DR

SafeRL methods can struggle to achieve near-zero costs in complex tasks, including vision-only settings. SafeDreamer combines Lagrangian methods with safety-reward planning in a world model, and reports near-zero-cost performance across low-dimensional and vision-only Safety-Gymnasium tasks.

  • Problem

    Existing SafeRL approaches can fail to achieve zero-cost performance in complex scenarios, especially vision-only tasks, motivating safety-aware planning that balances long-term rewards and costs.

    Source: §1 · p. 1

  • Method

    SafeDreamer integrates Lagrangian methods with online and background safety-reward planning in a world model to balance rewards and costs.

    Source: §1 · pp. 1–2

  • Results

    SafeDreamer achieves nearly zero-cost performance across low-dimensional and visual Safety-Gymnasium tasks and surpasses competing model-based and model-free methods.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors report SafeDreamer as the first algorithm to achieve nearly zero-cost final performance with vision-only input in Safety-Gymnasium.

    Source: §7 · p. 9

    The formulation assumes the initial state is sampled from the true initial-state distribution and restricts discussion to one cost constraint.

    Source: §3 · p. 4

Abstract

from arXiv · show

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}.

1 INTRODUCTION

SafeDreamer combines world-model safety planning with Lagrangian methods to balance rewards and costs. It targets nearly zero-cost performance across low-dimensional and vision-only Safety-Gymnasium tasks.

1 INTRODUCTION

  • Existing Lagrangian methods can miss near-zero cost constraints or meet them without completing tasks, while finite-horizon planning can produce locally optimal, unsafe solutions.Obtaining a ground-truth dynamics model for complex vision-only scenarios is described as impractical.Source: §1 · p. 1
  • Prior world-model methods have not achieved near-zero-cost performance across several environments, which the authors attribute to inaccurate safety modeling of current or future states.SafeDreamer integrates safety planning and Lagrangian methods within a world model to balance errors between cost models and critics.Source: §1 · p. 1
  • OSRP plans safety-reward action trajectories within the world model, using the Constrained Cross-Entropy Method for vision-only tasks.The introduction presents this as evidence that online planning in a world model can satisfy constraints in vision-only tasks.Source: §1 · pp. 1–2
  • OSRP-Lag and BSRP-Lag integrate Lagrangian methods with online and background world-model planning to balance long-term rewards and costs.Online planning generates actions, whereas background planning updates a safe actor.Source: §1 · p. 2, Figure 1 · p. 2
  • SafeDreamer achieves nearly zero-cost performance on low-dimensional and visual Safety-Gymnasium tasks and outperforms competing model-based and model-free methods.The conclusion identifies vision-only near-zero-cost performance in the benchmark as a first for the algorithm.Source: §1 · p. 2

2 RELATED WORK

Prior SafeRL work includes constrained policy optimization, safety certificates, and model-based planning, but the surveyed methods leave gaps in zero-violation performance, long-term safety, or vision input. SafeDreamer combines planning and Lagrangian methods in a framework intended for both input modalities.

2 RELATED WORK

  • Prior SafeRL methods struggle to achieve zero violations, while model-free safe algorithms require substantial environmental interaction and model-based methods can improve sample efficiency.The comparison table frames SafeDreamer as accommodating both low-dimensional and vision inputs while unifying planning-based and Lagrangian-based approaches.Source: §2 · p. 2, Table 1 · p. 3
  • Online planning methods such as MPC can generate safe actions, but limited planning scope and absent critics can make decisions short-sighted; terminal values address long-term reward, not long-term safety.The surveyed background-planning methods use safety value functions with PPO or SAC, but adapting them to vision input remains challenging.Source: §2 · pp. 2–3
  • LAMBDA extends DreamerV1 with Lagrangian and Bayesian principles, but the review cites instability and missing adaptive modifications as limits, while Safe SLAC does not use online or background planning.Safe SLAC achieves performance comparable to LAMBDA on vision-only tasks, according to the review.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries formulate SafeRL as a constrained Markov decision process and define reward maximization subject to cost limits. For model-based RL, imagined world-model trajectories provide the returns used in policy optimization.

3 PRELIMINARIES

  • A CMDP specifies states, actions, transition probabilities, rewards, costs with thresholds, an initial-state distribution, and a discount factor.The policy objective is to maximize expected reward while satisfying cost constraints.Source: §3 · p. 3
  • In the safety-reward planning process, the encoder maps observations to latent states, and the world model predicts trajectory rewards, costs, and their values for policy updates.TD(λ) estimates reward and cost returns for each trajectory.Source: Figure 2 · p. 4
  • The safe model-based formulation rolls out imagined trajectories in a parameterized world model to estimate reward and cost returns for policy optimization.The discussion restricts attention to one cost function and threshold and assumes the initial state comes from the true initial-state distribution.Source: §3 · p. 4

4 METHODS

SafeDreamer combines a learned world model with safety-reward planning and Lagrangian methods to balance rewards and costs. Its methods use latent rollouts for online or background planning.

4 METHODS

  • The framework combines world-model safety-reward planning with Lagrangian methods for online or background planning.The world model is trained from a replay buffer as the agent interacts with the environment.Source: §4 · p. 4

4.1 MODEL COMPONENTS

  • The world model encodes observations as discrete representations, predicts subsequent representations, and uses model states to predict observations, rewards, costs, value estimates, and actions.A model state concatenates the recurrent state and discrete representation.Source: §4.1 · p. 4

4.2 ONLINE SAFETY-REWARD PLANNING (OSRP) VIA WORLD MODEL

  • OSRP samples trajectories from an action distribution and a reward-driven actor, evaluates their rewards and costs, then selects an action trajectory using constrained cross-entropy planning.When fewer than the desired number of trajectories satisfy the cost threshold, it ranks candidates by lower cost; otherwise, OSRP ranks safe trajectories by reward, and the first action of the final sampled trajectory is executed.Source: §4.2 · pp. 5–6

4.3 LAGRANGIAN METHOD WITH THE WORLD MODEL PLANNING

  • OSRP-Lag applies a Lagrangian multiplier in online planning to balance reward and cost, with episode costs informing multiplier updates.Its stated intuition is to optimize a conservative policy under comparatively safe conditions while considering long-term risks.Source: §4.3 · p. 6
  • BSRP-Lag uses background imagined rollouts and an Augmented Lagrangian objective to train a safe actor, avoiding online planning during action generation.The method is presented as an extension for scenarios with real-time requirements; its imagined latent rollouts have length T = 15.Source: §4.3 · p. 6

5 PRACTICAL IMPLEMENTATION

SafeDreamer extends DreamerV3 with a world model and actor-critic components trained from replay-buffer sequences. Its world model learns latent representations and predicts observations, rewards, and costs.

5 PRACTICAL IMPLEMENTATION

  • SafeDreamer extends the DreamerV3 architecture to build a safe model-based reinforcement-learning algorithm.The world model is trained using a variational auto-encoding approach that maps observations to latent representations used to evaluate reward and safety.Source: §5 · pp. 6–7
  • The world model uses a Recurrent State Space Model with a GRU for sequence modeling.Its latent representations are regularized toward a prior predicted by the sequence model.Source: §5 · p. 7
  • Training samples replay-buffer sequences to update the world model, while the safe actor and Lagrangian multiplier are updated through background safety-reward planning.In OSRP, a reward-driven actor guides planning; BSRP-Lag instead uses its safe actor to generate actions without online planning.Source: §5 · p. 7

6 EXPERIMENTAL RESULTS

In Safety-Gymnasium experiments, SafeDreamer achieved nearly zero-cost performance while maintaining reward, including across vision tasks. Results also varied by planning strategy and environment: background planning helped on precision tasks, while online planning performed better with moving obstacles.

6 EXPERIMENTAL RESULTS

  • Evaluation used five robotic environments, episodic reward and cost, and training cost rate; comparisons included model-free and model-based baselines.Reward and cost were averaged over ten 1,000-step evaluation episodes; cost rate used costs incurred during training.Source: §6 · pp. 7–8
  • After 2M steps, SafeDreamer (BSRP-Lag) achieved rewards similar to LAMBDA and reduced costs by 94.3%, attaining nearly zero-cost performance.Figure 3 summarizes normalized reward and cost for model-based methods on five vision tasks.Source: Figure 3 · p. 8, §6 · p. 8
  • Across five vision tasks, SafeDreamer combined minimal costs with higher rewards; Safe SLAC hit a cost floor, while MPC struggled with dense or dynamic obstacles.In SafetyPointGoal1 and SafetyPointPush1, SafeDreamer matched DreamerV3's reward while preserving nearly zero cost.Source: §6 · p. 8
  • On two vector-input tasks, SafeDreamer surpassed model-free algorithms and MBPPO-Lag on reward and cost; limited-data trial-and-error hindered model-free agents' search for higher rewards.MBPPO-Lag's reward stopped increasing once costs became relatively low.Source: §6 · p. 8
  • On dynamic-obstacle SafetyPointButton1, OSRP outperformed BSRP-Lag in reward; the authors suggest online planning may explain the advantage.Source: §6 · p. 8
  • In SafetyPointPush1, BSRP-Lag found the dexterous policy while online planners struggled; in SafetyPointGoal2, OSRP-Lag's cost critic made it more conservative than OSRP.The authors also report that world-model reconstruction loss and latent hidden-state size affect cost reduction.Source: §6 · p. 9

7 CONCLUSION

SafeDreamer combines world-model safety-reward planning with Lagrangian methods to balance rewards and costs, and reports nearly zero-cost final performance with vision-only input in Safety-Gymnasium. The conclusion notes that each task is trained independently, making training costly, and identifies multi-task pretraining as future work.

7 CONCLUSION

  • SafeDreamer is presented as the first algorithm to achieve nearly zero-cost final performance using vision-only input in Safety-Gymnasium.It uses safety-reward planning in world models and Lagrangian methods to balance rewards and costs.Source: §7 · p. 9
  • Training each task independently incurs substantial costs; the authors propose multi-task offline-data pretraining of the world model as future work.They also identify safe robot deployment and constrained-optimization approaches to LLM safety as potential directions.Source: §7 · p. 9
  • Figure 5 reports results for low-dimensional input tasks.The conclusion points to these results alongside its vision-only benchmark claim.Source: Figure 5 · p. 9, §7 · p. 9

9 REPRODUCIBILITY STATEMENT

The authors provide code and more than 80 model checkpoints for training and evaluating SafeDreamer agents. They report the benchmark environments and baseline implementations used in the experiments.

9 REPRODUCIBILITY STATEMENT

  • The authors release 80+ model checkpoints and code for training and evaluating SafeDreamer agents.The resources are available through the SafeDreamer repository.Source: §9 · p. 10
  • Experiments use Safety-Gymnasium, MetaDrive3, and Gymnasium, with baselines implemented from named repositories and Omnisafe.The listed baselines include LAMBDA, Safe SLAC, MBPPO-Lag, Dreamerv3, CPO, TRPO-Lag, PPO-Lag, Sauté PPO, and Sauté TRPO.Source: §9 · p. 10

A DETAILS OF SAFEDREAMER TRAINING PROCESS

SafeDreamer trains a world model and actor-critic from replayed experience, then uses latent rollouts and safe-actor or online-planner actions in environment interaction. Its PID Lagrangian planning variant addresses a planning-horizon limitation, while cost-threshold and critic choices vary by setting.

A.1 THE TRAINING PROCESS OF SAFEDREAMER

  • Training samples replay trajectories, updates the world model, encodes observations into latent states, generates actor rollouts, updates reward and cost critics, and collects new environment experience.Actions come from either the safe actor or Online Safety-Reward Planner.Source: §A.1 · p. 14

A.2 THE PROCESS OF PID LAGRANGIAN

  • OSRP-Lag adds PID Lagrangian to online safety-reward planning to address cases where hazard detection requires a longer horizon than the planner uses.The stated planning constraint is (J_C,H^ϕ/H)L < b.Source: §A.2 · p. 14
  • A cost threshold of 2.0 consistently converges toward nearly zero cost in vision tasks; BSRP-Lag can use zero on some tasks but needs more samples to converge.For BSRP-Lag, using summed rather than average cost returns over H steps helps the policy more closely approximate zero cost; low-dimensional tasks use observation reconstruction and a zero threshold.Source: §A.2 · pp. 14–15
  • The reward-driven actor balances expected reward and entropy, while reward and cost critics use discrete regression, twohot targets, and symlog smoothing.The cited techniques address sparse, heterogeneous costs that challenge direct cost-value regression.Source: §A.2 · p. 15

B HYPERPARAMETERS

The appendix reports experimental hardware and hyperparameter guidance for SafeDreamer and several baselines. It recommends tuning selected Lagrangian parameters across scales for other safety tasks and documents baseline-specific configuration changes.

B HYPERPARAMETERS

  • SafeDreamer experiments used Python 3, JAX 0.3.25, CUDA 11.7, Ubuntu 20.04.2, and a machine with 40 CPU cores, 251 GB RAM, and eight GeForce RTX 3090Ti GPUs.The reported software and hardware provide the experiment’s computational setting.Source: B HYPERPARAMETERS · p. 16
  • For other safety tasks, the authors recommend tuning the initial Lagrangian multiplier and PID coefficients at various scales.The coefficients are proportional, integral, and difference.Source: Table 2 · p. 16
  • Safe SLAC settings change action repeat from 2 to 4, sampled sequence length from 10 to 15, and cost limit from 25.0 to 2.0 for comparison.The LAMBDA cost limit is also changed from 25.0 to 2.0; model-free baselines use larger MLPs and a cost discount factor of 0.999 for Sauté RL methods.Source: Table 3 · p. 17, Table 4 · p. 17, Table 6 · p. 18

C.1 EXPERIMENTS ON METADRIVE

In MetaDrive, SafeDreamer reached a 100% success rate in 10 evaluation tests after 4M environment steps and reduced costs more efficiently than PPO-Reward Shaping. The experiment assesses driving to a destination while penalizing collisions and leaving the road.

C.1 EXPERIMENTS ON METADRIVE

  • MetaDrive combines vector and first-person-view observations; collisions or leaving the road incur a cost penalty.The driving objective is to navigate successfully to a predetermined destination.Source: §C.1 · p. 19, Figure 6 · p. 19
  • 100% success in 10 evaluation tests after 4M environment steps; SafeDreamer reduced costs more efficiently than PPO-Reward Shaping.Training and testing used the same roadway, with vehicle positions randomized at the start of each episode.Source: Table 7 · p. 19, §C.1 · p. 19

C.2 EXPERIMENTS ON CAR-RACING

The experiments cover Car-Racing, FormulaOne, and Safety-Gymnasium, with additional descriptions of the benchmark setup and ablations. SafeDreamer reaches near-zero cost in the two reported driving tasks, while the ablations examine factors associated with cost reduction.

C.2 EXPERIMENTS ON CAR-RACING

  • SafeDreamer reaches near-zero cost within 0.2M steps on Car-Racing and significantly reduces vehicle slippage compared with DreamerV3.The passage describes the reduction in slippage as enhancing driving robustness.Source: §C.2 · p. 20, Table 8 · p. 20

C.3 FORMULAONE

  • SafeDreamer attains nearly zero cost after 1.5M training steps in FormulaOne while completing the task with obstacle avoidance.The task uses 64x64x3 image inputs and requires reaching the goal while avoiding barriers and racetrack fences.Source: §C.3 · p. 21, Table 9 · p. 21

C.4 EXPERIMENTS ON SAFETY-GYMNASIUM

  • Safety-Gymnasium experiments use five tasks, forward and rear camera views, and a modified state representation intended to facilitate model learning.The tasks are SafetyPointGoal1, SafetyPointGoal2, SafetyPointPush1, SafetyPointButton1, and SafetyCarGoal1; all experiments use identical settings.Source: §C.4 · p. 22, Figure 11 · p. 22, Figure 12 · p. 22

C.4.1 AGENT SPECIFICATION

  • The benchmark describes Point, Car, and Racecar robots, with continuous actions scaled to [-1, +1].Their control setups range from separate rotation and linear actuators for Point to coordinated steering and speed control for Racecar.Source: §C.4.1 · pp. 22–23

C.4.2 TASK REPRESENTATION

  • Goal, Push, and Button tasks use changing targets and combine sparse completion rewards with dense progress rewards.Push rewards include progress narrowing the agent-box distance and moving the box toward the target; Button tasks advance through highlighted buttons.Source: §C.4.2 · p. 23

C.4.3 CONSTRAINT SPECIFICATION

  • Safety costs arise from contact with pillars, hazards, or vases, and from contact with gremlins in Button tasks.Vases are used only in Goal tasks; gremlins are dynamic objects used specifically in Button tasks.Source: §C.4.3 · pp. 23–24

C.4.4 EVALUATION METRICS

  • Evaluation ranks safety-constraint satisfaction first, then return among feasible agents, and cost convergence speed when feasible agents have similar rewards.The reported metrics are average episodic return, average episodic cost, and cost rate.Source: §C.4.4 · p. 24

C.4.5 SAFEDREAMER TRAINING RESULTS WITHIM 8M STEPS

  • Video-prediction figures show task-specific forecasts for Point and Racecar, while Figure 15 covers training within 8M steps and Figures 18–20 compare BSRP-Lag with DreamerV3 across three agents.The model predicts the next 45 steps from the previous 25 frames and an action sequence, without intermediate image access.Source: Figure 15 · p. 24, §C.4.6 · p. 25, Figure 16 · p. 25
Loading 2307.07176v3…