Skip to content

Improving Vision-Language-Action Model Fine-Tuning with Structured Stage and Keyframe Supervision

🕒 Published (v1): 2026-06-25 09:38 UTC · Source: Arxiv · link

Why this paper was selected

Structured stage+keyframe supervision for VLA fine-tuning; practical reproducible recipe

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

StaKe is a plug-in auxiliary supervision framework for fine-tuning Vision-Language-Action (VLA) models that automatically extracts two complementary training signals—stage labels and keyframe targets—from demonstration gripper states without manual annotation. It adds lightweight auxiliary heads (stage classifier + keyframe predictor) on top of a frozen inference-time policy, improving success rates by 14% (bimanual simulation) and 56% (real-robot single-arm) relative to the base π0.5 policy. Gains scale with task horizon length.

Problem

Standard VLA fine-tuning applies a uniform continuous action loss across all timesteps, ignoring the inherent stage structure of manipulation trajectories (free-space motion vs. contact-constrained skill execution). Gripper-state transitions are critical failure points, yet the policy receives no explicit supervision about which stage it is in or what configuration it must reach at the next transition. This causes failures to cluster around stage boundaries.

Method

StaKe augments the π0.5 VLA architecture (ViT encoder + PaLM-style backbone + flow-matching action expert) with two learnable query tokens \(e_{ss}\) and \(e_{kf}\) appended to the backbone's input sequence. These tokens produce contextualized representations fed into two lightweight MLP auxiliary heads:

Stage Supervision (SS): Stage labels \(s_t \in \{0,1\}\) (motion/skill) are extracted automatically from demonstration gripper states. Frames where any gripper is closed are labeled skill=1, with a \(K\)-frame expansion margin on both sides to capture pre-contact and post-contact phases (Algorithm 1). The stage head predicts \(\hat{s}_t\) via binary cross-entropy: $\(\mathcal{L}_{\text{stage}} = -\frac{1}{T}\sum_{t=1}^{T}\left[s_t \log \hat{s}_t + (1-s_t)\log(1-\hat{s}_t)\right]\)$

Keyframe Supervision (KS): At each timestep \(t\), the keyframe target \(q_t^{kf}\) is the joint configuration at the next gripper-state transition (open↔closed), with the terminal frame included as a fallback (Algorithm 2). The keyframe head predicts \(\hat{q}_t^{kf}\) via L1 loss: $\(\mathcal{L}_{kf} = \frac{1}{T}\sum_{t=1}^{T}\|\hat{q}_t^{kf} - q_t^{kf}\|_1\)$

Total training loss: $\(\mathcal{L} = \mathcal{L}_{\text{policy}} + \lambda_s \mathcal{L}_{\text{stage}} + \lambda_k \mathcal{L}_{kf}, \quad \lambda_s = \lambda_k = 0.01\)$

Both heads are training-only; at inference the query tokens remain in the sequence but the heads are not invoked, leaving the π0.5 inference loop unchanged.

Key Contributions

  • StaKe framework: plug-in structured supervision (SS + KS) for VLA fine-tuning, architecture- and inference-agnostic.
  • Automatic label extraction: stage labels and keyframe targets derived solely from gripper open/close sequences—no manual annotation.
  • Event-driven keyframe targets: next gripper-transition frame as supervision anchor, outperforming fixed-horizon targets (\(t+H\)).
  • Dedicated learnable query tokens: specialized representation pathways that decouple auxiliary prediction from action-generation features.
  • Validation on bimanual simulation (RoboTwin 2.0, 10 tasks) and single-arm real-robot (Franka, 4 tasks) with consistent and scalable gains.

Results

  • Simulation (RoboTwin 2.0, 10 bimanual tasks, 100 episodes each): StaKe achieves 59.8% average success rate vs. Ï€0.5 52.4%, DP 42.9%, ACT 42.7%, RDT-1B 49.7%; best or tied-best on 7/10 tasks.
  • Average improvement over Ï€0.5: +7.4 pp absolute across all 10 tasks (14% relative).
  • Long-horizon scaling: +8 pp and +12 pp on two 4-keyframe tasks; 46% vs. 42% on the 5-keyframe Stack Blocks Two task.
  • Ablation (4-task subset): Full StaKe 58.5% avg; w/o stage supervision 54.3%; w/o keyframe supervision 53.3%; fixed-horizon target (\(t+H\)) 50.0%; last VL-prefix token instead of dedicated query 50.8%.
  • Expansion margin \(K\): \(K=25\) optimal (86.0% avg on 2-task subset); \(K=0\) and \(K=30\) both degrade; stable across \(K\in[20,30]\).
  • Real-robot (Franka, 4 tasks, 10 trials each): 56% relative improvement over Ï€0.5 (exact per-task numbers cut off in provided text).
  • Keyframe prediction error: L1 error \(\sim10^{-5}\) (Lift Pot, Stack Blocks Two) and \(\sim10^{-4}\) (Handover Block) in normalized joint space throughout rollouts.

Limitations

  • Evaluated only on top of Ï€0.5; generalizability to other VLA architectures (e.g., OpenVLA, RT-2) is not demonstrated.
  • Assumes gripper state is observable and discretizable—may not transfer cleanly to dexterous hands or soft-body contacts without re-engineering the label extraction.
  • Bimanual experiments merge gripper transitions from both arms into a single keyframe sequence; complex asynchronous coordination scenarios may need finer-grained handling.
  • Real-robot evaluation uses only 10 trials per task, limiting statistical confidence.
  • Loss weights \(\lambda_s = \lambda_k = 0.01\) are empirically fixed across all tasks; no principled tuning strategy is provided.
  • The expansion margin \(K\) requires domain knowledge to set (≈1 s pre/post contact), which may vary across robot morphologies or task speeds.

Relevance to Vision-Language Models

StaKe directly addresses a weak link in the VLM→VLA transfer pipeline: while VLMs provide rich cross-modal representations, the standard fine-tuning signal is temporally undifferentiated, wasting the opportunity to exploit the structured temporal grammar of manipulation. The auxiliary supervision acts as a curriculum signal that shapes the backbone's representations toward stage-aware and goal-directed features—a form of task-aware representation learning on top of a pre-trained VLM that complements rather than replaces the underlying model. For VLM researchers, the finding that two small scalar losses (\(\lambda=0.01\)) applied to dedicated learnable tokens can improve a flow-matching policy's success rate by 14–56% underscores how much structured prior knowledge remains unexploited during standard instruction-following fine-tuning. This connects to the broader question of how to inject temporal and causal structure into VLM-based agents without modifying their pre-trained priors.