ElevateHead: Decide Once, Then Commit: When the Action Head Matters for Multimodal Imitation

Tsung-Lung Yang

cs.ROOct 25, 2026

ElevateHead: Decide Once, Then Commit: When the Action Head Matters for Multimodal Imitation

Authors: Tsung-Lung Yang

Organizations: National University

Abstract

Robots that learn from people often see several valid ways to do the same task: pass an obstacle on the left or on the right, open the microwave or move the kettle first. A policy that re-decides at every step can waver between these ways mid-task. We ask when the action head (the part of a policy that turns the state into actions) matters, and how it should make this choice. On a controlled task, heads that sample one way succeed far more often than a head that averages the ways when the demonstrations contain several ways and the state gives no hint which to take (+0.72 success with two ways); this trend was not reproduced on the other tasks we ran. We then present ElevateHead, a head that decides once at a learned decision point, commits to that choice for the episode, and exposes the choice as an input a user can set. We designed it by diagnosis: each version was frozen and tested on tasks it had not seen, and the next version changed only what a measured failure pointed to; predictions were registered before each run. The first four versions each fell below BeT on at least one held-out task. The fifth completed all four subtasks of the held-out Franka Kitchen task in 45.2% of episodes against 29.3% for BeT, the strongest baseline across our tasks (+0.160, 95% CI [+0.129, +0.188]); when a mode was forced, the first completed subtask matched it in 95% of successful episodes. Diagnosis also located where commitment breaks: varied start positions (repaired on the controlled task: commitment 0.22 to 0.81, success 0.74 to 0.91) and varied object placement (partly repaired: in an exploratory variant, commitment 0.03 to 0.59 on aligning with no detectable change in success). On held-out block stacking, version 4 was comparable to BeT and did not exceed it. A real-robot test with predictions fixed in advance is planned for a follow-up version of this paper.

Figures & tables

Panel of Figure 1 — Fig. 1: The idea in one picture. Dotted: the two demonstrated ways around the obstacle (gray). (a) A head that samples its choice again at every step can switch ways mid-route. (b) ElevateHead chooses once, at a learned decision point (diamond), and keeps the choice. (c) Because the choice is an explicit input, a user can set it.

Question	Answer	Key evidence
When does the head matter? (§ III )	With several ways and no hint in the state (controlled task); not reproduced on other tasks	gap $+0.72$ at two ways, $-0.04$ when the average path works
Does deciding once help on an unseen task? (§ VI )	Yes on Kitchen; not shown on stacking	$p_{4}$ 0.452 vs. 0.293 for BeT, $+0.160$ [+0.129, +0.188]
Does the mode input set the behavior?	On Kitchen, mostly (descriptive)	first subtask matches the forced mode in 95.4% of successes
Where does commitment break?	Varied starts (repaired on one task), few demonstrations	commitment 1.00 to 0.22; gate ( 1 ) restores 0.81
Varied object placement (toward clutter)?	Partly repaired; success unchanged within uncertainty (exploratory)	commitment 0.03 to 0.59 (aligning), 0.32 to 0.64 (stacking)
Which baseline is strongest? (Table III )	BeT (or tied) on every task, among eight heads	Kitchen $p_{4}$ : BeT 0.293, flow 0.111, diffusion 0.000

TABLE I: Main findings at a glance.

Panel of Figure 3 — Fig. 2: When the head matters (controlled task, 5 seeds, 95% CI). $\Delta$ : success of the five sampling heads minus that of an averaging (L1) head. “avg. works”: two ways whose average path also reaches the goal; “hint”: random start offset that tells which way to take; open loop (OL): no re-planning. The gap appears with several ways and disappears when averaging works or a hint exists.

Panel of Figure 4 — Fig. 3: ElevateHead in four ideas. Offline (gray): find the ways and the decision points in unlabeled demonstrations. Online (blue): at the first decision point a gate checks whether the choice is genuinely open; if so, one mode is drawn and kept for the rest of the episode; if not, the head acts without a mode.

Panel of Figure 5 — Fig. 4: The design record: every held-out result against BeT, by version (95% CI). Red: below BeT; gray: within uncertainty; blue: superior (registered). Gray text at top: the change each diagnosis led to in the next version (Table II ). Points of different versions are on different tasks, so this is a record, not a learning curve. Stacking: registered two-box outcome (three-box success below 0.10 for every head). Pushing was added as a second v2 test after one sorting-4 BeT cell signalled a floor risk, before any ElevateHead sorting-4 result.

Ver.	Diagnosis (registered experiments)	Change in the next version
v1	On sorting-2 the proposal was confident, yet the head committed and timed out; re-selection, its threshold and the mode choice were ruled out as causes	v2: commit only when the proposal is uncertain; train a no-mode branch (mode dropout 0.25)
v2	On pushing the regression loss limited it, not commitment (plain L1 0.10 vs. L2 0.68)	v3: L2 action loss
v3	On aligning the gate never committed (certain proposal), so success rested on an action network that trailed BeT	v4: BeT-style action network
v4	On development tasks: commitment helped on avoiding; re-selecting a mode when the arm “stalls” hurt (turning it off: $+0.062$ on pushing); the gate closed as clustering found more modes	v5: no stall re-selection; gate ( 1 ); training unchanged

TABLE II: From diagnosis to design: what each failure revealed and what the next version changed.

Panel of Figure 7 — Fig. 5: Held-out Franka Kitchen (10 seeds $\times$ 500 episodes; same network, inputs and training for every head; state: hand position and 30 robot and object coordinates). (a) Share of episodes completing at least $n$ of the four subtasks; both ElevateHead versions commit in every episode. (b) Share of each first subtask: version 5 stays closer to the demonstrated mix than BeT.

Task	L1	L2	MIP	AR	VQ-BeT	Flow	Diff.	BeT	ElevateHead
avoiding	.667	.000	.000	.807	.610	.490	.300	.813	.967 †
sorting-2	.667	.643	.743	–	–	–	–	.776	.619 (v1)
sorting-4	.043	.034	.041	.083	–	–	–	.166	.054 (v2)
pushing	.102	.683	.154	.602	–	–	–	.787	.764 †
aligning	.001	.134	.020	.246	.131	.247	.032	.342	.402 †
stacking a	.140	.017	.000	.107	.130	.016	.000	.139	.154 (v4)

TABLE III: Every head on every task (success, seed means). Shading: success relative to the best baseline in the row; bold: best baseline. BeT is the best baseline, or within noise of it, on every task, so every held-out test compares with BeT.

Panel of Figure 9 — Fig. 6: Where commitment breaks and how far diagnosis repaired it. (a,b) Varied starts, controlled task (two ways, 400 demonstrations, 5 seeds; development, gate change adopted on these data). (c) Varied object placement, a step toward cluttered scenes; success did not change detectably, though non-inferiority is not shown. “v4 gate”: threshold 0.5 of versions 2–4.

Id	Task	Test	Value [lo, hi]	Verdict
P1	C-S	slope of $\Delta$ on $\log_{2}k$	$+0.333$ [+0.134, +0.437]	supported
P2	C-S $k{=}1$	$\|\Delta\|\leq\delta$	$-0.032$ [-0.039, -0.025]	inconclusive
P5	C-S	$\Delta(k{=}2)-\Delta(k{=}2\mathrm{F})\geq\delta$	$+0.768$ [+0.363, +0.981]	supported
HM-1	C-S	$\Delta$ falls with start jitter (0, 0.005, 0.02)	$+0.724>+0.126>-0.056$	supported
HM-3	C-S	$\Delta$ (open loop) $>\Delta$ (closed loop) $+\delta$ , jitter 0.02	$+0.143$ [+0.097, +0.189]	supported
HM-4	C-S	blocked route: deterministic $\leq 0.1$ , noise heads $\geq 0.3$	0.00 vs. 0.42	supported

TABLE IV: Registered tests used in the paper. v5-X: version 5 with a regressor retrained with token dropout 0.5. v5-S: proposal trained on the mode mix of the 10 nearest demonstrations. Oracle: version 5 re-deciding the mode whenever the environment reports a completed subtask.

Task (role for v5)	v5	BeT	v5 $-$ BeT
C-S, $J{=}0$ (development)	1.000	1.000	$+0.000$
avoiding (development)	0.967	0.813	$+0.153$ [+0.097, +0.217]
pushing (development)	0.764	0.787	$-0.022$ [-0.076, +0.029]
aligning (development)	0.402	0.342	$+0.060$ [+0.016, +0.104]
stacking, 3 boxes (held-out)	0.044	0.043	–
Kitchen $p_{4}$ (held-out)	0.452	0.293	$+0.160$ [+0.129, +0.188]

TABLE V: Version 5 against BeT on every task it was run on.

C-S dataset	commit v4	commit v5	succ. v4	succ. v5	BeT	steer (seeds)
$J{=}0$	1.000	1.000	1.000	1.000	1.000	1.000 (5/5)
$J{=}0.005$	1.000	1.000	1.000	1.000	1.000	1.000 (5/5)
$J{=}0.02$	1.000	1.000	1.000	1.000	1.000	0.989 (5/5)
$J{=}0.05$	0.620	1.000	1.000	1.000	1.000	0.953 (5/5)
$J{=}0.1$	0.342	0.950	1.000	1.000	1.000	0.931 (5/5)
$J{=}0.2$	0.216	0.812	0.744	0.912	0.740	0.921 (5/5)

TABLE VI: C-S dose study at

e{=}10

(5 seeds). Steering columns refer to version 4 (forcing bypasses the gate, so version 5 steers identically).

Head (one CPU thread, batch 1)	stacking	aligning
BeT	0.495	0.524
ElevateHead (v4 networks)	0.659	0.631
VQ-BeT	2.185	1.615
Flow matching (10 steps)	2.119	1.892
Diffusion (10 DDIM steps)	2.604	2.331

TABLE VII: Median time per chunk (ms), Intel Core i7-13700HX, float32; one sample call, excluding observation reading and environment steps.

Panel of Figure 14 — Fig. 7: Version 5 against BeT on every task where both ran (95% CI; dashed: $\pm\delta$ ). Development tasks were used for design, so only Kitchen is a held-out test. Stacking is omitted because version 5 was not compared with BeT there by a registered test.

Robot Policy Learning

Robot policy learning develops controllers that map sensory observations, task goals, and interaction history to actions for embodied robots.

See all 369 papers

Demo-JEPA: Joint-Embedding Predictive Architecture for One-shot Cross-Embodiment Imitation
May 20, 2026Jingyang He, Guangrun Li, Jieyu Zhang +3Robot Skill Learning World Models for Robotics
State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University · Beijing Innovation Center of Humanoid Robotics · University of Washington
Conditional Trajectory Peaks: Single-Pass Multimodal Policies over Action Chunks
Oct 5, 2026Di Wu, Rongtian Shen, Ping Liu +4Efficient Multimodal Inference Robot Policy Learning
Magiclab Robotics Technology Co., Ltd., China · Southeast University, Nanjing, China
Commit While Futures Agree: Consequence-Aware Adaptive Action Chunking for Robot Manipulation
Oct 6, 2026Yuyan Li, Yujia Wang, Yusong Huang +8Robotic Manipulation
The Hong Kong University of Science and Technology (Guangzhou) · ANU Intelligent · Victoria University of Wellington +2

May 20, 2026J/K move · Enter open · S save

Demo-JEPA: Joint-Embedding Predictive Architecture for One-shot Cross-Embodiment Imitation

Jingyang He, Guangrun Li, Jieyu Zhang, Chengkai Hou+2

State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University · Beijing Innovation Center of Humanoid Robotics · University of Washington

Robot Skill Learning World Models for Robotics Joint-Embedding Predictive Architecture Generalization in Robotic Manipulation Cross-Embodiment Robot Learning Imitation Learning+1

Open paper AskPDF arXiv

Robotic imitation learning is often treated as reproducing demonstrated actions, but actions are inherently embodiment-specific. When demonstrations come from humans or robots with different morphology, kinematics, or action spaces, this action-centric view requires shared action spaces, heuristic retargeting, or large-scale multi-embodiment co-training. We instead view demonstrations as implicit specifications of future goals: the target agent should infer what state the demonstrator is trying to realize, rather than how the demonstrator executes it. We propose Demo-JEPA, a cross-embodiment imitation framework that decouples demonstration intent from embodiment-specific execution. Built on a JEPA-based world model, Demo-JEPA translates source visual demonstrations into target-compatible future latent trajectories in a shared predictive representation space. The target agent then uses these latent trajectories as subgoals and realizes them through planning under its own learned forward dynamics. Because Demo-JEPA avoids action-level correspondence and requires only visual demonstrations plus the target agent's own interaction experience, it supports flexible imitation across heterogeneous embodiments. Experiments on RLBench and real-world manipulation tasks show that Demo-JEPA matches specialized in-domain planners and generalizes to unseen tasks and embodiment configurations where prior methods fail.

Abstract

Robot Policy Learning

Explore similar work

Demo-JEPA: Joint-Embedding Predictive Architecture for One-shot Cross-Embodiment Imitation

Conditional Trajectory Peaks: Single-Pass Multimodal Policies over Action Chunks

Commit While Futures Agree: Consequence-Aware Adaptive Action Chunking for Robot Manipulation