Strata–trajectory correspondence in RL latent spaces

Establish that distinct strata in the stratified token-embedding latent space of the Transformer-XL-based Proximal Policy Optimization agent trained on the Searing Spotlights two-coins environment correspond to different state–action trajectories, and determine whether increases in local dimension occur when the agent’s latent representation lies near intersections of strata.

Background

The paper adapts the Volume Growth Transform to analyze the token embedding space of a transformer-based PPO agent trained on a visual reinforcement learning game (a two-coins variant of Searing Spotlights). The authors show the latent space departs from manifold and fiber-bundle behavior and is better modeled as a stratified space with varying local dimensions across tokens.

By examining local dimension along trajectories during gameplay, they observe spikes in local dimension near coin collections and in scenes with greater environmental complexity or policy ambiguity. These findings motivate the conjecture that strata encode different state–action trajectories and that intersections of strata are associated with increased local dimension.