Interpreting Emergent Planning in Model-Free Reinforcement Learning

TL;DR

This work provides mechanistic evidence that a model-free reinforcement learning agent can learn to plan in Sokoban by revealing internally formed, square-level CA (Agent Approach Direction) and CB (Box Push Direction) concepts via linear probes. The authors show that these concepts form iterative, bidirectional plans that resemble parallelized search, can be refined with additional test-time compute, and causally influence behavior through targeted interventions. By linking concept representations to planning behavior and demonstrating their emergence during training, the study argues for an implicit world model-like mechanism within a generic, model-free agent. The findings advance understanding of emergent planning in RL and offer a methodology for diagnosing planning-like processes in other architectures and environments, with implications for scalable reasoning in RL systems.

Abstract

We present the first mechanistic evidence that model-free reinforcement learning agents can learn to plan. This is achieved by applying a methodology based on concept-based interpretability to a model-free agent in Sokoban -- a commonly used benchmark for studying planning. Specifically, we demonstrate that DRC, a generic model-free agent introduced by Guez et al. (2019), uses learned concept representations to internally formulate plans that both predict the long-term effects of actions on the environment and influence action selection. Our methodology involves: (1) probing for planning-relevant concepts, (2) investigating plan formation within the agent's representations, and (3) verifying that discovered plans (in the agent's representations) have a causal effect on the agent's behavior through interventions. We also show that the emergence of these plans coincides with the emergence of a planning-like property: the ability to benefit from additional test-time compute. Finally, we perform a qualitative analysis of the planning algorithm learned by the agent and discover a strong resemblance to parallelized bidirectional search. Our findings advance understanding of the internal mechanisms underlying planning behavior in agents, which is important given the recent trend of emergent planning and reasoning capabilities in LLMs through RL

Figures (63)

• Figure 1: Examples of the DRC agent internally forming plans to push boxes to targets. A purple arrow on a square means that a linear probe decodes that the agent plans to push a box off of that square in the associated direction. No arrow on a square means that the probe decodes that agent does not plan to push a box off of that square. (A) The agent evaluates a naively-appealing route, concludes it is infeasible, and forms a longer alternate path. (B) The agent adapts its plan and changes the target it plans to push the left-most box to. (C) The agent extends part of its plan backward from a target. (D) The agent extends part of its plan forward from a box. (E) The agent extends many parts of its plan in parallel. We provide further examples in Appendices \ref{['ap-planform1']}-\ref{['ap-planform5']}.
• Figure 2: Pixel and symbolic representations of a Sokoban board.
• Figure 3: Examples of the classes assigned to the squares of a Sokoban board over 6 transitions (from left to right) by the concepts 'Agent Approach Direction' ($C_{\text{A}}$) and 'Box Push Direction' ($C_{\text{B}}$). An arrow corresponds to the assignment of the associated directional class. The lack of an arrow in a square indicates the assignment of the class NEVER.
• Figure 4: Macro F1s achieved by probes when predicting $C_{\text{A}}$ and $C_{\text{B}}$ using the cell state at each layer, or, for the baseline probes, using the observation. Error bars show $\pm1$ standard deviation.
• Figure 5: Examples of internal plans computed by the agent. An internal plan corresponds to the agent's combined square-level representations of $C_{\text{A}}$ and $C_{\text{B}}$. That is, an internal plan corresponds to the classes the agent represents these concepts as mapping squares of observed boards to. These internal plans are decoded from the agent's final layer cell state by a 1x1 probe. Teal and purple arrows respectively indicate the agent expects to next step on to, or push a box off, a square in the associated direction. No arrow indicates the agent does not plan to step onto, or push a box off, a square again. Further examples of internal plans are given in Figures \ref{['fig:p_corr_l0']}, \ref{['fig:p_corr_l1']} and \ref{['fig:p_corr_l2']} in Appendix \ref{['ap-intplan']}.