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

Research Question and Approach

Vision-Language-Action (VLA) models have become a standard for robotic manipulation, yet they often struggle with long-horizon tasks. A primary limitation is that standard fine-tuning applies a uniform action loss across all timesteps, failing to account for the distinct phases of a task—such as free-space motion versus contact-constrained skill execution. The authors propose StaKe, a plug-in framework that introduces structured auxiliary supervision to help the model better understand these manipulation stages and critical gripper-event transitions.

StaKe adds two lightweight auxiliary heads to the VLA model during training: a stage classifier that identifies whether the robot is in a motion or skill phase, and a keyframe predictor that estimates the target joint configuration for the next gripper event. These labels are derived automatically from demonstration gripper states, eliminating the need for manual annotation. Crucially, these heads are only active during training; the base VLA policy architecture and the inference loop remain entirely unchanged.

The authors evaluated StaKe on both bimanual simulation tasks (RoboTwin 2.0) and single-arm real-robot tasks (Franka). In simulation, StaKe consistently outperformed the base VLA policy (π0.5) and other strong baselines like Diffusion Policy and RDT, achieving an average success rate of 59.8% compared to 52.4% for the baseline. The improvements were particularly pronounced in longer-horizon tasks that involved more frequent gripper-state transitions. Real-world experiments further confirmed these gains, showing a 56% relative improvement in success rates. Qualitative analysis and ablation studies validated that the auxiliary heads successfully enrich the learned representations, allowing the model to more faithfully track the manipulation stages.

This work demonstrates that VLA models can be significantly improved by injecting domain-specific structure into the training process without increasing the complexity of the final deployed model. By leveraging the inherent structure of robotic trajectories—specifically the transition between free-space motion and contact—StaKe provides a scalable way to enhance performance on complex, multi-stage tasks. Because the auxiliary supervision is fully automated and does not alter the inference procedure, it offers a highly practical, plug-and-play strategy for researchers looking to boost the reliability of existing VLA policies in real-world environments.