CAST: Counterfactual Labels Improve Instruction Following in Vision-Language-Action Models

Technical Summary: CAST – Counterfactual Labels Improve Instruction Following in Vision-Language-Action Models

Problem Statement
Generalist robots require the ability to follow diverse, open-vocabulary language instructions. While Vision-Language-Action (VLA) models offer a promising architecture for mapping language to robot actions, they currently struggle with fine-grained instruction following. A primary limitation is the lack of semantic diversity and language grounding in existing robot datasets. In many datasets, observations are highly correlated with specific actions (e.g., a chest of drawers implies "open the drawer"), leading to posterior collapse. In this phenomenon, the policy learns to ignore language inputs and infer tasks solely from visual observations because the language provides no additional predictive information. Furthermore, existing datasets often lack fine-grained task diversity for similar observations, and some navigation datasets lack language labels entirely. Scaling up data curation to address these issues is expensive and tedious.

Methodology: Counterfactual Augmentation with Synthetic Trajectories (CAST)
The authors propose CAST, a data augmentation framework that leverages Vision-Language Models (VLMs) to generate counterfactual language labels and synthetic actions for existing, uncurated robot trajectory datasets. The core insight is that by training policies on multiple distinct actions associated with the same observation (but driven by different instructions), the model is forced to attend to the language input to predict the correct action, thereby mitigating posterior collapse.

The CAST pipeline operates in four main stages:

1. Decision Point Identification: Existing trajectories are segmented into atomic chunks based on behavioral changes (e.g., turning, stopping). These boundaries serve as decision points.
2. Counterfactual Instruction Generation: At each decision point, a VLM (specifically Gemini 2.5 Pro) is queried with the current observation and the original trajectory context. The VLM generates plausible counterfactual instructions (e.g., "Turn left" instead of the executed "Turn right") and corresponding atomic commands (e.g., "turn left").
3. Synthetic Action Generation: A pre-trained atomic policy ($\pi_a$), trained on real robot data to follow simple atomic commands (e.g., "go forward," "turn left"), is used to generate the actual action chunks for these counterfactual instructions. This policy branches off from the original trajectory to create a new, plausible "counterfactual ending."
4. Dataset Augmentation: The original trajectories are augmented with these new $(observation, counterfactual\_instruction, counterfactual\_action)$ tuples. This expanded dataset is used to train the final VLA policy, termed CounterfactualVLA.

The authors provide a theoretical motivation based on mutual information. They argue that language following is improved by maximizing the mutual information $I(a; \ell|o)$ between action ($a$) and language ($\ell$) conditioned on observation ($o$). By increasing the entropy of atomic commands ($H(\ell_a|o)$) for a given observation while ensuring the atomic command is predictable from the full instruction ($H(\ell_a|\ell, o)$), CAST effectively pushes up a lower bound on this mutual information.

Implementation Details
The authors instantiate CAST in two domains:

• Navigation: Using the GNM dataset (indoor/outdoor), they employ a hindsight relabeling pipeline to generate initial labels for unlabeled data. They segment trajectories into atomic behaviors (turn left/right, go forward, stop) to train an atomic policy. The VLM generates counterfactual navigation prompts, and the atomic policy synthesizes the corresponding action chunks.
• Manipulation: Using the Bridge dataset, they utilize existing task-level labels and subtask decompositions (e.g., "pick up," "move to") as atomic commands. The VLM generates counterfactual tasks (e.g., placing an object on a different surface), and a pre-trained Steerable Policy generates the actions.

In both cases, the final CounterfactualVLA is a fine-tuned version of the 3B PaliGemma VLM, trained on the augmented dataset using action chunking to ensure temporal consistency.

Experimental Results
The authors evaluate the method through several lenses:

• Label Quality: Human judges evaluated the generated labels for accuracy regarding objects, motion, and scene grounding. The counterfactual labels were found to be accurate and plausible.
• Navigation Performance: CounterfactualVLA was tested on 27 instruction-following tasks across three environments (office, kitchen, outdoor park). The tasks included object navigation, referential navigation (spatial relationships), and continuous navigation.
  • CounterfactualVLA outperformed VLAs trained on unaugmented data by 27%.
  • It achieved a success rate double that of the baseline.
  • It outperformed state-of-the-art methods in the domain.
• Manipulation Performance: Evaluated on pick-and-place tasks with distractors, the method showed a 27% performance improvement over a VLA trained with standard data.

Significance and Claims
The paper claims that CAST provides a scalable, data-efficient method to improve the steerability and instruction-following capabilities of VLA policies without requiring additional real-world data collection. By leveraging off-the-shelf VLMs to generate counterfactual data, the method addresses the "posterior collapse" issue inherent in current datasets. The authors position this work as a "first step toward truly open-set robotic language following," demonstrating that augmenting existing datasets with synthetic counterfactuals can significantly enhance a robot's ability to follow complex, fine-grained commands in diverse environments. The code, data, and checkpoints are open-sourced to facilitate further research.