Learning Discrete State Abstractions With Deep Variational Inference
This work addresses the challenge of abstraction in reinforcement learning for domains with high-dimensional states, such as robotics, by providing a method to learn discrete models that improve planning efficiency, though it appears incremental as it builds on existing bisimulation and information bottleneck concepts.
The paper tackles the problem of learning discrete state abstractions for efficient sequential decision-making in large state spaces by proposing a deep variational inference method that maps high-dimensional states to discrete representations using an action-conditioned hidden Markov model, enabling efficient planning for unseen goals in multi-goal reinforcement learning settings, with testing in robotic manipulation domains and comparisons to previous model-based approaches.
Abstraction is crucial for effective sequential decision making in domains with large state spaces. In this work, we propose an information bottleneck method for learning approximate bisimulations, a type of state abstraction. We use a deep neural encoder to map states onto continuous embeddings. We map these embeddings onto a discrete representation using an action-conditioned hidden Markov model, which is trained end-to-end with the neural network. Our method is suited for environments with high-dimensional states and learns from a stream of experience collected by an agent acting in a Markov decision process. Through this learned discrete abstract model, we can efficiently plan for unseen goals in a multi-goal Reinforcement Learning setting. We test our method in simplified robotic manipulation domains with image states. We also compare it against previous model-based approaches to finding bisimulations in discrete grid-world-like environments. Source code is available at https://github.com/ondrejba/discrete_abstractions.