Research map / Reinforcement learning / Visual and representation RL
Building Minimal and Reusable Causal State Abstractions for Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Causal Bisimulation Modeling (CBM), which learns causal dynamics with implicit (contrastive) models and a causal reward model to derive minimal task-specific state abstractions for sample-efficient, generalizable RL.
Problem
Learning state abstractions in factored MDPs that are minimal and reusable across tasks in the same environment, so RL is sample-efficient and generalizes to out-of-distribution irrelevant variables; existing explicit causal dynamics models (CDL) are inaccurate and abstractions are non-minimal.
Approach
Learn shared causal dynamics with implicit InfoNCE-based models (with a regularized CMI estimator fixing overfitting and importance-sampling issues) plus a per-task causal reward model, then take the causal ancestors of the reward as a bisimulation-equivalent minimal abstraction.
Key results
- Implicit dynamics model achieves 90.5% causal graph accuracy on block vs 87.5% for explicit, and 85.5% vs 82.6% on tool-use.
- Implicit models yield state abstraction accuracy at least 20% higher than explicit models across tasks.
- CBM achieves near-Oracle sample efficiency and stays robust to OOD values of task-irrelevant variables on Pick and Stack, where CDL fails to generalize.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use implicit rather than explicit dynamics models | CBM dynamics learning · Block causal graph | accuracy: 87.5 ± 0.1 -> 90.5 ± 0.4 | Helped |
| Use implicit rather than explicit dynamics models | CBM dynamics learning · Block Pick abstraction | accuracy: 53.2 ± 4.6 -> 95.7 ± 6.0 | Helped |
| Use implicit rather than explicit dynamics models | CBM dynamics learning · Block Stack abstraction | accuracy: 59.6 ± 4.6 -> 95.7 ± 6.0 | Helped |
| Use implicit rather than explicit dynamics models | CBM dynamics learning · Tool-use causal graph | accuracy: 82.6 ± 0.2 -> 85.5 ± 0.1 | Helped |
| Use implicit rather than explicit dynamics models | CBM dynamics learning · Tool-use Series abstraction | accuracy: 80.0 ± 1.5 -> 98.8 ± 1.3 | Helped |
| Use implicit rather than explicit dynamics models (ablation) | CBM task learning · Pick, Stack, Series, and DMC tasks | sample efficiency | Helped |
| Use task-specific causal reward information to refine the state abstraction | CBM task learning · All evaluated tasks | sample efficiency | Helped |
| Use a minimal task-specific abstraction instead of the full state | CBM · Series | episode reward: 200 over 40k training episodes | Helped |
| Use a minimal task-specific abstraction | CBM · Pick and Stack OOD states | task performance | Helped |
| Use a minimal task-specific abstraction | CBM · Walker | episode reward: 1000 by 20k episodes | No clear effect |
Surprising findings
- TIA and Denoised MDP learn meaningless abstractions when their dynamic-independence assumptions fail in the manipulation environments.
Limitations
- Assumes a pre-defined factored state space; not applicable to high-dimensional images directly.
- Assumes independent per-variable transitions, which is violated for e.g. quaternion rotation variables.
- Some baselines (TIA, Denoised MDP) fail to learn meaningful tasks in the manipulation environments.
Experimental setup
Datasets: Robosuite block environment (Pick, Stack) Robosuite tool-use environment (Series) DeepMind Control Suite Cheetah and Walker
Benchmarks: Robosuite block (Pick, Stack) Robosuite tool-use (Series) DeepMind Control Suite Cheetah DeepMind Control Suite Walker
Metrics: causal graph accuracy state abstraction accuracy success rate episode reward sample efficiency OOD generalization
Baselines: CDL TIA Denoised MDP Oracle (ground-truth abstraction) Full state RL explicit dynamics model
Builds on: Causal Dynamics Learning (CDL) soft actor-critic implicit behavioral cloning / implicit models bisimulation
Closest papers by meaning
- Causal Dynamics Learning for Task-Independent State Abstraction (ICML 2022 · similarity 0.92)
- Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning (ICML 2022 · similarity 0.89)
- Learning Causal Dynamics Models in Object-Oriented Environments (ICML 2024 · similarity 0.88)
- Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning (ICML 2024 · similarity 0.87)
- Causal Information Prioritization for Efficient Reinforcement Learning (ICLR 2025 · similarity 0.86)
- Reconciling Spatial and Temporal Abstractions for Goal Representation (ICLR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.