🤖 AI

SafeBranch: Branch-Pair Safety Alignment for Embodied Agents

SafeBranch is a framework that enhances the safety of vision-language-model-based embodied agents by training them on branch pairs derived from their own unsafe rollouts, enabling them to reliably handle safety-critical steps without compromising task success or requiring a critic during deployment.

Original authors: Hyunse Lee, Jiwoo Jeong, Haneul Lee, Kyochul Jang, Youngjae Yu, Woojin Lee

Published 2026-08-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Hyunse Lee, Jiwoo Jeong, Haneul Lee, Kyochul Jang, Youngjae Yu, Woojin Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Robots that can see, understand, and move through our homes are no longer just a dream of science fiction; they are becoming a reality powered by advanced computer models that process both images and language. These systems can follow instructions like "put the peach on the plate" and navigate a kitchen to find the fruit and the dish. However, there is a critical gap between simply finishing a task and doing so safely. A robot might successfully place a peach on a plate, but if it leaves the refrigerator door wide open or places the fruit on a dirty surface, it has failed the most important test of all: protecting the home and its inhabitants. This specific challenge, known as interactive safety, arises because the robot's own actions change the environment, creating new dangers that did not exist before. If a robot leaves a stove burner on after cooking or touches an electrical outlet with wet hands, the hazard is a direct result of its choices. The core difficulty lies in the fact that safety is not a constant requirement for every single move a robot makes; rather, safety is decided at a few critical moments where the robot must choose between a safe path and a dangerous shortcut.

For a long time, researchers tried to solve this by adding a separate safety guard to watch the robot at every step, like a human supervisor who constantly interrupts to say "stop" or "fix that." While this works, it is slow, expensive, and leaves the robot itself unchanged and dependent on the guard. A more difficult but promising approach is to teach the robot to be safe on its own, but standard teaching methods have struggled. Simply showing a robot a video of a safe task does not explain why a specific step was safe, and comparing a safe path to a completely different unsafe path often confuses the robot because the two paths differ in too many ways. The researchers behind a new framework called SafeBranch realized that to truly learn safety, a robot needs to see two versions of the exact same situation: one where it makes a safe choice and one where it makes a dangerous choice, with everything else remaining identical.

To achieve this, the team developed a clever method to create these specific teaching moments from the robot's own mistakes. Instead of waiting for a robot to fail and then hoping to find a safe alternative, they let the robot attempt a task and, when it makes a dangerous move, they rewind the simulation back to the exact moment the mistake happened. At that precise point, they ask the robot to try again, this time with a hint about the danger it just created. The robot then generates a safe action to replace the dangerous one. By pairing the original unsafe action with this new safe action, and ensuring both start from the exact same state, the researchers created a "branch pair." This pair isolates the single decision that matters, showing the robot that in this specific context, one choice leads to safety and the other to a hazard. They then used a training process to teach the robot to consistently choose the safe branch, effectively internalizing the safety lessons so that the external guard is no longer needed.

The results of this approach were striking when tested on a variety of household tasks involving unseen objects and new scenarios. In standard tests, robots that had not been trained with this method succeeded in completing tasks safely only about 3 percent of the time. After training with the SafeBranch method, the robots completed tasks safely nearly 28 percent of the time, a tenfold improvement. Even more importantly, when the robots faced completely new objects they had never seen before, the untrained robots were safe only 5 percent of the time, while the SafeBranch-trained robots achieved safety in nearly 47 percent of attempts. Crucially, the robots did not become overly cautious or refuse to do the work; they learned to take the necessary risks to finish the job while avoiding the specific steps that caused harm. Unlike other methods that require a separate safety module to run in the background during every step, the SafeBranch-trained robots operate independently, carrying the safety knowledge within their own decision-making process.

This work demonstrates that safety in robots is not a matter of adding more rules or external checks, but of teaching the agent to recognize the precise moment where a choice matters. By constructing training data that contrasts safe and unsafe actions in the exact same context, the researchers showed that robots can learn to navigate the complex, changing world of a human home without constant supervision. The method proved effective across different types of hazards, from electrical shocks and fire risks to food contamination and tripping hazards, suggesting a robust path forward for deploying helpful robots that can be trusted to act safely on their own.

Technical Summary: SafeBranch

Problem Statement

Vision-language model (VLM)-based embodied agents are capable of executing multi-step tasks in interactive environments but frequently violate safety constraints during execution. This issue, termed interactive safety, arises because the agent's own actions can create emergent hazards (e.g., leaving a stove on, touching an outlet with wet hands). Unlike static safety constraints, interactive safety hazards are concentrated at a sparse set of safety-critical steps within a trajectory. At these specific decision points, the agent's choice determines whether the trajectory branches toward a safe or unsafe outcome.

Existing approaches face two primary limitations:

Inference-time Guardrails: Methods that use external safety checkers or search-based planners to block or revise actions at every step add deployment overhead and leave the underlying actor unchanged.
Standard Supervision:
- Imitation (SFT): Training on safe trajectories teaches what to do but fails to explain why an alternative was unsafe, leading to context-less application of safe behaviors.
- Trajectory-level Preference (DPO): Contrasting entire safe and unsafe trajectories mixes the safety signal with unrelated differences across many steps, failing to isolate the specific decision point where safety was determined.

The core challenge is to provide supervision that explicitly contrasts safe and unsafe actions at the same safety-critical step (same context, different actions) to teach the agent the step-level safety signal.

Methodology: SafeBranch

The authors propose SafeBranch, a framework that constructs branch pairs from the actor's own unsafe rollouts and aligns the actor using Branch Preference Optimization (BranchPO).

1. Branch Pair Construction via Environment Rollback

Since natural interactive environments do not provide two trajectories that diverge at a single step, SafeBranch synthesizes them:

Rollback: When a safety violation occurs in an agent rollout, a safety critic (an external LLM) identifies the violated constraint and the specific step ( $h_k$ ) where the decision was made.
Repair: The environment is rolled back to $h_k$ . The actor is then queried for a safe alternative ( $y^+$ ) conditioned on the original context $h_k$ plus the critic's feedback ( $f_k$ ).
Hindsight Relabeling: To ensure the actor learns safety from the context rather than the feedback cue, the critic's feedback ( $f_k$ ) is removed from the training prompt. The resulting branch pair consists of the original unsafe action ( $y^-$ ) and the repaired safe action ( $y^+$ ), both conditioned on the identical context $h_k$ .
Filtering: A judge LLM filters pairs to ensure the safe action is justified by the context, executable, and resolves the constraint.

2. Branch Preference Optimization (BranchPO)

The framework trains the actor using a preference learning objective adapted for step-level alignment. Given a dataset of branch pairs $D_{branch}$ , the objective maximizes the log-probability margin between the safe action and the unsafe action at the same anchor step:
$\mathcal{L}_{BranchPO} = -\mathbb{E}_{D_{branch}} \left[ \log \sigma(\beta(r^+_\theta - r^-_\theta)) \right]$
where $r_\theta(h, y) = \log \pi_\theta(y|h) - \log \pi_{ref}(y|h)$ . This internalizes the critic's safety judgments into the actor itself.

3. Deployment

Unlike inference-time methods, the trained actor deploys critic-free. It handles safety-critical situations autonomously without external modules, rollback mechanisms, or feedback loops during runtime.

Key Contributions

Formulation of Interactive Safety as a Step-Level Problem: The paper identifies that safety alignment requires isolating the decision at safety-critical steps, rather than optimizing over entire trajectories.
SafeBranch Framework: A novel method to construct same-state branch pairs by rolling back unsafe rollouts and eliciting safe alternatives, solving the data scarcity problem for step-level safety supervision.
BranchPO Objective: A preference optimization objective specifically designed to train actors on these branch pairs, internalizing safety constraints.
Critic-Free Deployment: The method shifts the cost of the safety critic from runtime (inference) to training, resulting in a standalone actor that acts safely without external intervention.

Experimental Results

The method was evaluated on IS-Bench, SafetyALFRED, and out-of-distribution (OOD) variants involving unseen tasks and objects.

Performance Gains:
- On IS-Bench, SafeBranch increased the Safe Success Rate (SSR) from 0.031 (untrained baseline) to 0.281.
- On the OOD-ObjectShift variant (unseen objects), SSR improved from 0.048 to 0.355.
- On the OOD-TaskShift variant (unseen tasks), SSR improved from 0.048 to 0.469.
- On SafetyALFRED, hazard accuracy increased from 0.274 to 0.438.
Comparison to Baselines:
- SafeBranch outperformed inference-time baselines (Self-Verification, Lookahead) which failed to generalize to OOD settings.
- It significantly outperformed standard preference learning variants (Trajectory DPO, SFT-only), which struggled to concentrate the safety signal at the correct decision points.
Efficiency: The branch-pair construction pipeline generates usable pairs approximately 5.2× faster than natural best-of-N sampling under matched compute budgets.
Ablation: Removing the hindsight relabeling (retaining critic feedback in training) caused the model to fail at deployment when the feedback was absent, confirming the necessity of internalizing the safety signal.

Significance and Claims

The paper claims that SafeBranch offers a data-efficient route to internalizing interactive safety. By framing safety as a step-level branching problem and constructing specific branch pairs, the method enables embodied agents to learn where and how to act safely without sacrificing task success.

The authors emphasize that this approach moves safety alignment from an external, runtime-dependent process to an internal capability of the actor. The results suggest that training on same-state branch pairs is essential for learning the step-level safety signal, a capability that standard trajectory-level supervision fails to provide. The method achieves state-of-the-art safety performance across in-distribution, out-of-distribution, and cross-simulator settings while eliminating the need for external safety modules during deployment.