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

Source: arXiv:2606.26801 · Published 2026-06-25 · By Yuan Xu, Yixiang Chen, Kai Wang, Jiabing Yang, Peiyan Li, Qisen Ma et al.

TL;DR

This paper addresses a key limitation in fine-tuning Vision-Language-Action (VLA) models for robotic manipulation, where existing methods apply uniform supervision across all timesteps without distinguishing different manipulation stages or explicitly supervising transitions between gripper events. Such uniform treatment often results in failures concentrated near these stage boundaries. The authors propose StaKe, a plug-in auxiliary supervision framework that automatically derives two complementary signals from demonstration gripper states without requiring manual annotation: a stage classifier to identify the current manipulation phase (motion vs. skill) and a keyframe predictor that estimates the target joint configuration at the upcoming gripper transition. These auxiliary heads are lightweight, added to the pre-trained VLA model only during training, and do not affect the model’s architecture or inference efficiency.

StaKe was validated extensively on the RoboTwin 2.0 bimanual simulation benchmark with 10 diverse tasks and on real-world single-arm Franka robot tasks of varying horizon lengths. It consistently outperformed strong baselines including the state-of-the-art π0.5 model and other diffusion or transformer-based policies, with relative success rate gains averaging +7.4 percentage points across all simulation tasks and +22.5 percentage points (from 40% to 62.5%) on real robot tasks. Improvements were especially notable on longer-horizon tasks with multiple gripper transitions, confirming that the structured supervision effectively guides the policy through challenging stage boundaries. Ablations confirmed both auxiliary losses contributed complementary improvements, and qualitative results showed accurate stage classification and keyframe prediction tracking during manipulation execution. Overall, StaKe provides a general, annotation-free method to inject structured temporal awareness into VLA fine-tuning, significantly boosting long-horizon robotic manipulation success.

Key findings

• StaKe improves average success rate from 52.4% (π0.5 baseline) to 59.8% on 10 RoboTwin 2.0 bimanual simulation tasks (Table I).
• Relative to other baselines: DP (42.9%), ACT (42.7%), RDT (49.7%), StaKe outperforms all.
• StaKe improves all tasks over π0.5 baseline, with gains ranging from +2 pp (Open Laptop) up to +14 pp (Click Alarmclock), +14 pp (Lift Pot).
• On longer-horizon tasks with more keyframes (4-5), gains are larger, e.g., Stack Blocks Two task success improves from 42% (π0.5) to 46% (StaKe).
• Removing stage supervision or keyframe supervision reduces average success by 4-5 percentage points (Table II), showing both losses are complementary.
• Event-driven keyframe targets outperform fixed-horizon targets by 8.5 percentage points (58.5% vs 50.0%, Table II).
• The learnable query tokens for auxiliary heads outperform replacing with the last VL-prefix token (+7.7 percentage points).
• Real Franka robot tasks saw final success improve from 40% (π0.5) to 62.5% (StaKe) over 4 tasks of increasing keyframe complexity (Fig 5).

Methodology — deep read

The authors target the problem of better fine-tuning VLA models for long-horizon robotic manipulation tasks where the uniform continuous action loss fails to supervise the critical gripper-event transition stages explicitly.

Threat model & assumptions: The adversary is n/a since this is a supervised learning improvement. The setting assumes access to demonstration trajectories with gripper open/close states and joint configurations but no manual labels for segmentation or keyframes.

Data: Experimental datasets include RoboTwin 2.0 bimanual simulated tasks (10 tasks with 50 demonstrations each) and real robot Franka R3 arm tasks (4 tabletop tasks with 50 teleoperated demos each). Instructions are language-conditioned but the stage labels and keyframe targets are automatically extracted from gripper states from the same data.

Architecture/algorithm: Builds on the π0.5 VLA model consisting of a pre-trained vision transformer encoder for images, a language model backbone for instructions, and a flow-matching action expert that predicts joint action chunks via diffusion processes.

StaKe adds two learnable query tokens per timestep to the input token sequence given to the language model backbone. Each of these goes to lightweight auxiliary prediction heads:

• Stage classification head: MLP on the stage query token output, outputs probability of 'skill' vs 'motion' stage, trained by binary cross-entropy loss.
• Keyframe prediction head: MLP on the keyframe query token output, predicts the joint configuration of the next gripper-transition keyframe, trained with L1 loss.

Both auxiliary heads are invoked only during training and do not alter the inference loop.

Training regime: Models are fine-tuned for 20,000 steps on 4 NVIDIA H20 GPUs, batch size 64, learning rate 5e-5. Auxiliary loss weights λs and λk are both set to 0.01 to balance with the main flow-matching policy loss.

Stage labels are extracted by parsing gripper open/closed states across all arms, labeling frames around each closed interval (with ±K frame expansion) as 'skill' stage and the rest as 'motion'. A tuned margin K=25 frames (~1 second) was found best.

Keyframe targets are extracted as the joint configuration at each gripper open/close transition frame (including terminal frame), and per timestep the label is set to the nearest upcoming such keyframe (event-driven supervision).

Evaluation: Each task is evaluated for 100 episodes in simulation, using the built-in benchmark success checker. Comparison is made against π0.5 baseline and three other strong manipulation policies (Diffusion Policy, ACT, RDT). Ablations remove each auxiliary loss and vary design choices like keyframe target definition and query tokens.

Qualitative evaluation visualizes predicted stage labels and keyframe errors during rollouts. In real robot experiments, success is reported cumulatively at each keyframe over 10 trials.

Code and pretrained models were not mentioned as publicly released at the time of writing, but a project website is provided for supplementary materials and videos.

Example end-to-end: For the RoboTwin 2.0 'Stack Blocks Two' 5-keyframe task, gripper open/close timings from demos are used to generate stage labels with K=25 to mark skill intervals. Keyframe joint targets at the 5 gripper transitions plus task end frame are assigned to future frames. During fine-tuning on these demos, the model predicts actions via the flow matching loss while simultaneously predicting the current stage and next keyframe via the auxiliary heads. This structured supervision guides the learned representations to capture transitions, resulting in improved policy success from 42% (π0.5) to 46% (StaKe) as shown in Table I and qualitative figures.

Technical innovations

• Automatically derives structured auxiliary supervision signals (stage labels and keyframe targets) from demonstration gripper states without manual annotation.
• Introduces two learnable query tokens appended to the VLA model backbone for specialized auxiliary heads without altering policy architecture or inference.
• Combines flow-matching policy fine-tuning with auxiliary stage classification and keyframe prediction losses for enhanced temporal structure learning.
• Defines event-driven keyframe targets as joint configurations at the next gripper-transition frame, outperforming fixed-horizon targets.

Datasets

• RoboTwin 2.0 — 10 tasks, 50 demonstrations each — simulated dual-arm manipulation benchmark
• Franka Real Robot Tasks — 4 tabletop tasks, 50 teleoperated demonstrations each — single-arm real robot

Baselines vs proposed

• Diffusion Policy (DP): average success = 42.9% vs StaKe: 59.8%
• ACT: average success = 42.7% vs StaKe: 59.8%
• RDT: average success = 49.7% vs StaKe: 59.8%
• π0.5 baseline: average success = 52.4% vs StaKe: 59.8%
• Ablation w/o stage supervision: average success 54.3% vs full StaKe 58.5%
• Ablation w/o keyframe supervision: average success 53.3% vs full StaKe 58.5%
• Keyframe target fixed-horizon: average success 50.0% vs event-driven StaKe 58.5%
• Query token replaced by last VL prefix token: 50.8% vs dedicated token 58.5%

Figures from the paper

Figures are reproduced from the source paper for academic discussion. Original copyright: the paper authors. See arXiv:2606.26801.

Fig 1: Manipulation failures concentrate around gripper-state transi-

Fig 2: Auxiliary supervision introduced by StaKe. Two learnable query

Fig 3: Overview of the StaKe framework. Two learnable query tokens are appended to the pre-trained VLM backbone, each feeding a lightweight

Fig 4 (page 3).

Fig 5 (page 3).

Fig 6 (page 3).

Fig 7 (page 3).

Fig 8 (page 3).

Limitations

• StaKe relies on binary gripper open/close events, so its approach is tailored to prehensile pick-and-place and not directly applicable to non-prehensile or dexterous manipulation.
• The auxiliary supervision is only used during training; inference does not leverage the stage or keyframe predictions which may leave potential gains untapped.
• Stage boundaries are defined with a fixed expansion margin K, which though robustly tuned here, may not generalize perfectly across all tasks or robotic platforms.
• Experiments focus on specific bimanual simulated and Franka single-arm tasks; generalization to other robot morphologies or more diverse environments remains to be validated.
• Code and pretrained model release status is unclear, which may hamper immediate reproducibility and adoption.
• No explicit adversarial or distribution shift evaluation was conducted to assess robustness to environmental or sensor noise.

Open questions / follow-ons

• Can the auxiliary stage and keyframe representations be leveraged during inference to adapt policy execution in real-time?
• How can StaKe be extended to handle non-prehensile manipulation or more complex tactile-based skill transitions where gripper binary states are insufficient?
• Would combining StaKe with reinforcement learning fine-tuning further enhance long-horizon task robustness and generalization?
• Can automated selection or learned adaptation of the stage label expansion margin K improve performance across diverse task domains without manual tuning?

Why it matters for bot defense

While StaKe is not directly related to bot defense or CAPTCHA security, its core idea of injecting structured, event-driven auxiliary supervision to improve sequence modeling in VLA tasks offers insights relevant to sequential decision tasks under partial observability. For CAPTCHA developers exploring robotic interaction tests, the methodology highlights the power of grounding supervision signals in semantically meaningful discrete events rather than uniform time steps, which could inspire more precise human-bot interaction modeling. Additionally, the efficient use of lightweight auxiliary heads appended to pretrained vision-language backbones without disrupting inference could inform the design of robust multi-modal models that must operate under strict deployment constraints. However, direct applicability is limited since robotic manipulation temporal segmentation differs significantly from bot detection challenges.

Cite

@article{arxiv2606_26801,
  title={ Improving Vision-Language-Action Model Fine-Tuning with Structured Stage and Keyframe Supervision },
  author={ Yuan Xu and Yixiang Chen and Kai Wang and Jiabing Yang and Peiyan Li and Qisen Ma and Yan Huang and Liang Wang },
  journal={arXiv preprint arXiv:2606.26801},
  year={ 2026 },
  url={https://arxiv.org/abs/2606.26801}
}

Read the full paper

• arXiv:2606.26801 (abstract)
• arXiv:2606.26801 (PDF)