Source-linked AI summary

CounterAlign: Counterfactual Supervision for Vision-Language-Action Models

Haru Kondoh, Kei Ota, Asako Kanezaki, Yueh-Hua Wu

arXiv:2608.21740v1cs.RO

TL;DR

VLA behavior cloning lacks explicit negative supervision, while robotics RL rewards and non-expert data are costly. CounterAlign relabels successful expert demonstrations to create counterfactual supervision and combines adversarial discriminators with offline RL, improving robustness to position and task perturbations and outperforming strong baselines. Its benefits are especially relevant when large, diverse robot datasets are expensive or impractical to collect.

  • Problem

    Behavior cloning supplies only positive supervision, while RL rewards and curated non-expert trajectories are costly to obtain for robotics.

  • Method

    CounterAlign pairs expert action chunks with mismatched instructions and uses relabeling and adversarial discriminators to learn a dense instruction-grounded reward for offline RL.

  • Results

    CounterAlign improves robustness and generalization under object-position and task perturbations, maintains original-distribution performance, and outperforms strong baselines in simulation and real-robot experiments.

  • Takeaways & Limitations

    Extracting denser corrective supervision from existing demonstrations can complement collecting additional data for data-constrained VLA learning.

  • Takeaways & Limitations

    The method could be strengthened with larger and more diverse datasets spanning more objects, environments, embodiments, and linguistic variations.

Abstract

from arXiv · show

Vision-Language-Action (VLA) models are typically trained with behavior cloning (BC) on expert demonstrations. However, BC provides only positive supervision for expert actions, without explicit negative supervision indicating which actions are instruction-inconsistent or otherwise inappropriate. Reinforcement learning (RL) can provide such corrective signals, but often relies on externally specified rewards or curated non-expert data, both of which are costly to obtain in robotics. We show that offline RL for VLA models need not rely on curated non-expert trajectories: successful expert demonstrations alone can be transformed into dense corrective supervision through instruction relabeling. Specifically, by pairing expert actions with mismatched alternative instructions, we synthesize counterfactual instruction-observation-action tuples from the dataset and combine them with adversarial discriminator training to learn an instruction-grounded reward model for offline RL, without collecting additional rollouts or annotations. On the robustness-focused LIBERO-PRO benchmark, our method improves robustness to object position and task perturbations over a strong state-of-the-art baseline. It also outperforms competitive baselines in real-robot experiments on the TX-G2 (compatible with AGIBot G2). More broadly, our results suggest that, for data-constrained VLA learning, extracting denser supervision from each demonstration can complement collecting additional data.

1. Introduction

CounterAlign addresses the limits of behavior cloning and costly robotics supervision by converting expert demonstrations into counterfactual, instruction-grounded corrective signals for offline RL. It combines relabeling and adversarial discriminators to improve semantic robustness and generalization without additional data collection or annotation.

  • Motivation: Behavior cloning provides positive expert-action supervision but does not identify instruction-inconsistent alternatives, contributing to brittleness under paraphrases, object-position shifts, and scene variations.
  • Motivation: Offline RL can provide corrective supervision without interaction, but requires rewards that capture instruction satisfaction and datasets with sufficiently informative behavioral support.
  • CounterAlign: CounterAlign pairs successful expert action chunks with semantically inappropriate instructions to synthesize counterfactual instruction–observation–action tuples without interaction, annotations, or corrective demonstrations.
  • CounterAlign: nnPU learning treats jointly relabeled samples as unlabeled because different instructions can share valid action chunks, reducing label noise in the instruction-grounded reward.
  • Results: CounterAlign alternates discriminator refinement with offline RL policy optimization, improving robustness and generalization under object-position and task perturbations while maintaining original-distribution performance and outperforming strong simulation and real-robot baselines.
  • CounterAlign: The framework uses adversarial and relabeling discriminators: one contrasts expert with policy-generated actions, while the other evaluates samples formed by replacing instructions or actions from the dataset.

4. Experiments

CounterAlign is evaluated on simulated LIBERO-PRO and real TX-G2 experiments, with emphasis on robustness to positional and task perturbations. Ablations and learned-advantage visualizations examine which components contribute and whether rewards track instruction-relevant behavior.

  • Simulation Experiments: In Sem and Obj settings, CounterAlign shows no substantial performance gap over behavior cloning because replaying familiar action sequences can overestimate BC performance.The reported explanation is that these settings may not require semantic generalization beyond training-like observations.
  • Simulation Experiments: CounterAlign’s largest LIBERO-PRO improvements occur in Pos and Task settings, which require adapting actions to changed object configurations or task conditions.These settings more directly test instruction understanding and situation-appropriate action selection than Sem and Obj settings.
  • Ablation: The ablation finds that offline RL already outperforms behavior cloning in diagnostic Pos and Task settings, while discriminator and critic relabeling, entropy regularization, and unlabeled samples improve robustness.Applying relabeling directly to the actor does not consistently provide additional gains.
  • Real-World Experiments: On TX-G2, CounterAlign outperformed baselines across trials with perturbations to the target object’s initial position and orientation.The real-robot evaluation used 10 trials, including two perturbations outside the training distribution, and measured mean success rate across primitive tasks.
  • Qualitative Results: Learned advantage values tend to decrease during failure segments and increase when the robot correctly follows the instruction, while remaining relatively low at episode boundaries.The authors associate low initial values with ambiguity and low final values with weakly instruction-related hand-lifting motion.
  • Real-World Experiments: OpenVLA-OFT could not be evaluated on the real robot because its parameter count exceeded the available GPU memory.

5. Conclusion

CounterAlign is an offline reinforcement-learning framework that learns instruction–observation–action alignment rewards from offline data. Experiments show improved robustness under position and task perturbations, including on the TX-G2 real robot, while broader and more diverse datasets remain a future direction.

  • CounterAlign learns an instruction–observation–action alignment reward with adversarial and relabeling-based discriminators trained entirely from offline data.The reward is integrated into an IQL-style objective with advantage-weighted flow matching.
  • CounterAlign improves over behavior cloning particularly under position and task perturbations, where policies must adapt to changed object positions or task contexts.
  • The gains carry over to TX-G2 real-robot experiments, where CounterAlign outperforms competitive baselines.
  • Future strengthening is proposed through larger, more diverse datasets covering broader object categories, environments, embodiments, and linguistic variations.The paper also identifies combining counterfactual supervision with complementary data collection or adaptation strategies as a direction.

Supplementary Material for CounterAlign: Counterfactual Supervision for

The supplementary material is titled for Vision-Language-Action Models and provides additional details on the paper’s topics. The supplied material does not specify those details further.

  • The supplementary material is associated with Vision-Language-Action Models.
  • The supplementary material provides additional details on several topics, which are not enumerated in the supplied passage.

1. Definitions of similarity functions for relabeling

The relabeling procedure uses offline heuristic similarity functions computed from expert data to compare reference samples with samples associated with different instructions. These functions cover observations, actions, instructions, and proprioceptive robot-state information.

  • Similarity scores for relabeling are heuristic and computed offline from the expert dataset.
  • For a reference instruction and observation–action sample, candidate samples are associated with a different instruction for comparison.
  • Observation similarity and action similarity compare the reference sample with observation–action pairs associated with another instruction.An observation feature extractor is used in defining observation similarity.
  • Instruction similarity is defined from the observation and action similarities, with a joint-relabeling variant that omits the action term.
  • Proprioceptive similarity filters candidate action chunks during joint relabeling using the robot-state component of the observation.

2. Details of Relabeling Strategies

The relabeling strategies create informative mismatches by altering instructions, actions, or both while preserving plausible task context. Similarity filtering and unlabeled treatment reduce ambiguity and label noise, and all strategies operate offline on expert data.

  • Offline Construction: All three relabeling strategies are applied offline to the expert dataset and require no environment interaction.The resulting datasets supply labeled negatives or unlabeled samples for discriminator training.
  • Instruction Relabeling: Instruction relabeling pairs a semantically different instruction with the original observation and action to create plausible hard negatives.Moderate instruction similarity avoids trivial negatives and preserves meaningful semantic distinctions.
  • Action Relabeling: Action relabeling keeps the instruction and observation fixed while substituting an action chunk from another expert sample.The substituted action is selected to be physically distinguishable but not trivially incorrect.
  • Joint Relabeling: Joint relabeling replaces both instruction and action chunks, using observation and proprioceptive similarity to construct physically plausible but ambiguous tuples.Because these tuples may remain semantically consistent, they are treated as unlabeled rather than definitive negatives.

3. Actor-Side Relabeling (Ablation Variant)

The actor-side ablation extends relabeling beyond critic training by supplying near-positive relabeled samples to the policy objective. This variant differs from the main method, where only the critic uses relabeled data.

  • Design: In the main method, only the critic is trained with relabeled data, whereas the actor uses the expert dataset.The ablation tests whether relabeled samples should also enter actor training.
  • Dataset Construction: Actor-side relabeling uses near-positive similarity thresholds rather than the hard-negative thresholds used for discriminator relabeling.Instruction- and action-relabeled datasets are constructed for this extended variant.
  • Policy Objective: The extended variant replaces the expert dataset in the policy loss with an actor-training dataset containing relabeled samples.The construction follows the discriminator relabeling procedures but selects highly similar alternatives.

4. Complete Algorithm of the Proposed Method

The proposed algorithm first constructs relabeled datasets offline, then alternates discriminator, critic, and actor updates for K iterations. Its advantage-weighted policy objective uses offline samples while retaining positive action weighting.

  • Training Procedure: Phase 1 augments the expert dataset through instruction, action, and joint relabeling to produce negatives, unlabeled samples, and critic-training data.Phase 2 alternates updates to discriminators, critic, and actor using the induced reward.
  • Optimal Policy: The KL-constrained optimal policy weights behavior-policy actions by exp(βA(l, o_t, a_t:t+h)) and normalizes them with Z(l, o_t).The behavior policy μ and context-observation distribution d_μ define the offline policy-improvement problem.
  • Policy Objective: Because samples from the optimal policy are unavailable, importance sampling rewrites the policy objective as an expectation over behavior-policy samples from the offline dataset.Substitution cancels behavior-policy terms, leaving the partition function in the importance weight.
  • Policy Objective: Omitting Z preserves relative advantage-based action weighting within each context while changing only marginal context weighting.Explicitly estimating Z is unreliable in offline continuous control with unknown behavior policies and limited action samples.
  • Offline Stability: Using exp(βA) gives every offline update a positive weight, avoiding the gradient reversal and potential out-of-distribution behavior associated with negative advantage weights.This distinguishes the objective from online AWM, which weights flow matching by the raw advantage A.

6. Model Architecture and Parameter Counts of BC and the Proposed Method

The proposed method adds discriminator and critic networks while leaving the policy architecture unchanged. Backbone sharing reduces the resulting parameter overhead, with separate expert heads preserving module specialization.

  • Parameter Counts: The BC baseline has approximately 3.4B trainable parameters, compared with approximately 12.7B for the proposed method.The increase comes from additional discriminator and critic networks; the policy itself is unchanged.
  • Backbone Sharing: Auxiliary networks share VLM backbones extensively while retaining lightweight, separate expert heads for task-specific objectives.When backbones are separate, token embeddings are shared.
  • Critic Architecture: The Q- and V-networks share one VLM backbone, and the two double-Q networks share another backbone while differing in critic expert heads.This design shares representation computation while separating critic outputs.
  • Target and Policy: The target network uses a fully separate VLM backbone and critic heads so its estimates remain stable under EMA updates.The policy retains exactly the BC baseline architecture without structural modifications.
  • Optimization: All parameters are fine-tuned rather than adapted with LoRA or other parameter-efficient methods.This follows the original π0.5 setup.

7. Training Configuration

Models use different training configurations for LIBERO and TX-G2, with substantially longer training and more total compute for the real-robot setting.

  • LIBERO: LIBERO models train for 15K gradient steps with batch size 256 on 8 NVIDIA H200 GPUs for approximately 28 hours.Training approximately converges within 10K–15K gradient steps.
  • TX-G2: TX-G2 models train for 150K gradient steps with batch size 64 on 8 H200 GPUs for approximately 96 hours.The 150K-step checkpoint is used for evaluation.
  • Comparison: TX-G2 training uses ten times as many gradient steps as LIBERO and takes substantially longer.This comparison follows directly from the reported 150K versus 15K steps and 96 versus 28 hours.

8. Detailed Ablation Analysis

The ablations attribute gains to offline RL, relabeling, and applying relabeling to the discriminator and critic, while results vary with benchmark data diversity and evaluation horizon.

  • Offline RL: The basic offline RL configuration improves over standard BC, using adversarial discriminator training, IQL for the critic, and advantage-weighted flow matching for the actor.This configuration demonstrates policy improvement without the additional relabeling components.
  • Discriminator relabeling: Relabeling helps the discriminator evaluate instruction–action semantic consistency, encouraging instruction-conditioned action selection rather than observation-based imitation.Entropy regularization and unlabeled treatment of ambiguous relabeled samples further improve the reward signal and performance.
  • Critic relabeling: Applying relabeling to the critic produces additional performance gains by transferring language–action correspondence into more appropriate advantage signals.The reported transfer proceeds from the relabeled discriminator knowledge through the critic to the actor’s advantage estimates.
  • Actor relabeling: Relabeling the actor provides no further improvement and may introduce noise because the learned advantage function is imperfect.The ablation therefore favors incorporating relabeling indirectly through the discriminator and critic.
  • CALVIN: CALVIN ABC →D shows a small but consistent advantage over BC, with average length increasing from 3.93 to 4.03 and improvements at every evaluation stage.The benchmark evaluates generalization to unseen environment D across 1,000 trials.
  • Data diversity: The smaller CALVIN gap is attributed to its richer, more diverse training data, whereas the method is particularly useful when large-scale data collection is costly or impractical.CALVIN provides approximately 41.8K steps per task versus approximately 6.8K for LIBERO, according to the comparison passage.
  • Long-horizon evaluation: Across 1 to 5 consecutive CALVIN tasks, success rates decrease monotonically as errors accumulate, while the proposed method remains ahead of BC at every level.The passage characterizes this decline as a natural consequence of multiplying per-task success probabilities rather than method-specific degradation.
Loading 2608.21740v1…