Stagewise Reinforcement Learning and the Geometry of the Regret Landscape

TL;DR

The paper extends singular learning theory to deep reinforcement learning and proves that the concentration of the generalized posterior over policies is governed by the local learning coefficient (LLC), an invariant of the regret geometry. It predicts Bayesian phase transitions where simpler, higher-regret policies can be favored at finite data, a phenomenon formalized through a free-energy tradeoff between regret and complexity. The authors validate these predictions in a gridworld exhibiting stagewise policy development, observing opposing staircases of decreasing regret and rising LLC during phase transitions, with LLC detecting changes even when behavior is indistinguishable in regret. This work links the geometry of the regret landscape to Bayesian learning and SGD dynamics, offering a principled lens on complexity biases, policy emergence, and potential alignment implications in RL. The results lay groundwork for applying singular learning tools to more complex RL settings and to alignment-relevant phenomena in AI safety contexts.

Abstract

Singular learning theory characterizes Bayesian learning as an evolving tradeoff between accuracy and complexity, with transitions between qualitatively different solutions as sample size increases. We extend this theory to deep reinforcement learning, proving that the concentration of the generalized posterior over policies is governed by the local learning coefficient (LLC), an invariant of the geometry of the regret function. This theory predicts that Bayesian phase transitions in reinforcement learning should proceed from simple policies with high regret to complex policies with low regret. We verify this prediction empirically in a gridworld environment exhibiting stagewise policy development: phase transitions over SGD training manifest as "opposing staircases" where regret decreases sharply while the LLC increases. Notably, the LLC detects phase transitions even when estimated on a subset of states where the policies appear identical in terms of regret, suggesting it captures changes in the underlying algorithm rather than just performance.

Figures (14)

• Figure 1: Opposing staircases of regret and complexity: Comparison of the regret and complexity (as estimated by the local learning coefficient) across training for an agent optimized to follow shortest paths towards the goal location (cheese), which was located in the top-left corner in $\sim\frac{1}{3}$ of sampled episodes and uniformly across all locations in the remaining episodes. Phases are indicated along with visualizations of policies at the indicated checkpoints (i.e. policy gradient steps). In the first phase the agent moves up and left with equal probability; in the second phase it moves deterministically towards the top-left region, passing through the goal if possible, and finally in the third phase it moves directly toward the goal. For each location $s$, the arrow drawn points in the direction of the vector $v_s = [\pi(\rightarrow | s, \varnothing) - \pi(\leftarrow | s, \varnothing), \pi(\uparrow | s, \varnothing) - \pi(\downarrow | s, \varnothing)]$ and has size proportional to $||v_s||_1$ representing the expected direction of movement if the agent spawns in cell $s$. Policy representation taken from mini2023understanding. Phase transitions are associated with rapid decreases in regret and rapid increases in the LLC estimate.
• Figure 2: Two possible states the environment can be in, together with the observation seen by the agent.
• Figure 3: A series of plots illustrating the model transitions into phases $\pi^1$, $\pi^{2b}$ and $\pi^3$ respectively, as a function of mixing parameter $\alpha$ and discount rate $\gamma$. We report the number of checkpoints (policy gradient steps) since training began before we enter each phase. The red line in phase 3 represents when training runs were terminated.
• Figure 4: Distribution of average LLC estimates across phases. The LLC has been estimated on distribution with $\alpha=0.68$ and $\gamma=0.975$.
• Figure 5: LLC and regret over training for a model trained with $\alpha=0.68$ and $\gamma=0.975$. The LLC has been computed with $\alpha=0$ and $\gamma=0.98$. We show two regret curves, the orange one with the same $\alpha$ and $\gamma$ as was used for training, and the green one with the same $\alpha$ and $\gamma$ as was used for LLC estimation.

Theorems & Definitions (59)

• Definition 2.1
• Theorem 2.3
• proof
• Remark 3.1
• Remark 3.2: Off-distribution LLC estimation
• Example E.1
• Remark E.2
• Definition E.3
• Definition E.4
• Definition E.5