Action-Sufficient State Representation Learning for Control with Structural Constraints
Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bernhard Schölkopf, Kun Zhang
Abstract
Perceived signals in real-world scenarios are usually high-dimensional and noisy, and finding and using their representation that contains essential and sufficient information required by downstream decision-making tasks will help improve computational efficiency and generalization ability in the tasks. In this paper, we focus on partially observable environments and propose to learn a minimal set of state representations that capture sufficient information for decision-making, termed Action-Sufficient state Representations (ASRs). We build a generative environment model for the structural relationships among variables in the system and present a principled way to characterize ASRs based on structural constraints and the goal of maximizing cumulative reward in policy learning. We then develop a structured sequential Variational Auto-Encoder to estimate the environment model and extract ASRs. Our empirical results on CarRacing and VizDoom demonstrate a clear advantage of learning and using ASRs for policy learning. Moreover, the estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space to improve sample efficiency.
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 6b22da28-cb41-49fc-971d-883d17848c8cCited by top-tier papers13
- MoCoDA: Model-based Counterfactual Data AugmentationSilviu Pitis, Elliot Creager, Ajay Mandlekar, Animesh GargNeurIPS 2022 · 60 citations
- Learning World Models with Identifiable FactorizationYuren Liu, Biwei Huang, Zhengmao Zhu, Hong-Long Tian et al.NeurIPS 2023 · 32 citations
- ACE: Off-Policy Actor-Critic with Causality-Aware Entropy RegularizationTianying Ji, Yongyuan Liang, Yan Zeng, Yu Luo et al.ICML 2024 · 20 citations
- Policy-Guided Causal State Representation for Offline Reinforcement Learning RecommendationSiyu Wang, Xiaocong Chen, Lina YaoWWW 2025 · 5 citations
- BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement LearningHaohong Lin, Wenhao Ding, Jian Chen, Laixi Shi et al.NeurIPS 2024 · 5 citations
Builds on9
- Dream to Control: Learning Behaviors by Latent ImaginationDanijar Hafner, Timothy P. Lillicrap, Jimmy Ba, Mohammad NorouziICLR 2020 · 1,852 citations
- CURL: Contrastive Unsupervised Representations for Reinforcement LearningMichael Laskin, Aravind Srinivas, Pieter AbbeelICML 2020 · 1,261 citations
- 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
- Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable ModelAlex X. Lee, Anusha Nagabandi, Pieter Abbeel, Sergey LevineNeurIPS 2020 · 437 citations
Related papers
- Variational Recurrent Models for Solving Partially Observable Control TasksDongqi Han, Kenji Doya, Jun TaniICLR 2020 · 75 citations
- Learning Task-Sufficient World Models by Synergizing Agentic Exploration and Structured ModelingFan Feng, Yujia Zheng, Minghao Fu, Yongqiang Chen et al.ICML 2026
- Which Mutual-Information Representation Learning Objectives are Sufficient for Control?Kate Rakelly, Abhishek Gupta, Carlos Florensa, Sergey LevineNeurIPS 2021 · 44 citations
- Value-Consistent Representation Learning for Data-Efficient Reinforcement LearningYang Yue, Bingyi Kang, Zhongwen Xu, Gao Huang et al.AAAI 2023 · 19 citations
- Scalable Decision-Making in Stochastic Environments through Learned Temporal AbstractionBaiting Luo, Ava Pettet, Aron Laszka, Abhishek Dubey et al.ICLR 2025
