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

May 20, 2026

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

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.

Jingyang He, Guangrun Li, Jieyu Zhang +3

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

Ask PDF

Oct 5, 2026

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

Multimodal imitation learning requires diverse executable futures under the same observation and consistent behavior across replanning cycles. We present Conditional Trajectory Peaks (CTP), a single-pass policy framework that jointly predicts complete action-chunk candidates, probability masses, and trajectory scales. Distribution-Aware Peak Specialization (DAPS) specializes trajectory peaks using trajectory-level posterior responsibilities and mass- and scale-modulated overlap constraints. Evidence-Gated Trajectory Belief Transport (ETBT) maintains cross-chunk consistency through geometric correspondence between exchangeable candidate sets, while allowing current policy evidence to override historical constraints. CTP achieves a coverage score of 91.40% on Push-T; success rates of 100.0%, 79.72%, and 84.44% on D3IL Avoiding, Aligning, and Sorting-2, respectively. On LIBERO, CTP achieves an average success rate of 97.25%. In real-world dual-arm experiments, CTP preserves both placement modes in a two-plate task, succeeding in all 50 trials. On bottle uprighting and pen placement into a holder, it maintains success rates comparable to

π_{0.5}

while reducing policy inference latency from 218.24 ms to 75.80 ms. These results demonstrate that single-pass trajectory modeling can combine multimodal behavior, closed-loop consistency, and efficient inference.

Di Wu, Rongtian Shen, Ping Liu +4

Magiclab Robotics Technology Co., Ltd., China · Southeast University, Nanjing, China

Efficient Multimodal Inference Robot Policy Learning Imitation Learning Robotic Manipulation Action Chunking Multimodal Trajectory Prediction

Ask PDF

Oct 6, 2026

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

Action-chunking policies predict multi-step control sequences, but a fundamental question remains: how much of a predicted action chunk should be committed before replanning? Existing systems typically execute a fixed-length prefix, implicitly assuming that the same execution horizon remains trustworthy across states. Some adaptive methods estimate this horizon from the similarity or stability of predicted actions. However, different actions may lead to the same successful outcome, whereas similar actions can produce different futures, suggesting that commitment should be determined by agreement among imagined futures rather than by similarity in action space. To this end, we propose Consequence-Aware Adaptive Action Chunking (CA

^3

C), an inference-time framework built on a simple principle: commit while imagined futures agree, and replan when they diverge. Without modifying or retraining the base policy, CA

^3

C uses an action-conditioned world model to imagine the future consequences of multiple candidate action chunks under the same sampling noise. Using these imagined consequences, we formulate execution-horizon estimation as a Bayesian change-point inference problem and select the execution candidate through future consensus. Across multiple simulation benchmarks and real-world robot manipulation tasks, CA

^3

C consistently improves diverse action-chunking policies, achieving up to a 71.8% relative reduction in failure rate over the corresponding base policies.

Yuyan Li, Yujia Wang, Yusong Huang +8

The Hong Kong University of Science and Technology (Guangzhou) · ANU Intelligent · Victoria University of Wellington +2

Robotic Manipulation

Ask PDF

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