PaperlayerView exact revision on arXiv

Source-linked AI summary

Disentangling Spurious Correlations in Vision-Language-Action Models via Predicting Domain-Invariant Latent Lookahead

Junghyun Kim, Ngseo Kim, ChungWoo Lee, Seoyeon Lee, Woo-Jeong Baek, Adam Zhou, Chip Huyen, Jun-Ki Lee, Gi-Cheon Kang, Byoung-Tak Zhang

arXiv:2609.37165v1cs.ROcs.LG

TL;DR

VLA models can remain brittle under visual shifts because incidental appearance may be confused with task-relevant structure. DILL trains policies with domain-invariant future latents and disentangled current representations. It reduces counterfactual task substitution and achieves 69.1% average success on LIBERO-Plus visual perturbations, 11.4 percentage points above the strongest external baseline.

  • Problem

    VLA models remain brittle under visual distribution shifts, and future-prediction targets can retain domain-specific appearance rather than task-relevant structure.

    Source: §1 · pp. 1–2

  • Method

    DILL learns task and domain latents from transformed trajectories, then supervises VLA policies to predict future task latents while disentangling current representations from domain latents.

    Source: §1 · p. 2

  • Results

    DILL reduces counterfactual task substitution and achieves 69.1% average success across LIBERO-Plus visual perturbations, 11.4 percentage points above OpenVLA-OFT.

    Source: §4.1 · p. 6, §4.2 · p. 7

  • Takeaways & Limitations

    Physical-robot experiments provide evidence that DILL’s benefits extend beyond controlled simulation.

    Source: §6 · p. 9

    DILL targets visual-domain shortcuts rather than other possible spurious factors, and overly strong invariance may suppress task-useful cues and reduce peak in-distribution performance.

    Source: §5 · p. 9

Abstract

from arXiv · show

Vision-Language-Action (VLA) models remain brittle under visual distribution shifts, often relying on spurious correlations tied to domain-specific factors rather than task-relevant structure. We propose Domain-Invariant Latent Lookahead (DILL), a representation-learning framework that mitigates shortcut learning in VLA policies. Our key idea is to supervise policies with domain-invariant future latents learned from domain-transformed trajectory data. A Task-Domain Encoder is trained with contrastive objectives and Gaussian disentanglement regularization to separate task-relevant structure from domain-specific visual variation. The learned encoder then provides future latents for VLA policy learning through lookahead prediction and domain disentanglement, encouraging the policy to focus on task-relevant structure rather than incidental visual factors. Counterfactual task-view evaluations show that DILL reduces shortcut reliance, while LIBERO-Plus evaluations demonstrate improved visual robustness, with 69.1% average success, 11.4 percentage points above the strongest baseline. Real-world manipulation experiments further support DILL's applicability beyond controlled simulation. Complementary latent-space diagnostics show that these behavioral gains are accompanied by representations that better preserve task-consistent structure while suppressing domain-specific variation. Our project page is available at https://dill-vla.github.io/.

1 Introduction

VLA models can rely on visual cues that correlate with actions in training, leaving them brittle under visual distribution shifts. DILL addresses this by supervising policies with domain-invariant future latents; the paper reports reduced shortcut reliance and improved visual robustness.

1 Introduction

  • VLA models remain brittle under shifts in viewpoint, background, lighting, or camera configuration, despite these variations being irrelevant to task success.The paper frames this brittleness as an obstacle to reliable out-of-distribution generalization and practical deployment.Source: §1 · p. 1
  • Limited dataset diversity and pretrained representations retaining domain-specific information can let policies use visual cues as shortcuts instead of task-relevant semantics.Under counterfactual task–domain recomposition, a policy may execute the task associated with the observed domain rather than the commanded task.Source: §1 · pp. 1–2
  • DILL learns task and domain latents from domain-transformed trajectory chunks, then trains policies to predict future task latents and disentangle current representations from domain latents.Contrastive objectives distinguish task-positive and domain-positive pairs, while Gaussian disentanglement regularization separates the latents.Source: §1 · p. 2
  • Under counterfactual task–view compositions, DILL reduces shortcut reliance and improves out-of-distribution task success, with policies less likely to substitute the task associated with the observed view.The paper also reports highest average success among matched-input baselines on LIBERO-Plus visual perturbations and a smaller average performance drop across shifts.Source: §1 · pp. 2–3
  • Physical-robot experiments and latent-space diagnostics provide additional evidence of DILL’s applicability beyond controlled simulation and task-consistent representation structure.The diagnostics report representations organized more by task-consistent trajectory content than shared visual appearance.Source: §1 · p. 3

2 Related Work

Prior work addresses spurious correlations through data diversity and representation design, while future-prediction methods differ in whether they predict pixels or latents. DILL combines predictive task supervision with explicit suppression of domain-specific visual factors.

2.1 Spurious Correlations in Robot Learning

  • Data-centric methods expand nuisance diversity, while representation-centric approaches suppress task-irrelevant factors or inject task-relevant structure; DILL bridges these directions with predictive future latents and domain-invariant disentanglement.The cited examples include novel-view synthesis, synthetic viewpoint augmentation, viewpoint-invariant representations, and spatial information.Source: §2.1 · p. 3

2.2 World Models and Future Prediction in VLA Models

  • Future-prediction VLA methods use either pixel-level imagination or latent-space prediction, whereas DILL focuses on future targets that separate task structure from domain-specific visual variation.The passage notes pixel rollouts can be costly and prone to compounding errors; latent prediction avoids explicit pixel rollouts.Source: §2.2 · p. 3

3 Method

DILL first learns separated task and domain latents from transformed trajectories, then uses them to supervise policy representations and action prediction. The deployed policy uses only the current observation and instruction.

3.1 Overview

  • The Task-Domain Encoder learns task latents invariant to domain changes and domain latents that capture visual variation using contrastive pair objectives and Gaussian disentanglement.The encoder supplies the future-latent supervision used in the subsequent policy-learning stage.Source: §3.1 · p. 3
  • During policy learning, a lookahead latent is aligned with a future task latent, while a current representation is regularized against the domain latent; both condition the action head.At deployment, the policy receives only the current observation and instruction, without the Task-Domain Encoder.Source: §3.1 · pp. 3–4, §3.3 · p. 6

3.2 Task-Domain Latent Disentanglement

  • Task-positive pairs preserve an observation chunk across domain changes, while domain-positive pairs share a domain condition across different chunks.The transformations include viewpoint, environment appearance, camera heterogeneity, and visual degradation.Source: §3.2 · p. 4
  • The task and domain encoders use InfoNCE objectives with their respective positive pairs, and SIGReg on the concatenated latent encourages approximate disentanglement by discouraging cross-covariance.The joint regularizer matches the task-domain latent to an isotropic Gaussian approximation.Source: §3.2 · p. 5

3.3 VLA Policy Learning with Domain-Invariant Latent Lookahead

  • The pretrained encoder extracts future task and domain targets, and the policy aligns its predicted lookahead latent with the future task target.The lookahead latent provides predictive task information about the future observation chunk.Source: §3.3 · p. 5
  • The policy disentangles its current representation from the domain latent and combines it with the lookahead latent for behavior-cloning action prediction.The encoder is not used at inference; the deployed policy receives only the current observation and instruction.Source: §3.3 · pp. 5–6

4 Experiments

The experiments evaluate whether domain-invariant latent lookahead mitigates shortcut learning through controlled simulation and real-world manipulation experiments.

4 Experiments

  • The evaluation combines controlled simulation, where task–domain correlations and visual shifts can be manipulated systematically, with real-world manipulation experiments.The experiments are organized around testing whether domain-invariant latent lookahead mitigates shortcut learning in VLA policies.Source: §4 · p. 6

3. Do the learned latents successfully separate task structure from task-irrelevant domain factors?

Across LIBERO-Plus, counterfactual task–view swaps, latent diagnostics, and real-world tests, DILL shows stronger visual robustness and less shortcut reliance. The latent diagnostics show task-pair similarity remains higher than domain-pair similarity in the final VLA representation.

4.1 LIBERO Shortcut Diagnostic Under Counterfactual Task–View Compositions

  • Under swapped task–view associations, DILL reaches 0.58 OOD success and 0.05 shortcut degree, versus Base VLA’s zero OOD success and 0.73 shortcut degree.The diagnostic trains Task-L and Task-R only in their respective left- and right-view ranges, then swaps those associations at test time.Source: §4.1 · p. 6, Figure 3 · p. 6

4.2 LIBERO-Plus Zero-Shot Robustness Evaluation Under Visual Distribution Shifts

  • 69.1% average success on LIBERO-Plus visual perturbations puts DILL 11.4 percentage points above OpenVLA-OFT, the strongest external baseline.DILL’s average performance drop is 12.5 points, versus 37.6 for OpenVLA-OFT and 52.3 for UniVLA; the matched-augmentation Base VLA control reaches 38.2% average perturbed success.Source: Table 1 · p. 7, §4.2 · p. 7

4.1 LIBERO Shortcut Diagnostic Under Counterfactual Task–View Compositions

  • With matched architecture and augmentation, disentangled future task latents raise OOD success from zero to 0.44; disentangling the current representation raises it further to 0.58.Entangled LA remains at zero OOD success; shortcut degree falls from 0.65 to 0.06 with disentangled future targets, then changes little with the current-head addition (0.06 to 0.05).Source: §4.1 · p. 7

4.3 Pairwise Diagnostics of Learned Latents

  • The final VLA latent has mean similarity 0.94 for task pairs and 0.55 for domain pairs, while the task and domain encoders show complementary similarity patterns.The task encoder’s means are 0.88 for task pairs and 0.35 for domain pairs; the domain encoder’s are 0.94 for domain pairs and 0.34 for task pairs.Source: §4.3 · p. 8, Figure 4 · p. 8

4.4 Real-World Experiments

  • In real-world tests, DILL reaches 80.2% counterfactual command-following with zero observed shortcut execution; at matched 88.9% unperturbed success, it beats DILL-Current by 5.6 and 4.8 points under predefined and new perturbations.DILL-Current reaches 83.3% counterfactual success with no observed shortcut execution; the new perturbations were absent from encoder pair construction.Source: §4.4 · pp. 8–9, Table 2 · p. 9

5 Limitations

DILL targets shortcuts tied to visual-domain factors, not every possible spurious correlation. Strong invariance may also trade off with in-distribution performance when visual cues are useful for control.

5 Limitations

  • DILL addresses visual-domain shortcuts, but the experiments do not cover biases from initial states, object layouts, language templates, task frequencies, or demonstrator styles.Addressing those non-visual shortcuts would require additional nuisance factors, paired interventions, or objectives that discover them more automatically.Source: §5 · p. 9
  • Overly strong invariance may suppress visual information useful for control in some contexts, reducing peak in-distribution performance even as robustness improves when cues mislead.The paper proposes adaptive objectives that preserve action-relevant factors while suppressing action-irrelevant shortcuts in a context-dependent way.Source: §5 · p. 9
  • A visual factor can be a spurious cue in one setting yet part of the task-relevant scene configuration in another.The paper therefore cautions against uniformly removing domain-specific information.Source: §5 · p. 9

6 Conclusion

DILL addresses VLA shortcut learning by encouraging policies to use task-relevant structure stable across visual domains. The appendix details its data, objectives, implementation, costs, and counterfactual-view diagnostic.

6 Conclusion

  • DILL encourages VLA policies to rely less on incidental appearance and more on task-relevant structure, with simulation and physical-robot experiments supporting benefits beyond simulation.The conclusion frames robustness as shaping what the policy learns, not only scaling data or architectures.Source: §6 · p. 9

A.1 Domain Transformations and Positive Pair Mining

  • Task-positive pairs preserve an observation chunk across domain changes, while domain-positive pairs apply a shared condition to different clips; transformations span four visual-shift families.The encoder is pretrained on MimicGen, ManiSkill, and OXE trajectory sources, then kept fixed during downstream policy learning.Source: §A.1 · pp. 14, 16

A.2 Task-Domain Latent Disentanglement Objectives

  • Gaussian regularization of concatenated task and domain latents encourages approximate disentanglement, alongside separate InfoNCE objectives for task-positive and domain-positive pairs.Under a joint Gaussian approximation, isotropic matching discourages cross-covariance; it does not guarantee exact independence for arbitrary distributions.Source: §A.2 · pp. 16–17

A.3 Policy Objectives

  • The policy aligns a predicted lookahead latent with a future task target and regularizes its current representation against the domain latent before action prediction.The lookahead loss uses other future task targets in the batch as negatives and stops gradients into the pretrained task encoder.Source: §A.3 · pp. 17–18

A.4 Architecture and Implementation Details; A.5 Training Stage Summary

  • Separate task and domain video-encoder branches use VJEPA2, LoRA adapters, and attention-pooling heads; policy heads predict lookahead and current latents for an action head.The pretrained encoder is reused and frozen for downstream policy training, separating encoder pretraining from downstream tuning.Source: §A.4 · p. 18, §A.5 · p. 18

A.6 Computational cost

  • 273 A100 GPU-hours are required once for Task-Domain Encoder pretraining; versus Base VLA, DILL adds 42% downstream tuning cost, 1.5 GB peak memory, and 0.3 ms per 30-action chunk latency.The encoder is not used at deployment, so the reported additional cost is concentrated in training.Source: §A.6 · p. 19, Table 7 · p. 19

B.1 LIBERO Shortcut Diagnostic

  • The LIBERO diagnostic swaps the training associations between task groups and left/right views, measuring commanded-task success under counterfactual views and shortcut-driven task substitution.Task-L is tested at 70° and Task-R at 10° after training only in their respective left- and right-view ranges.Source: §B.1 · p. 20

B.2 Ablation Study

The ablations test whether DILL’s components prevent shortcut use under counterfactual task–view swaps, while the following analyses characterize its visual-robustness scope and latent similarity diagnostic.

B.2 Ablation Study

  • 82% in-distribution success with source augmentation still yields only 2% Counter OOD success, while Entangled LA reaches 24% Center OOD but 0% Counter OOD.The results show that added visual diversity and generic future prediction do not resolve counterfactual task–view generalization.Source: §B.2 · pp. 21–22
  • DILL w/o CH reaches 44% Center OOD and 44% Counter OOD success with 6% shortcut degree; DILL w/o LA reaches 62% Counter OOD and 3% shortcut degree but only 50% in-distribution success.The disentangled current head suppresses shortcut behavior, while lookahead alignment helps balance nominal execution and counterfactual generalization.Source: §B.2 · p. 22
  • 62% average success, 62% Center OOD success, 58% Counter OOD success, and 5% shortcut degree give full DILL the best overall tradeoff.It matches the best Center OOD result and remains close to the best Counter OOD result.Source: Table 8 · p. 21, §B.2 · p. 22

B.3 LIBERO-Plus evaluation details

  • DILL performs best on visual-appearance and viewpoint shifts, scoring 68.4% on Camera and 68.7% on Noise, but its 50.8% Total remains below OpenVLA-OFT.The Total comparison is chiefly limited by Language and Robot categories, which require invariances beyond visual-domain invariance.Source: Table 9 · p. 22, §B.3 · p. 23

C.1 Pairwise latent similarity protocol

  • The pairwise diagnostic compares cosine similarity for task pairs that preserve trajectory content across domain changes against domain pairs that preserve visual conditions across different trajectories.A larger positive task–domain separation gap indicates a more task-centric representation.Source: §C.1 · p. 24

C.2 Final action-conditioning latent

The final action-conditioning latent is the representation supplied to the action head. DILL shifts its similarity structure toward shared trajectory content and away from shared visual domain, while probes show domain cues remain partly accessible.

C.2 Final action-conditioning latent

  • The final action-conditioning latent directly feeds the action head; DILL forms it by concatenating the current representation and predicted lookahead latent.The task-centric diagnostic expects higher similarity for shared trajectory content across domains than for shared domains across different trajectories.Source: §C.2 · p. 24
  • 0.94 task-pair versus 0.55 domain-pair similarity gives DILL a 0.39 gap, reversing Base VLA’s 0.52 versus 0.70 ordering and −0.18 gap.The domain-pair similarity remains nonzero, but shared visual factors no longer dominate the action-conditioning representation.Source: Figure 8 · p. 25, Table 10 · p. 25, §C.2 · p. 25

C.3 Domain predictability with linear probes

  • 64.4% mean probe accuracy for Base VLA versus 25.4% for DILL, against a 21.7% chance reference, indicates reduced linear access to visual-domain cues.The probe result complements the similarity analysis but does not establish that all domain information is removed.Source: §C.3 · p. 25

D Real-World Evaluation Protocol

The real-world protocol tests both counterfactual color–viewpoint swaps and visual perturbations while keeping task instructions and manipulation goals unchanged. It also specifies robot setup, training, and encoder conditions.

D Real-World Evaluation Protocol

  • A UR5e with a Robotiq gripper is trained using 50 demonstrations per task; the source-pretrained encoder stays frozen without real-world fine-tuning or paired real-view data.RGB observations come from Intel RealSense D435i cameras at 640 × 480 resolution.Source: D Real-World Evaluation Protocol · p. 26
  • Training pairs red-target instructions only with the left view and blue-target instructions only with the right; counterfactual tests reverse those pairings while preserving the commanded task.Both colors and viewpoints appear during training, but their counterfactual combinations are held out.Source: D Real-World Evaluation Protocol · p. 26, Figure 9 · p. 27
  • Counterfactual success measures execution of the commanded task, while shortcut degree measures execution of the target associated with the training viewpoint.The latter captures task substitution tied to the viewpoint association, not all unsuccessful trials.Source: D Real-World Evaluation Protocol · p. 26
  • Shoe placement, tissue pulling, and laptop closing are evaluated without perturbation, under predefined visual transformations, and under new transformations absent from encoder pair construction.All conditions preserve each task’s instruction and manipulation goal; the new transformations include dynamic backgrounds and foreground clutter.Source: D Real-World Evaluation Protocol · p. 26, Figure 10 · p. 28

Related papers

Loading 2609.37165v1…