SafeDojo: Safe Reinforcement Learning for VLA via Interactive World Model

arXiv:2606.20698 · cs.RO · Submitted 2026-06-15 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: Robotics Radio. Generated commentary on the latest robotics and control papers.

Rosa: We're looking at the paper: "SafeDojo: Safe Reinforcement Learning for VLA via Interactive World Model".

Dev: Here's the gist of the paper. The gist: SafeDojo proposes a world model-based safe reinforcement learning framework for vision-language-action models that learns safe actions through world model-based imagination, enabling simultaneous assessment of task execution and trajectory safety Problem Definition and Preliminaries The paper addresses the challenge of learning collision-free safe VLA policies in obstacle-rich environments, where existing methods either require costly realworld exploration or depend on hand-crafted safety functions <ref:2606.20698#pg2>.

Rosa: Standard VLA MDPs are extended to Constrained MDPs (CMDPs) by introducing a safety cost function c: S × A →

0, 1: , requiring the policy πθ to maximize J r (πθ) = Eπθ

P t γ t rt: subject to J c (πθ) = Eπθ

P t γ t ct: ≤ d, where d denotes the allowable safety budget <ref:2606.20698#pg4>.

Dev: Model-based safe RL for VLA improves the policy using imagined trajectory feedback, rather than relying on static safety rules or risky real-world exploration <ref:2606.20698#pg3>.

Rosa: A trajectory τ is considered safe-successful only if it completes the instructed task without incurring collision cost, making safe-success rate a stricter criterion than task success alone <ref:2606.20698#pg4>.

Title and authors: Rosa: So we've seen the basics of what this paper is about, but let's talk about the title and who wrote it. This paper is "SafeDojo: Safe Reinforcement Learning for VLA via Interactive World Model", and it’s by Kai Tang, Peidong Jia, Zhong Chu, Jixian Wu, Rui Ma, Jiajun Cao, Fangyuan Zhao, Sixiang Chen, Yichen Guo Xiaowei Chi Chun-Kai Fan Kevin Zhang Jinchang Xu Fubing Yang Weishi Mi Xiaozhu Ju Jian Tang B Shanghang Zhang.

Dev: It's a big team of researchers from places like Peking University and Nanyang Technological University, so this is definitely coming from a solid group in the field. The title itself points directly to the method they use, which is building a world model to guide safe reinforcement learning for VLA policies.

Rosa: That's right. What this means in plain terms is they are proposing a new way for these models to learn how to act safely without needing massive amounts of real-world testing or manual safety rules that we have to write by hand.

Dev: They’re suggesting that instead of just trying to maximize the task reward, the policy learns by rolling out imagined trajectories through a world model and simultaneously assessing how safe those imagined actions are in terms of collision costs.

Taro: So, if I'm hearing this right, it’s about using imagination—the world model—to figure out what moves are okay before the robot actually tries them in the real world?

The paper's summary: Rosa: That’s a perfect way to put it. The paper explains that standard vision-language-action models are usually set up as Markov Decision Processes, but they extend this to Constrained MDPs, or CMDPs.

Dev: In a CMDP setup, the policy has two jobs: maximize the task reward J r while keeping a separate safety cost J c under a budget d. The budget is how much total "danger" we allow during the entire task.

Rosa: The paper then introduces model-based safe reinforcement learning for VLA, which uses this imagined trajectory feedback to improve the policy instead of relying on static safety rules or just risky real-world exploration.

Dev: They focus on the idea that a trajectory is only considered safe if it completes the task without exceeding that collision cost budget d, which is a stricter condition than just succeeding at the task itself.

Taro: So, they’re not just asking, "Did we finish?" but also, "Did we do it safely?" across the whole imagined path?

The paper's improvements: Rosa: Now let's look at what they actually did to make this work. They introduced three main innovations to handle the tension between exploring and causing unsafe contacts.

Dev: First, they have chunk-wise interactive world model rollouts where the model predicts future visual observations and latent dynamics based on the current action chunk. They even use a static-video augmentation occasionally to help keep those long rollouts from getting messed up by prediction errors.

Rosa: That’s clever because it’s trying to make the imagined path more robust against those long-horizon errors that usually plague these kinds of models, which is a real concern in practice.

Dev: Second, they decouple the task reward and safety cost estimation. They get dense feedback at every single control step by estimating the task reward r i,h and safety cost c i,h separately during the imagination process.

Rosa: And then they aggregate those per-step costs and rewards into trajectory-level pairs using a topM aggregation with M set to sixteen <ref:2606.20698#pg2>. This lets them look at short high-risk intervals rather than averaging the danger over the entire path, which is a smart way to focus on immediate risks.

Taro: So, they are getting two different signals—one for success and one for safety—and then smartly combining them based on where they matter most in the trajectory?

Conclusion: Rosa: So, to wrap up, SafeDojo establishes world model-based safe RL as a way to learn both task progress and safety signals from imagined trajectories. They show that this approach can be highly effective.

Dev: Their results on benchmarks like SafeLIBERO Level I show an eight point two five percentage-point improvement over the strongest baseline on that level in the average safe-success rate, which is a big jump for model-based methods compared to others tested <ref:2606.20698#pg2>.

Rosa: On real-world Franka deployments, they achieved the best average task and safe-success rates across five different manipulation tasks, showing it works outside of just lab settings.

Taro: I'm interested in their limitation here, though. They mentioned that the safety head is trained on binary collision labels. That means it might not be capturing finer details like force limits or proximity constraints during actual deployment, which is a real concern for robustness when the world misbehaves.

Dev: And they also pointed out that prediction errors in the world model can accumulate over longer rollout horizons, which could degrade the reliability of both task progress and safety evaluation if you go too far into imagination.

Rosa: So, while it’s a strong method for learning safe VLA policies—the SafeDojo: Safe Reinforcement Learning for VLA via Interactive World Model—they need to be careful about how those learned signals translate when the real world throws something unexpected at them.

Kai Tang, Peidong Jia, Zhong Chu, Jixian Wu, Rui Ma, Jiajun Cao, Fangyuan Zhao Sixiang Chen Yichen Guo Xiaowei Chi Chun-Kai Fan Kevin Zhang Jinchang Xu Fubing Yang Weishi Mi Xiaozhu Ju Jian Tang B Shanghang Zhang

State Key Laboratory of Multimedia Information Processing Peking University · Beijing Innovation Center of Humanoid Robotics · Nanyang Technological University · Hong Kong University of Science and Technology · University of Electronic Science and Technology of China

cs.RO

Submitted: 2026-06-15

Updated: 2026-10-21

Comments: CoRL 2026

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 91/100

The gist: The gist: SafeDojo proposes a world model-based safe reinforcement learning framework for vision-language-action models that learns safe actions through world model-based imagination, enabling

Key concepts

World Model: This is a learned representation of the environment that predicts future visual observations and latent dynamics based on current states, language instructions, and actions. SafeDojo uses this model for 'imagined rollouts'—simulating potential future scenarios internally before taking real-world actions.
Constrained MDP (CMDP): This is a standard reinforcement learning problem extended to include a safety constraint. The goal is not just to maximize task reward, but also to ensure the total accumulated 'safety cost' over time stays below a predefined limit. This forces the agent to learn policies that are both effective and safe.
Lagrangian-based Safety Constrained GRPO: This is a specific optimization technique used to update the policy. Instead of simple reward shaping, it uses a Lagrange multiplier ($\eta$) to combine the task reward and safety cost into a single 'safe advantage.' This allows the agent to learn how much it can sacrifice task progress for safety, dynamically adjusting based on constraint violations.
Decoupled Signal Estimation: SafeDojo estimates two separate signals at every step: a task reward ($r_{i,h}$) using a visual classifier and a safety cost ($c_{i,h}$) using an encoder. These dense predictions are combined via topM aggregation to distinguish between trajectories that are both successful and safe.

Terminology

Reinforcement Learning — 6× in this paper · explained in 4 episodes
World Model — 5× in this paper · explained in 8 episodes
GRPO — 1× in this paper · explained in 6 episodes
GRPO Objective — 1× in this paper · explained in 2 episodes
Online Reinforcement Learning — 1× in this paper · explained in 2 episodes

Summary

The gist: SafeDojo proposes a world model-based safe reinforcement learning framework for vision-language-action models that learns safe actions through world model-based imagination, enabling simultaneous assessment of task execution and trajectory safety

Problem Definition and Preliminaries

The paper addresses the challenge of learning collision-free safe VLA policies in obstacle-rich environments, where existing methods either require costly realworld exploration or depend on hand-crafted safety functions <ref:2606.20698#pg2>. Standard VLA MDPs are extended to Constrained MDPs (CMDPs) by introducing a safety cost function c: S × A → [0, 1], requiring the policy πθ to maximize J r (πθ) = Eπθ [P t γ t rt] subject to J c (πθ) = Eπθ [P t γ t ct] ≤ d, where d denotes the allowable safety budget <ref:2606.20698#pg4>. Model-based safe RL for VLA improves the policy using imagined trajectory feedback, rather than relying on static safety rules or risky real-world exploration <ref:2606.20698#pg3>. A trajectory τ is considered safe-successful only if it completes the instructed task without incurring collision cost, making safe-success rate a stricter criterion than task success alone <ref:2606.20698#pg4>.

Methodology

SafeDojo introduces three key innovations to address the tension between policy exploration and unsafe contacts <ref:2606.20698#pg2>.

Chunk-wise Interactive World Model Rollouts: SafeDojo performs model-in-the-loop rollouts inside an interactive world model pϕ, which acts as an action-conditioned video predictor that generates future visual observations and latent dynamics <ref:2606.20698#pg2>. The world model recursively rolls out each candidate by (ˆoi,k+1, zˆi,k+1) = pϕ(ˆoi,k, zˆi,k, ai,k, l), k = 0..., K − 1 <ref:2606.20698#pg5>. To improve rollout robustness against train-test mismatch caused by exposure error in long-horizon rollouts, a static-video augmentation is adopted to occasionally freeze the conditioning frames <ref:2606.20698#pg5>.
Decoupled Task Reward and Safety Cost Estimation: Given an imagined trajectory τˆi, SafeDojo assigns dense feedback by estimating task reward ri,h and safety cost ci,h at every single control step <ref:2606.20698#pg6>. These dense predictions are then aggregated into trajectory-level task reward and safety cost pairs (ri, ci) using a topM aggregation where M = 16 <ref:2606.20698#pg6>. The task reward branch ftask is instantiated with a ResNet success classifier finetuned on rollout frames, while the safety cost branch fsafe is a compact convolutional latent encoder followed by an action-conditioned MLP head that outputs per-step obstacle-contact probability <ref:2606.20698#pg6>.
Lagrangian-based Safety Constrained GRPO: Instead of reward shaping, SafeDojo decouples task reward and safety cost and optimizes them with a Lagrangian-based safety-constrained GRPO objective <ref:2606.20698#pg7>. The group-relative advantages are computed separately as Aˆr i = ri − µG(r) / σG(r) + ϵ and Aˆc i = ci − µG(c) / σG(c) + ϵ <ref:2606.20698#pg7>. The safety constraint is enforced by combining these advantages using a non-negative Lagrange multiplier η to compute the safe advantage Aˆsafe i = Aˆr i − ηAˆc i <ref:2606.20698#pg7>. The multiplier η is updated according to the batch-level constraint violation: η ← "η + αη / (1/G X G i=1 ci > d, else η ↓" <ref:2606.20698#pg8>. The final objective is LSafe-GRPO = −Ei,t h min ρi,t(θ)Aˆsafe i, clip (ρi,t(θ), 1 − ϵ, 1 + ϵ) Aˆsafe i i + βDKL πθ∥π ref θ <ref:2606.20698#pg9>.

Empirical Performance

SafeDojo demonstrates superior empirical performance across various benchmarks <ref:2606.20698#pg3>. On SafeLIBERO Level I, SafeDojo achieves the best aggregate task success, safe success, and execution efficiency among inference-time safety, model-free RL, and model-based RL baselines <ref:2606.20698#pg2>. Specifically on SafeLIBERO Level I, it achieved an 8.25 percentage-point improvement over the strongest baseline on Level I in the average safe-success rate <ref:2606.20698#pg8>. In realworld Franka deployment, SafeDojo further shows the best average task and safe-success rates across five tasks <ref:2606.20698#pg8>.

Ablation Study

Component ablation studies confirm that both components are critical for performance <ref:2606.20698#pg2> The safety head is shown to be critical as removing it reduces SSR to 38.0% <ref:2606.20698#pg2>. Removing the task progress reward branch collapses TSR and SSR to 3.0% and increases ETS to 298.3, indicating that task-reward feedback is essential for goal-directed policy improvement <ref:2606.20698#pg2>. Replacing GRPO with unconstrained optimization decreases TSR and SSR to 40.0% and 34.0%, validating the importance of safety-constrained policy updates <ref:2606.20698#pg2>.

Conclusion

SafeDojo establishes world model-based safe RL as a scalable and generalizable path toward safe embodied intelligence by learning both task progress and safety signals from imagined trajectories <ref:2606.20698#pg2>. Experiments on SafeLIBERO and real-world deployment demonstrate that SafeDojo achieves the best aggregate TSR, SSR, and ETS on both levels <ref:2606.20698#pg8>. The results establish world model-based safe RL as a viable paradigm for safe embodied intelligence <ref:2606.20698#pg2>.

Limitations and Future Work

Limitations include the safety head being trained on binary collision labels which may not capture finer-grained safety requirements such as force limits or proximity constraints <ref:2606.20698#pg2>. Furthermore, the world model introduces prediction errors that accumulate over longer rollout horizons, potentially degrading the reliability of both task progress and safety evaluation <ref:2606.20698#pg2>. The current evaluation focuses on tabletop manipulation with static obstacles, and deploying SafeDojo in more dynamic environments remains an open challenge <ref:2606.20698#pg2>. Finally, SafeDojo currently operates as a training-time safety mechanism and does not provide hard safety guarantees at inference time <ref:2606.20698#pg2>. Future work includes integrating learned safety signals with inference-time safeguards such as control barrier functions <ref:2606.20698#pg3>.

Broader Impact

SafeDojo advances the deployment of vision-language-action models on physical robots by enabling safer interaction with real-world environments, contributing to making embodied AI systems more reliable and trustworthy for applications such as household assistance, industrial automation, and human-robot collaboration <ref:2606.20698#pg2>. By reducing the risk of physical violations during policy learning, our framework lowers the barrier for deploying learning-based robotic systems beyond controlled laboratory settings <ref:2606.20698#pg2>. It is also noted that safety mechanisms learned from data may exhibit unexpected failure modes when deployed in scenarios that differ substantially from training distributions, and overreliance on learned safety signals without external safeguards could lead to false confidence in autonomous systems <ref:2606.20698#pg3>. Therefore, we advocate for combining SafeDojo with complementary safety measures, including runtime monitoring and human oversight, in real-world deployment <ref:2606.20698#pg3>. The technologies enabled by safer VLA deployment should be developed with consideration for ethical, economic, and societal implications <ref:2606.20698#pg3>.

Implementation Details

The policy actions are executed in chunks of H = 8 low-level control steps, matching the world-model rollout interface used by SafeDojo <ref:2606.20698#pg2>. The world model is initialized from Wan2.2 and fine-tuned on SafeLIBERO rollout trajectories converted into the Wan/RLinf video-action format <ref:2606.20698#pg2>. During imagined rollouts, the world model uses the current imagined observation, the latent context, the language instruction, and the next 8-step action chunk to predict future visual and latent states <ref:2606.20698#pg2>. The safety cost head operates on Wan latent features instead of RGB frames <ref:2606.20698#pg2>. The task reward head is implemented as a ResNet-style binary success classifier compatible with the reward model interface used in the world-model environment <ref:2606.20698#pg3>. The safety supervision is obtained from SafeLIBERO contact-collision auditing <ref:2606.20698#pg3>. Trajectory aggregation uses the topM aggregation in Eq. 4 with M = 16, emphasizing short high-risk intervals instead of diluting them over the full horizon <ref:2606.20698#pg2>. The reward and safety heads produce dense per-step signals along an imagined trajectory <ref:2606.20698#pg3>.

Real-World Results

SafeDojo achieves the best average TSR (76.00) and SSR (70.00) across five real-world Franka manipulation tasks <ref:2606.20698#pg2>. SafeDojo consistently generates direct yet collision-aware trajectories, avoiding both obstacle contacts and unnecessary detours <ref:2606.20698#pg2>.

Improvements for AI systems

World Model based Reward-Cost Estimation: The system can estimate per-step task progress from imagined frames using a finetuned ResNet success classifier and predict per-step safety costs from latent context and action chunk, enabling simultaneous assessment of task execution and trajectory safety.
Safety Constrained Policy Update: The system uses Lagrangian based constrained GRPO to optimize policies under explicit safety constraints, coordinating task success and collision avoidance under explicit constraints without relying on fixed penalty functions.
Decoupled Signal Optimization: By estimating dense feedback via ri,h = ftask(ˆoi,k(h), l), ci,h = fsafe(ˆzi,k(h), ai,h) and aggregating them using a topM selection (M=16), the system can distinguish between successful and safe execution and safe but ineffective trajectories.
Robustness to Rollout Errors: The system incorporates static-video augmentation during world-model training to make the model more tolerant to imperfect intermediate predictions during autoregressive chunk-wise rollout, addressing the train-test mismatch.
Safety Signal Sensitivity Tuning: The system can dynamically adjust the safety weight eta based on constraint violation, with a mechanism where eta increases to penalize unsafe trajectories more strongly when average safety cost exceeds the budget d=0.2.

Abstract

Safe control is a prerequisite for real-world embodied intelligence, for which safe reinforcement learning has emerged as a promising paradigm. However, existing safe reinforcement learning methods either require costly real-world exploration or depend on hand-crafted safety functions. Neither scales to vision-language-action models deployed in open-world physical environments. We propose SafeDojo, a model-based safe reinforcement learning framework that improves vision-language-action policies through world-model imagination. SafeDojo performs online reinforcement learning on top of an interactive video world model. A ResNet success classifier scores predicted future frames, while a lightweight safety head predicts per-step contact probabilities from the current latent context and proposed action chunk, jointly assessing task execution and trajectory safety. These decoupled signals are balanced through Lagrangian-constrained GRPO. On SafeLIBERO Level I, SafeDojo improves average safe-success rate by 8.25 percentage points over the strongest baseline; on the zero-shot Level II setting, it is competitive with the strongest model-free baseline. Real-world Franka deployment further shows the best average task and safe-success rates across five tasks. These results support world-model-based safe reinforcement learning as a promising path toward safer embodied intelligence.