Technical Summary: SafeDojo

Problem Definition

The deployment of Vision-Language-Action (VLA) models in real-world physical environments faces a critical safety challenge: collision avoidance during contact-rich manipulation. Unlike low-dimensional control systems where safety constraints can be explicitly defined or compactly learned, VLA policies must infer safety boundaries from high-dimensional visual observations across diverse and previously unseen obstacle configurations.

Existing approaches to safe robotic manipulation suffer from significant limitations:

• Inference-time interventions (e.g., Control Barrier Functions) require precise geometric modeling and camera calibration, which are impractical for open-world VLA deployment.
• Model-free constrained Reinforcement Learning (RL) enforces safety through penalty-based optimization but requires costly and potentially unsafe real-world exploration. A single unsafe contact can cause task failure, irreversible hardware damage, or environmental disruption.
• Model-based safe RL typically operates in low-dimensional state spaces with simulator-provided cost functions, making it difficult to transfer to high-dimensional VLA models where safety signals must be learned from visual data rather than analytically defined.

The core open problem addressed is how to learn collision-free VLA policies in obstacle-rich environments without relying on hand-crafted safety functions or risky real-world exploration.

Methodology: SafeDojo

SafeDojo is proposed as the first model-based safe reinforcement learning framework specifically designed for VLA policies. It enables online policy optimization entirely within an interactive video world model, avoiding real-world risk. The framework consists of three key components:

1. Chunk-wise Interactive World Model Rollouts

Instead of collecting data in the physical world, SafeDojo performs "model-in-the-loop" rollouts inside an interactive world model $p_\phi$.

• Architecture: The world model is instantiated as an action-conditioned video predictor (based on Wan 2.2) that generates future visual observations and latent dynamics.
• Process: Given an initial observation $o_0$ and language instruction $l$, the current VLA policy $\pi_\theta$ samples candidate trajectories. The world model recursively rolls out these trajectories to generate imagined future frames $\hat{o}_{i,k+1}$ and latent states $\hat{z}_{i,k+1}$ conditioned on action chunks.
• Robustness: To mitigate train-test mismatch and exposure error in long-horizon rollouts, the training includes static-video augmentation that occasionally freezes conditioning frames, making the model tolerant to imperfect intermediate predictions.

2. Decoupled Task Reward and Safety Cost Estimation

SafeDojo evaluates imagined trajectories using a dual-branch evaluator that estimates dense, per-step signals rather than collapsing them into a single scalar reward.

• Task Reward ($f_{task}$): A fine-tuned ResNet success classifier predicts per-step task progress from imagined future frames.
• Safety Cost ($f_{safe}$): A lightweight safety head estimates per-step safety costs (obstacle-contact probability) from predicted latent contexts and the proposed action chunk.
• Aggregation: These dense signals are aggregated into trajectory-level rewards ($r_i$) and costs ($c_i$). The safety cost specifically aggregates the top-$M$ high-risk steps to emphasize critical intervals.

3. Lagrangian-based Safety-Constrained GRPO

To optimize the policy under explicit safety constraints, SafeDojo employs a Group Relative Policy Optimization (GRPO) variant with a Lagrangian formulation.

• Advantage Calculation: Task rewards and safety costs are normalized separately within a group of sampled trajectories to compute relative advantages $\hat{A}^r_i$ and $\hat{A}^c_i$.
• Constrained Objective: The final advantage is computed as $\hat{A}^{safe}_i = \hat{A}^r_i - \eta \hat{A}^c_i$, where $\eta$ is a non-negative Lagrange multiplier.
• Dynamic Balancing: The multiplier $\eta$ is updated based on the batch-level constraint violation (average safety cost vs. allowable budget $d$). If the average cost exceeds the budget, $\eta$ increases to penalize unsafe trajectories more strongly; otherwise, it decreases to allow task-oriented improvement.
• Update Rule: The policy is updated using a clipped objective with the safety-constrained advantage, ensuring coordinated improvement of task success and safety without fixed penalty weights.

Key Contributions

1. World Model-based Reward-Cost Estimation: Development of an interactive video world model coupled with a dual-branch evaluator that decouples task progress and safety risk estimation from imagined dynamics, eliminating the need for hand-crafted safety functions.
2. Safety-Constrained RL: Introduction of a Lagrangian-based safety-constrained GRPO objective that explicitly enforces safety budgets while coordinating task success and collision avoidance, avoiding the sensitivity of fixed penalty weights.
3. Empirical Performance: Demonstration of superior performance on the SafeLIBERO benchmark and real-world Franka robot deployments, achieving the best aggregate task success, safe success, and execution efficiency among inference-time safety, model-free RL, and model-based RL baselines.

Experimental Results

The framework was evaluated on SafeLIBERO (a safety-critical benchmark built on Robosuite) and real-world Franka Emika Panda robots.

Simulation Results (SafeLIBERO)

• Safe Success Rate (SSR): SafeDojo achieved the highest average SSR on both Level I (53.25%) and Level II (49.62%). On Level I, this represents an 8.25 percentage-point improvement over the strongest baseline (SafeVLA at 45.00%).
• Execution Efficiency: SafeDojo achieved the lowest average Execution Time Steps (269.41 on Level I), outperforming the best model-based baseline (WMPO at 274.02), indicating that safety-aware optimization does not lead to overly conservative or inefficient policies.
• Generalization: Trained only on Level I demonstrations, SafeDojo maintained top performance on Level II (obstacles placed further along the motion path), demonstrating that safety signals learned via imagined rollouts transfer to unseen obstacle configurations.

Real-World Results

• Deployed on five real-world manipulation tasks, SafeDojo achieved the best average Task Success Rate (76.00%) and Safe Success Rate (70.00%).
• Qualitative analysis showed SafeDojo generated direct, collision-aware trajectories, whereas baselines either collided with obstacles, became overly conservative, or failed to recover from unsafe states.

Ablation Studies

• Component Importance: Removing the safety head dropped SSR to 38.0%; removing the task reward branch collapsed performance to near zero. Replacing constrained GRPO with unconstrained optimization significantly reduced safety.
• Input Sources: Using the current world-model state plus the proposed action chunk provided the most stable reward-cost cues compared to using future predictions alone.
• Lagrange Multiplier ($\eta$): An initial $\eta$ of 0.2 provided the best balance; higher values led to over-constrained, inefficient policies, while lower values failed to suppress risky actions.

Significance and Claims

The paper positions SafeDojo as a scalable and generalizable path toward safe embodied intelligence. By leveraging world models to perform "anticipatory rollouts" in imagination, SafeDojo addresses the fundamental tension between the need for policy exploration and the irreversibility of physical safety violations.

The authors claim that SafeDojo is the first framework to successfully apply model-based safe RL to VLA models, enabling:

1. Risk-Free Learning: Policy improvement occurs entirely within a simulated world model, eliminating the need for costly and dangerous real-world exploration.
2. Data-Driven Safety: Safety signals are learned from high-dimensional visual observations rather than relying on predefined geometric constraints or hand-crafted cost functions.
3. Balanced Optimization: The Lagrangian-based approach allows for explicit constraint satisfaction while maintaining high task performance and efficiency.

The results suggest that world model-based safe reinforcement learning is a viable paradigm for deploying VLA policies in open-world physical environments where safety is a prerequisite.