Abstract
In this work, we explore the structure of the embedding space of atransformer model trained for playing a particular reinforcement learning (RL)game. Specifically, we investigate how a transformer-based Proximal PolicyOptimization (PPO) model embeds visual inputs in a simple environment where anagent must collect "coins" while avoiding dynamic obstacles consisting of"spotlights." By adapting Robinson et al.'s study of the volume growthtransform for LLMs to the RL setting, we find that the token embedding spacefor our visual coin collecting game is also not a manifold, and is bettermodeled as a stratified space, where local dimension can vary from point topoint. We further strengthen Robinson's method by proving that fairly generalvolume growth curves can be realized by stratified spaces. Finally, we carryout an analysis that suggests that as an RL agent acts, its latentrepresentation alternates between periods of low local dimension, whilefollowing a fixed sub-strategy, and bursts of high local dimension, where theagent achieves a sub-goal (e.g., collecting an object) or where theenvironmental complexity increases (e.g., more obstacles appear). Consequently,our work suggests that the distribution of dimensions in a stratified latentspace may provide a new geometric indicator of complexity for RL games.