Visualizing MuZero Models
This work addresses the lack of interpretability in value equivalent models for reinforcement learning, which is an incremental improvement for researchers and practitioners using MuZero.
The paper tackled the problem of understanding what representations MuZero's value equivalent models learn by visualizing their latent representations, finding that action trajectory divergence between observation embeddings and internal state transitions could cause instability, and proposed two regularization techniques to stabilize performance.
MuZero, a model-based reinforcement learning algorithm that uses a value equivalent dynamics model, achieved state-of-the-art performance in Chess, Shogi and the game of Go. In contrast to standard forward dynamics models that predict a full next state, value equivalent models are trained to predict a future value, thereby emphasizing value relevant information in the representations. While value equivalent models have shown strong empirical success, there is no research yet that visualizes and investigates what types of representations these models actually learn. Therefore, in this paper we visualize the latent representation of MuZero agents. We find that action trajectories may diverge between observation embeddings and internal state transition dynamics, which could lead to instability during planning. Based on this insight, we propose two regularization techniques to stabilize MuZero's performance. Additionally, we provide an open-source implementation of MuZero along with an interactive visualizer of learned representations, which may aid further investigation of value equivalent algorithms.