Building Minimal and Reusable Causal State Abstractions for Reinforcement Learning

Abstract

Two desiderata of reinforcement learning (RL) algorithms are the ability to learn from relatively little experience and the ability to learn policies that generalize to a range of problem specifications. In factored state spaces, one approach towards achieving both goals is to learn state abstractions, which only keep the necessary variables for learning the tasks at hand. This paper introduces Causal Bisimulation Modeling (CBM), a method that learns the causal relationships in the dynamics and reward functions for each task to derive a minimal, task-specific abstraction. CBM leverages and improves implicit modeling to train a high-fidelity causal dynamics model that can be reused for all tasks in the same environment. Empirical validation on manipulation environments and Deepmind Control Suite reveals that CBM's learned implicit dynamics models identify the underlying causal relationships and state abstractions more accurately than explicit ones. Furthermore, the derived state abstractions allow a task learner to achieve near-oracle levels of sample efficiency and outperform baselines on all tasks.

arXiv commentAccepted at AAAI24

References in corpus (26)

• Representation Learning with Contrastive Predictive Coding
• Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor
• Categorical Reparameterization with Gumbel-Softmax
• Reinforcement Learning with Deep Energy-Based Policies
• Model-Ensemble Trust-Region Policy Optimization
• dm_control: Software and Tasks for Continuous Control
• robosuite: A Modular Simulation Framework and Benchmark for Robot Learning
• Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning
• Tianshou: a Highly Modularized Deep Reinforcement Learning Library
• When to Trust Your Model: Model-Based Policy Optimization
• How to Train Your Energy-Based Models
• Sliced Score Matching: A Scalable Approach to Density and Score Estimation
• Deep Energy Estimator Networks
• Invariant Causal Prediction for Block MDPs
• Residual Energy-Based Models for Text
• On the model-based stochastic value gradient for continuous reinforcement learning
• ContactNets: Learning Discontinuous Contact Dynamics with Smooth, Implicit Representations
• The Primacy Bias in Deep Reinforcement Learning
• Decomposed Mutual Information Estimation for Contrastive Representation Learning
• Necessary and sufficient conditions for causal feature selection in time series with latent common causes
• Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning
• Causal Dynamics Learning for Task-Independent State Abstraction
• Denoised MDPs: Learning World Models Better Than the World Itself
• Action-Sufficient State Representation Learning for Control with Structural Constraints
• Joint Energy-based Model Training for Better Calibrated Natural Language Understanding Models
• CLOUD: Contrastive Learning of Unsupervised Dynamics