# Improving Generalization Ability of Robotic Imitation Learning by Resolving Causal Confusion in Observations

arXiv: https://arxiv.org/abs/2507.22380 · HTML: https://mlsanity.com/paper/2507.22380/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 30 Jul 2025.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 25 of 25)

- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [π0: A Vision-Language-Action Flow Model for General Robot Control](https://mlsanity.com/paper/2410.24164/) (2024)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations](https://mlsanity.com/paper/1811.12359/) (2018)
- [Solving Rubik's Cube with a Robot Hand](https://mlsanity.com/paper/1910.07113/) (2019)
- [Perceiver-Actor: A Multi-Task Transformer for Robotic Manipulation](https://mlsanity.com/paper/2209.05451/) (2022)
- [An Algorithmic Perspective on Imitation Learning](https://mlsanity.com/paper/1811.06711/) (2018)
- [Sim-to-Real Transfer in Deep Reinforcement Learning for Robotics: a Survey](https://mlsanity.com/paper/2009.13303/) (2020)
- [Transporter Networks: Rearranging the Visual World for Robotic Manipulation](https://mlsanity.com/paper/2010.14406/) (2020)
- [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [A Survey on Causal Inference](https://mlsanity.com/paper/2002.02770/) (2020)
- [Sim-to-Real via Sim-to-Sim: Data-efficient Robotic Grasping via Randomized-to-Canonical Adaptation Networks](https://mlsanity.com/paper/1812.07252/) (2018)
- [A Survey of Imitation Learning: Algorithms, Recent Developments, and Challenges](https://mlsanity.com/paper/2309.02473/) (2023)
- [On Feature Learning in the Presence of Spurious Correlations](https://mlsanity.com/paper/2210.11369/) (2022)
- [Foundations of Structural Causal Models with Cycles and Latent Variables](https://mlsanity.com/paper/1611.06221/) (2016)
- [ALOHA Unleashed: A Simple Recipe for Robot Dexterity](https://mlsanity.com/paper/2410.13126/) (2024)
- [Resnet in Resnet: Generalizing Residual Architectures](https://mlsanity.com/paper/1603.08029/) (2016)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Structured Domain Randomization: Bridging the Reality Gap by Context-Aware Synthetic Data](https://mlsanity.com/paper/1810.10093/) (2018)
- [Sim-to-Real Transfer for Vision-and-Language Navigation](https://mlsanity.com/paper/2011.03807/) (2020)
- [Learning Causally Disentangled Representations via the Principle of Independent Causal Mechanisms](https://mlsanity.com/paper/2306.01213/) (2023)

## Built on by (most cited first, 4 of 4)

- [Shortcut Learning in Generalist Robot Policies: The Role of Dataset Diversity and Fragmentation](https://mlsanity.com/paper/2508.06426/) (2025)
- [Do You Need Proprioceptive States in Visuomotor Policies?](https://mlsanity.com/paper/2509.18644/) (2025)
- [SEVO: Semantic-Enhanced Virtual Observation for Robust VLA Manipulation via Active Illumination and Data-Centric Collection](https://mlsanity.com/paper/2605.11114/) (2026)
- [What Matters, When? Diagnosing and Improving Conditional Visual Grounding in Visuomotor Imitation Policies](https://mlsanity.com/paper/2609.05376/) (2026)

## Abstract (arXiv)

Recent developments in imitation learning have considerably advanced robotic manipulation. However, current techniques in imitation learning can suffer from poor generalization, limiting performance even under relatively minor domain shifts. In this work, we aim to enhance the generalization capabilities of complex imitation learning algorithms to handle unpredictable changes from the training environments to deployment environments. To avoid confusion caused by observations that are not relevant to the target task, we propose to explicitly learn the causal relationship between observation components and expert actions, employing a framework similar to [6], where a causal structural function is learned by intervention on the imitation learning policy. Disentangling the feature representation from image input as in [6] is hard to satisfy in complex imitation learning process in robotic manipulation, we theoretically clarify that this requirement is not necessary in causal relationship learning. Therefore, we propose a simple causal structure learning framework that can be easily embedded in recent imitation learning architectures, such as the Action Chunking Transformer [31]. We demonstrate our approach using a simulation of the ALOHA [31] bimanual robot arms in Mujoco, and show that the method can considerably mitigate the generalization problem of existing complex imitation learning algorithms.