SafeBranch: Branch-Pair Safety Alignment for Embodied Agents

Hyunse Lee; Jiwoo Jeong; Haneul Lee; Kyochul Jang; Youngjae Yu; Woojin Lee

cs.AIAug 20, 2026

SafeBranch: Branch-Pair Safety Alignment for Embodied Agents

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

Organizations: Dongguk University · Seoul National University

Abstract

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 improves safe task completion, achieving roughly ten times more safe successes than the untrained baseline on the unseen-object variant. Project page: https://github.com/98hyunse/SafeBranch.

Figures & tables

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.

Training signal	Trajectory	Result
Imitation Supervision	open fridge $\rightarrow$ wipe plate $\rightarrow$ place peach on plate $\rightarrow$ close fridge $\rightarrow$ place peach on plate $\rightarrow$ place peach on plate $\rightarrow$ … (no Done )	! STALL Performs safe actions, but collapses into an action loop.
Trajectory-level Preference	open fridge $\rightarrow$ place peach on soiled plate $\rightarrow$ Done	$\times$ UNSAFE Reaches the goal through an unsafe shortcut.
Branch-pair Preference	open fridge $\rightarrow$ wipe plate $\rightarrow$ place peach on plate $\rightarrow$ close fridge $\rightarrow$ Done	$\surd$ SAFE Chooses the local safe action and completes the task.

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.

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.

	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

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.

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

Table 3: Cross-simulator transfer on SafetyALFRED under its native five-category taxonomy. Unsanitary † is pre-solved by the base VLM (

\sim

35%); Fall/Trip Hazard yields no training pairs under the SafeBranch recipe. For each column, bold marks the best.

Method	Pair type	Shared $h$	Task success
SFT	Imitation	—	—
Trajectory DPO	Contrast	$\times$	$\times$
+ success-matched	Contrast	$\times$	✓
BranchPO (ours)	Contrast	✓	✓

Table 4: Pair construction across preference-learning variants. Shared

h

y^{+}

and

y^{-}

share the same anchor. Task success : both branches complete the task.

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

\times

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 manual review 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.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Stage	Count
Safety-rule triggers	853
$\rightarrow$ Extracted preference pairs	753
$\rightarrow$ After duplicate merging	743
$\rightarrow$ LLM-judge accepted	675
$\rightarrow$ Author-reviewed pool	633
$\rightarrow$ Training-task pairs	475

Appendix

Table 12: SafeBranch data-construction funnel over 608 rollout episodes and four sampling temperatures.

Method	Variant	SR	SSR	SRec
SFT-only	Ours	0.594	0.219	0.422
SFT-only	+FB	0.714	0.000	0.270
BranchPO-only	Ours	0.656	0.250	0.488
BranchPO-only	+FB	0.600	0.133	0.409
SafeBranch	Ours	0.594	0.281	0.467
SafeBranch	+FB	0.690	0.138	0.425

Appendix

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

\rightarrow

BranchPO pipeline used in the body (Sec. 4 ); BranchPO-only drops the SFT warm-up.

Method	ID	OOD-ObjectShift	OOD-TaskShift
APO-zero w/o SFT	0.656 / 0.219 / 0.511	0.796 / 0.306 / 0.562	0.870 / 0.391 / 0.632
APO-zero w/ SFT	0.562 / 0.125 / 0.405	0.728 / 0.211 / 0.437	0.891 / 0.239 / 0.469
SafeBranch (SFT $\rightarrow$ BranchPO)	0.594 / 0.281 / 0.467	0.819 / 0.355 / 0.589	0.694 / 0.469 / 0.795

Appendix

Table 14: APO-zero comparison in SR / SSR / SRec. All methods run actor-only. TaskShift results use the same 138-task evaluation manifest as Table 2 .

	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	$\leq$ 9,510
SafeBranch (ours)	0.594	0.281	0.467	0.819	0.355	0.589	0.694	0.469	0.795	0

Appendix

Table 15: 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)\times 30=9{,}510

. SafeBranch adds no test-time critic calls.

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

Appendix

Table 16: Per-checkpoint training dynamics on the 32-task evaluation split (32 tasks, actor-only). Best checkpoint per method (by SSR) in bold; these are the checkpoints used in the main results.

	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

Appendix

Table 17: 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

\rightarrow

BranchPO at step 30 (see Table 16 ).

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

AI Agent Safety

AI agent safety studies how to prevent and detect harmful actions by autonomous systems that use tools, memory, and multi-step plans.

See all 326 papers

ActSafeGuard: Differentiable and Training-Aligned Constraint Enforcement for Flow-Matching Policies
Sep 10, 2026Jianming Ma, Rongjun Jin, Xiaxi Si +3Robotic Control Robot Foundation Models
1Shanghai Jiao Tong University · 2Shanghai Innovation Institute
NEXUS: Learning Safety Constraints and Action Models for Safe and Robust Embodied Planning
May 10, 2026Tiehan Cui, Peipei Liu, Yanxu Mao +3Continual Learning for LLM Agents Large Language Model-Based Robot Planning
School of Artificial Intelligence and Automation, Huazhong University of Science and Technology, China · School of Software, Henan University, China · Institute of Information Engineering, Chinese Academy of Sciences, China +2
RT-Safe: Benchmarking Agent Safety in Real-Time Embodied Environment
Oct 7, 2026Tianruo Rose Xu, Jiawei Ren, Yichi Yang +2AI Agent Safety Safe Robot Navigation
Cornell University · University of California, San Diego

Sep 10, 2026J/K move · Enter open · S save

ActSafeGuard: Differentiable and Training-Aligned Constraint Enforcement for Flow-Matching Policies

Jianming Ma, Rongjun Jin, Xiaxi Si, Yang Zhang+2

1Shanghai Jiao Tong University · 2Shanghai Innovation Institute

Robotic Control Robot Foundation Models Robot Safety Vision-Language-Action Models

Open paper AskPDF arXiv

Vision-Language-Action (VLA) and World-Action Models (WAMs) have demonstrated strong capabilities in general-purpose robotic manipulation, yet their generated actions may violate hard physical constraints and therefore be unsafe or infeasible for deployment. Existing safety approaches either optimize statistical safety objectives without deterministic per-step guarantees or correct unsafe actions only during inference, creating a mismatch between policy training and execution. We introduce ActSafeGuard, a differentiable and training-aligned safeguard layer for flow-matching based policies. ActSafeGuard integrates hard action feasibility into policy learning, not merely treating safety as an inference-time external component. Through an analytical ray-scaling operator design, ActSafeGuard enables boundary-aware gradients to guide the model to naturally learn constrained manifolds. Extensive experiments on multiple standard foundation backbones (

π_{0.5}

and Fast-WAM) across various tasks demonstrate that ActSafeGuard consistently achieves a

100\%

step safety rate while fully preserving or even boosting task success rates, providing a scalable and minimally invasive solution for safe embodied AI deployment.

Abstract

AI Agent Safety

Explore similar work

ActSafeGuard: Differentiable and Training-Aligned Constraint Enforcement for Flow-Matching Policies

NEXUS: Learning Safety Constraints and Action Models for Safe and Robust Embodied Planning

RT-Safe: Benchmarking Agent Safety in Real-Time Embodied Environment