Neural Discrete Abstraction of High-Dimensional Spaces: A Case Study In Reinforcement Learning

Abstract

We employ Neural Discrete Representation Learning to map a high-dimensional state space, made up from raw video frames of a Reinforcement Learning agent's interactions with the environment, into a low-dimensional state space made up from learned discrete latent representations. We show experimentally that the discrete latents learned by the encoder of a Vector Quantized Auto-Encoder (VQ-AE) model trained to reconstruct the raw video frames making up the high-dimensional state space, can serve as meaningful abstractions of clusters of correlated frames. A low-dimensional state space can then be successfully constructed, where each individual state is a quantized vector encoding representing a cluster of correlated frames of the high-dimensional state space. Experimental results for a 3D navigation task in a maze environment constructed in Minecraft demonstrate that this discrete mapping can be used in addition to, or in place of, the high-dimensional space to improve the agent's learning performance.