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

2026-06-25

AI Summary

Current vision-language-action (VLA) models lack structured supervision during fine-tuning to explicitly model manipulation phases and critical grasping events, often leading to failures in complex grasp transitions. This work proposes StaKe, a framework that, without altering the underlying VLA architecture, automatically infers gripper states to derive phase and keyframe labels, enabling a lightweight multi-task auxiliary head that generates structured supervisory signals without manual annotation. StaKe is the first to jointly integrate phase recognition and keyframe prediction into VLA fine-tuning, substantially enhancing stability in long-horizon manipulation tasks. Experiments demonstrate that StaKe achieves relative success rate improvements of 14% in dual-arm simulation and 56% on a real Franka robot, with particularly pronounced gains in extended tasks involving multiple grasp transitions.

Abstract

Vision-Language-Action (VLA) models have shown strong potential for generalizable robotic manipulation. During fine-tuning, however, action supervision applies equally across all timesteps, without structured supervision on which manipulation stage the robot is in or what the next gripper-event target should be. This causes failures to concentrate around challenging gripper-event transitions. To address this, we propose StaKe, a plug-in auxiliary supervision framework that automatically derives two complementary signals from demonstration gripper states without manual annotation: a stage classifier that identifies the current manipulation stage, and a keyframe predictor that estimates the target joint action at the next gripper transition. Both are modeled as lightweight auxiliary heads that enrich the learned representations during training, while leaving the base VLA policy architecture and inference loop unchanged. Experiments on bimanual simulation and single-arm Franka real-robot tasks show that StaKe consistently improves success rates (relative gains of 14% and 56%, respectively), with larger improvements on longer-horizon tasks that involve more gripper-event transitions. Ablation studies validate each design choice, and qualitative analysis confirms that the learned representations faithfully track manipulation stages. These results indicate that structured supervision is an effective and general strategy for enhancing VLA fine-tuning in long-horizon manipulation. Project website: https://hi-yuanxu.github.io/StaKe-Web/

Problem

• Vision-Language-Action models
• fine-tuning
• structured supervision
• gripper-event transitions
• robotic manipulation

Innovation

• structured supervision
• keyframe prediction
• stage classification
• vision-language-action models
• robotic manipulation