Towards Generalizable Reinforcement Learning via Causality-Guided Self-Adaptive Representations
Yupei Yang, Biwei Huang, Fan Feng, Xinyue Wang, Shikui Tu, Lei Xu
Abstract
General intelligence requires quick adaptation across tasks. While existing reinforcement learning (RL) methods have made progress in generalization, they typically assume only distribution changes between source and target domains. In this paper, we explore a wider range of scenarios where not only the distribution but also the environment spaces may change. For example, in the CoinRun environment, we train agents from easy levels and generalize them to difficulty levels where there could be new enemies that have never occurred before. To address this challenging setting, we introduce a causality-guided self-adaptive representation-based approach, called CSR, that equips the agent to generalize effectively across tasks with evolving dynamics. Specifically, we employ causal representation learning to characterize the latent causal variables within the RL system. Such compact causal representations uncover the structural relationships among variables, enabling the agent to autonomously determine whether changes in the environment stem from distribution shifts or variations in space, and to precisely locate these changes. We then devise a three-step strategy to fine-tune the causal model under different scenarios accordingly. Empirical experiments show that CSR efficiently adapts to the target domains with only a few samples and outperforms state-of-the-art baselines on a wide range of scenarios, including our simulated environments, CartPole, CoinRun and Atari games.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 36d3c630-9aa5-450e-bd1a-d1101271b3f2Cited by top-tier papers2
- Factored Causal Representation Learning for Robust Reward Modeling in RLHFYupei Yang, Lin Yang, Wanxi Deng, Lin Qu et al.ICML 2026 · 1 citation
- Reward-Preserving Counterfactual State Editing for Offline Reinforcement LearningSiyu Wang, Xiaocong Chen, Mingming Gong, Yong Li et al.ICML 2026
Builds on21
- Mastering Atari with Discrete World ModelsDanijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, Jimmy BaICLR 2021 · 1,170 citations
- Model Based Reinforcement Learning for AtariLukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski et al.ICLR 2020 · 969 citations
- Planning to Explore via Self-Supervised World ModelsRamanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel et al.ICML 2020 · 489 citations
- Data-Efficient Reinforcement Learning with Self-Predictive RepresentationsMax Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm et al.ICLR 2021 · 399 citations
- Mastering Atari Games with Limited DataWeirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel et al.NeurIPS 2021 · 345 citations
Related papers
- AdaRL: What, Where, and How to Adapt in Transfer Reinforcement LearningBiwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane et al.ICLR 2022 · 75 citations
- Meta-Reinforcement Learning Based on Self-Supervised Task Representation LearningMingyang Wang, Zhenshan Bing, Xiangtong Yao, Shuai Wang et al.AAAI 2023 · 22 citations
- Causality Meets Locality: Provably Generalizable and Scalable Policy Learning for Networked SystemsHao Liang, Shuqing Shi, Yudi Zhang, Biwei Huang et al.NeurIPS 2025 · 1 citation
- Curious Representation Learning for Embodied IntelligenceYilun Du, Chuang Gan, Phillip IsolaICCV 2021 · 50 citations
- CATAL: Causally Disentangled Task Representation Learning for Offline Meta-Reinforcement LearningShan Cong, Chao Yu, Xiangyuan LanAAAI 2026
