Deep Reinforcement Learning with Vector Quantized Encoding

TL;DR

This paper introduces VQ-RL, a novel deep reinforcement learning framework that uses vector quantization to cluster state features, enhancing interpretability, robustness, and generalization of RL agents.

Contribution

The paper proposes a plug-and-play VQ-based clustering method for deep RL that improves interpretability and state similarity understanding, with new regularization techniques to enhance cluster separation.

Findings

VQ-RL produces tighter, better-separated state clusters.

VQ-RL improves interpretability of RL models.

VQ-RL enhances robustness and generalization in simulations.

Abstract

Human decision-making often involves combining similar states into categories and reasoning at the level of the categories rather than the actual states. Guided by this intuition, we propose a novel method for clustering state features in deep reinforcement learning (RL) methods to improve their interpretability. Specifically, we propose a plug-and-play framework termed \emph{vector quantized reinforcement learning} (VQ-RL) that extends classic RL pipelines with an auxiliary classification task based on vector quantized (VQ) encoding and aligns with policy training. The VQ encoding method categorizes features with similar semantics into clusters and results in tighter clusters with better separation compared to classic deep RL methods, thus enabling neural models to learn similarities and differences between states better. Furthermore, we introduce two regularization methods to help increase the separation between clusters and avoid the risks associated with VQ training. In simulations, we demonstrate that VQ-RL improves interpretability and investigate its impact on robustness and generalization of deep RL.