COCOA: Action-Conditioned Correction for Vision-Language-Action Policies

Vision-language-action (VLA) models have shown strong performance on complex robot manipulation tasks by generating low-level action chunks from multimodal observations and language instructions. However, an erroneous action chunk can derail long-horizon execution and lead to task failure. To address this challenge, we propose $\text{COCOA}$, an action-conditioned correction training framework that learns a separate corrective policy to revise action chunks proposed by a base VLA before execution. Given the current visual observations, language instruction, robot state, and a proposed action chunk, the corrective policy generates a corrected action chunk for execution. To train the corrective policy, $\text{COCOA}$ introduces a two-stage training procedure that combines supervised flow-matching training with correction preference optimization for flow-based VLAs. We apply $\text{COCOA}$ to three flow-based VLA backbones, including ${\pi }_{0.5}$, FLOWER VLA, and Xiaomi-Robotics-0. Our experiments show that $\text{COCOA}$ achieves consistent improvements on two simulation benchmarks, SimplerEnv and CALVIN, as well as on real-robot tasks.