Exploring the Stratified Space Structure of an RL Game with the Volume Growth Transform
Justin Curry, Brennan Lagasse, Ngoc B. Lam, Gregory Cox, David Rosenbluth, Alberto Speranzon

TL;DR
This paper investigates the geometric structure of the embedding space in a transformer-based RL agent, revealing it as a stratified space with variable local dimensions, which correlates with the agent's behavior and environmental complexity.
Contribution
It extends the volume growth transform analysis from LLMs to RL models, demonstrating the stratified nature of embedding spaces and linking local dimension variations to agent strategies and game complexity.
Findings
Embedding space is a stratified space, not a manifold.
Local dimension varies with agent behavior and environmental complexity.
Dimension distribution may serve as a complexity indicator for RL environments.
Abstract
In this work, we explore the structure of the embedding space of a transformer model trained for playing a particular reinforcement learning (RL) game. Specifically, we investigate how a transformer-based Proximal Policy Optimization (PPO) model embeds visual inputs in a simple environment where an agent must collect "coins" while avoiding dynamic obstacles consisting of "spotlights." By adapting Robinson et al.'s study of the volume growth transform for LLMs to the RL setting, we find that the token embedding space for our visual coin collecting game is also not a manifold, and is better modeled as a stratified space, where local dimension can vary from point to point. We further strengthen Robinson's method by proving that fairly general volume growth curves can be realized by stratified spaces. Finally, we carry out an analysis that suggests that as an RL agent acts, its latent…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
