Deep Reinforcement Learning with Vector Quantized Encoding

Liang Zhang, Justin Lieffers, Pavithra Shivanna, Adarsh Pyarelal

Keywords: Reinforcement Learning, Interpretability, Clustering

Abstract: Human decision-making often involves combining similar states into categories and reasoning at the level of the categories rather than the actual states. Expanding on this notion, we introduce a novel approach to augment the interpretability, robustness, and generalization of deep reinforcement learning (RL) models through state feature clustering. Our method, termed \emph{Vector Quantized Reinforcement Learning} (VQ-RL), integrates an auxiliary classification task using vector quantized (VQ) encoding into conventional RL pipelines, improving the model's capability to group similar states into clusters with enhanced separation. Additionally, we propose two regularization techniques to bolster cluster separation and alleviate potential risks associated with VQ training. Our simulations demonstrate that VQ-RL enhances our understanding of the internal space of deep RL, as well as the robustness and generalization capabilities of state-of-the-art RL methods across various domains.

Submission Number: 12