Planning in a recurrent neural network that plays Sokoban

TL;DR

This paper investigates how a recurrent neural network trained on Sokoban internally represents planning, revealing causal plan structures, pacing behaviors, and extending capabilities to larger puzzles, thus advancing understanding of neural planning mechanisms.

Contribution

It uncovers the internal causal plan representations in an RNN trained on Sokoban and demonstrates how these representations enable solving larger, out-of-distribution puzzles.

Findings

RNN predicts actions 50 steps ahead

The quality of planning increases over initial steps

The RNN exhibits pacing behavior incentivized by training

Abstract

Planning is essential for solving complex tasks, yet the internal mechanisms underlying planning in neural networks remain poorly understood. Building on prior work, we analyze a recurrent neural network (RNN) trained on Sokoban, a challenging puzzle requiring sequential, irreversible decisions. We find that the RNN has a causal plan representation which predicts its future actions about 50 steps in advance. The quality and length of the represented plan increases over the first few steps. We uncover a surprising behavior: the RNN "paces" in cycles to give itself extra computation at the start of a level, and show that this behavior is incentivized by training. Leveraging these insights, we extend the trained RNN to significantly larger, out-of-distribution Sokoban puzzles, demonstrating robust representations beyond the training regime. We open-source our model and code, and believe the neural network's interesting behavior makes it an excellent model organism to deepen our understanding of learned planning.