SafeBranch: Branch-Pair Safety Alignment for Embodied Agents
arXiv:2608.197292026-08-21
Teaching a robot AI to know exactly which past decision was unsafe, by rewinding the sim to that moment and showing both choices side by side
Robot-body AI agents driven by vision-language models often finish household tasks while breaking safety rules along the way, like leaving a stove on or touching a wet outlet. The researchers built SafeBranch, which rewinds the agent's own failed attempt back to the exact moment it went wrong and trains it by comparing the unsafe choice against a safe one it generates itself at that same moment. This roughly tripled to tenfold improved safe task completion rates, including about ten times more safe successes than the untrained baseline when facing unfamiliar objects.
What they did
The problem: just imitating safe trajectories teaches the agent safe-looking behavior without explaining when it applies, causing it to repeat safety actions out of context; contrasting whole safe vs. unsafe trajectories spreads the signal across many unrelated steps, causing the agent to avoid finishing tasks just to seem safe.
The method: an unsafe rollout is rolled back in the simulator to the exact safety-critical step that caused the violation; an external critic (GPT-4o) points out what went wrong, the agent itself generates a safer alternative action at that same step, then the critic's hint is removed and only the original unsafe action and the agent's own safe alternative are kept as a matched training pair, called a branch pair.
Training: applying an objective called BranchPO to these branch pairs lets the trained agent act safely at deployment time with no external critic or safety checker running alongside it.
Results: on IS-Bench, the safe-success rate rose from 0.031 (untrained) to 0.281; on an unseen-object variant it rose from 0.048 to 0.469 (about tenfold); on the separate SafetyALFRED simulator, hazard-detection accuracy rose from 0.274 to 0.438.
Efficiency: the data-construction pipeline produced about 5.2 times more usable training pairs than a natural sampling approach within the same time budget, with human review confirming quality improved at each filtering stage.
Figure 1: Branch-pair supervision for interactive safety. At a safety-critical step, structuring supervision as two branches that share the same context but differ in the actor’s action makes the safety-determining choice explicit. This branch-pair form isolates the step-level safety signal, allowing the actor to learn where and how to take the safe branch.
Table 1: Qualitative comparison of supervision signals on a hygiene task. Task: put a peach from the fridge onto a soiled plate. Safety requirement: wipe the plate before placing the peach and close the fridge after retrieval. SafeBranch trains with BranchPO on branch pairs, contrasting safe and unsafe actions at the same decision point. Orange marks stalled continuations; red marks unsafe continuations; green marks safety-relevant actions.
Training signal
Trajectory
Result
Imitation Supervision
open fridge → wipe plate → place peach on plate → close fridge → place peach on plate → place peach on plate → … (no Done)
! STALL Performs safe actions, but collapses into an action loop.
Trajectory-level Preference
open fridge → place peach on soiled plate → Done
× UNSAFE Reaches the goal through an unsafe shortcut.
Branch-pair Preference
open fridge → wipe plate → place peach on plate → close fridge → Done
√ SAFE Chooses the local safe action and completes the task.
Figure 2: Overview of SafeBranch. SafeBranch constructs same-state branch pairs by rolling unsafe rollouts back to the safety-critical anchor step and eliciting a repaired output with critic feedback. The feedback is removed before training, and BranchPO aligns the actor on these branch pairs for critic-free deployment.
Table 2: Main results on IS-Bench and the two OOD benchmarks constructed from it (ObjectShift injects distractors; TaskShift substitutes target objects). The first block runs the untrained actor with optional inference-time safety machinery; the second block trains the actor and deploys it critic-free. For each column, bold marks the best and underline the second-best.
In-Distribution
OOD-ObjectShift
OOD-TaskShift
Method
SR
SSR
SRec
SR
SSR
SRec
SR
SSR
SRec
Inference-time / untrained actor
Baseline
0.656
0.031
0.273
0.899
0.051
0.243
0.803
0.048
0.295
Self-Verification
0.607
0.071
0.333
0.693
0.053
0.244
0.651
0.082
0.295
Lookahead
0.219
0.000
0.094
0.244
0.061
0.054
0.133
0.044
0.067
Actor-trained / critic-free deployment
SFT-only
0.594
0.219
0.422
0.714
0.347
0.390
0.755
0.434
0.689
Trajectory DPO
0.656
0.000
0.256
0.613
0.118
0.245
0.748
0.075
0.253
+ success-matched
0.594
0.000
0.333
0.796
0.097
0.309
0.774
0.063
0.237
BranchPO (ours)
0.594
0.281
0.467
0.819
0.355
0.589
0.694
0.469
0.795
Figure 3: Analysis of branch pair construction. (a) Data generation efficiency. Under the same DFS rollout budget, natural best-of-N sampling and critic-guided rollback are evaluated with a fixed GPT-4o judge for usable same-anchor branch pairs; SafeBranch produces such pairs 5.2× faster. (b) Pair quality through filtering. Starting from 753 raw pairs, the SafeBranch pipeline applies judge filtering and cross-temperature deduplication to retain reliable final branch pairs, with human verification showing higher usable-pair rates across stages. (c) Downstream effect of filtering. Training BranchPO on each filtering stage shows that SR, SSR, and SRec improve most after the final filtering stage, indicating that pair quality rather than raw pair count carries the downstream safety signal.
Table 3: Cross-simulator transfer on SafetyALFRED under its native five-category taxonomy. Unsanitary† is pre-solved by the base VLM (∼35%); Fall/Trip Hazard yields no training pairs under the SafeBranch recipe. For each column, bold marks the best.
Method
Appliance Misuse
Property Damage
Unsanitary†
Spoilage
Fall/Trip Hazard
All
Baseline
0.048
0.034
0.630
0.079
0.000
0.274
Self-Verification
0.089
0.159
0.663
0.053
0.000
0.323
Lookahead
0.024
0.028
0.683
0.053
0.000
0.287
BranchPO (ours)
0.202
0.428
0.711
0.079
0.078
0.438
Box 17: Per-risk-type cumulative bar plot, showing for each hazard category how many safety conditions SafeBranch satisfies vs. violates. Three panels stacked horizontally (ID, OOD, All).
Table 4: Pair construction across preference-learning variants. Shared h: y+ and y− share the same anchor. Task success: both branches complete the task.
Method
Pair type
Shared h
Task success
SFT
Imitation
—
—
Trajectory DPO
Contrast
×
×
+ success-matched
Contrast
×
✓
BranchPO (ours)
Contrast
✓
✓
Table 12: SafeBranch data-construction funnel from the 753 extracted preference pairs (608 rollout episodes, 4 temperatures; BeforeBDDL + TermSafety triggers). “Kept” is the fraction kept from the previous stage.
Stage
Count
Kept
Extracted preference pairs
753
—
→ Judge-kept (quality)
675
89.6%
→ Final (after dedup, train pool)
475
70.4%
Table 13: Removing critic feedback (Ours) vs. retaining it (+FB) in the training prompt, evaluated critic-free on the ID split. The two data variants are byte-identical apart from the critic guidance block. SafeBranch denotes the staged SFT→BranchPO pipeline used in the body (Sec. 4); BranchPO-only drops the SFT warm-up.
Method
Variant
SR
SSR
SRec
SFT-only
Ours
0.594
0.219
0.422
+FB
0.714
0.000
0.270
BranchPO-only
Ours
0.656
0.250
0.488
+FB
0.600
0.133
0.409
SafeBranch
Ours
0.594
0.281
0.467
+FB
0.690
0.138
0.425
Table 14: Runtime full-critic (GPT-4o, one critic call per decision step) vs. SafeBranch / BranchPO (critic-free) across IS-Bench and our controlled OOD extensions. The cost column reports the upper-bound number of additional GPT-4o critic calls under a 30-step budget: (32+147+138)×30=9,510. SafeBranch adds no test-time critic calls.
ID
OOD-ObjectShift
OOD-TaskShift
Test-time cost
Method
SR
SSR
SRec
SR
SSR
SRec
SR
SSR
SRec
Extra GPT-4o calls
Full-critic (GPT-4o)
0.656
0.406
0.680
0.762
0.524
0.742
0.723
0.616
0.793
≤9,510
SafeBranch (ours)
0.594
0.281
0.467
0.819
0.355
0.589
0.694
0.469
0.795
0
Table 15: Per-checkpoint training dynamics on the development split (32 tasks, actor-only). Best checkpoint per method (by SSR) in bold; these are the checkpoints used in the main results.
Method
Step
SR
SSR
SFT-only
10
0.656
0.031
20
0.688
0.125
30
0.594
0.219
60
0.094
0.094
90
0.281
0.156
BranchPO-only
30
0.688
0.031
50
0.656
0.188
70
0.594
0.156
90
0.656
0.250
140
0.594
0.188
210
0.594
0.188
SafeBranch
30
0.594
0.281
(SFT→BranchPO)
50
0.581
0.161
70
0.552
0.138
90
0.633
0.100
110
0.517
0.138
130
0.433
0.100
150
0.310
0.034
Table 16: SafeBranch per-risk-type safety recall on ID and OOD splits. Safe is the number of satisfied safety conditions, and Total is the number of required safety conditions. SRec is computed as Safe / Total. The All columns aggregate ID and OOD. SafeBranch checkpoint = SFT→BranchPO at step 30 (see Table 15).
ID
OOD
All
Risk type
SRec
Safe
Total
SRec
Safe
Total
SRec
Safe
Total
Collision/Tripping Hazard
0.429
9
21
0.560
190
339
0.552
199
360
Fire Hazard
0.167
1
6
0.233
7
30
0.222
8
36
Food Contamination
0.143
1
7
0.341
14
41
0.312
15
48
Chemical Hazard
0.750
6
8
0.610
86
141
0.617
92
149
Water Spill Damage
1.000
1
1
0.000
0
12
0.077
1
13
Falling Object Hazard
0.400
2
5
0.784
29
37
0.738
31
42
Electrical Shock
0.000
0
2
0.367
11
30
0.344
11
32
All
0.400
20
50
0.535
337
630
0.525
357
680
Why it matters
As embodied AI agents move into real homes and workplaces, completing a task safely matters as much as completing it at all, and this work offers a concrete recipe for training that safety directly into the agent rather than bolting on an external checker. Because the trained agent needs no critic running at deployment, it points toward a more practical, lower-cost path to safer robots in real-world use.
Terms in this paper
vision-language-model (VLM) agent · an AI system that combines visual perception with language understanding to decide actions
safety-critical step · a moment in a task where the agent's choice determines whether the outcome ends up safe or unsafe
branch pair · a matched pair of a safe action and an unsafe action recorded at the exact same decision moment
environment rollback · resetting a simulation back to an earlier point in time
BranchPO · a training method designed specifically to learn from branch pairs by preferring the safe action
critic-free deployment · running the trained agent without any external safety-checking module active during use
Original abstract (English)
Vision-language-model-based embodied agents can complete instructed tasks but often violate safety constraints in the process, a problem recently framed as interactive safety. Training such agents to act safely is difficult, since safety and task success are distinct objectives, and safety arises only at a small number of safety-critical steps within a trajectory. Standard supervision is insufficient: imitating safe trajectories teaches behavior without explaining why it is safe, and contrasting arbitrary safe and unsafe trajectories mixes the safety signal with unrelated differences. We propose SafeBranch, a framework that aligns an embodied actor on safety through branch pairs constructed from the actor's own unsafe rollouts via environment rollback. SafeBranch rolls each unsafe rollout back to the safety-critical step that caused the violation, queries the actor for a safe alternative, and pairs the original action with the alternative so that the two branches differ only at that step. The trained actor acts safely at deployment with no critic in the loop. On IS-Bench, SafetyALFRED, and out-of-distribution variants with unseen tasks and objects, it handles safety reliably without sacrificing task success, achieving roughly ten times more safe successes than the untrained baseline on the unseen-object variant.
