Planning in a recurrent neural network that plays Sokoban

## Abstract

Abstract How a neural network (NN) generalizes to novel situations depends on whether it has learned to select actions
heuristically or via a planning process. Guez et al. [ 2019 , “An investigation of model-free
planning”] found that a recurrent NN (RNN) trained to play Sokoban appears to plan, with extra
computation steps improving the RNN’s success rate. We replicate and expand on their behavioral analysis, finding the
RNN learns to give itself extra computation steps in complex situations by “pacing” in cycles. Moreover, we train
linear probes that predict the future actions taken by the network and find that intervening on the hidden state
using these probes controls the agent’s subsequent actions. Leveraging these insights, we perform model surgery,
enabling the convolutional NN to generalize beyond its 10 × 10 10 10 10\times 10 10 × 10 architectural limit to arbitrarily sized inputs.
The resulting model solves challenging, highly off-distribution levels. We open-source our model and code, and believe
the neural network’s small size (1.29M parameters) makes it an excellent model organism to deepen our understanding of learned
planning.

## 1 Introduction

In many tasks, the performance of both humans and sufficiently advanced neural networks (NNs) improves with more reasoning:
whether by giving a human more time to think before taking a chess move, or by prompting or training a large language model (LLM) to reason step
by step [Kojima et al., , OpenAI, ].

Among other reasoning capabilities, goal-oriented reasoning is particularly relevant to AI alignment. So-called
“mesa-optimizers” – AIs that have learned to pursue goals through internal reasoning [Hubinger et al., ] – may
internalize goals different from the training objective, leading to harmful misgeneralization
[Di Langosco et al., , Shah et al., ].

Within goal-oriented reasoning, this work distinguishes between *plans* and *search algorithms*. A *plan*
is a sequence of (perhaps abstract) actions which an agent intends to take. A *search algorithm* is an algorithm
that considers many plans, evaluates them according to their predicted outcomes, and picks the plan evaluated as best. Search
algorithms perform better the more plans they can evaluate and, with enough compute, generalize very well to novel
problem instances [Russell and Norvig, ].

This work investigates a recurrent neural network (RNN) that plays Sokoban [Guez et al., ], a challenging puzzle game that remains a benchmark for planning algorithms [Peters et al., ] and reinforcement learning [Chung et al., ]. Using
linear location-wise probes from concurrent work [Bush et al., ], we show that its activations represent a plan in a way which causes the RNN’s
actions.

The behavior of the RNN suggests that it also performs *search*. We find that the RNN sometimes takes actions which
return to previous states of the environment. This overwhelmingly happens at the beginning of episodes and can be
substituted by no-ops, suggesting the RNN deliberately goes in cycles to get more time to think. Giving the
RNN time to think with no-ops [Guez et al., ] disproportionately helps with levels that are difficult
(Figure 1), long, or that require waiting many steps before the first reward
(Figure 3).

Using insights from the plan probes and inspecting activations, we perform model surgery to uncover latent
generalization capabilities of the RNN. The RNN is trained only on $10\times 10$ levels with 4 boxes (and its output
MLP can only handle $10\times 10$ levels), but we are able to make it play on larger, nontrivial levels with more boxes
(Figure 2).

Altogether, evidence from probes and behavior show that the RNN causally represents a *plan*. The evidence highly
suggests that the RNN implements an online *search algorithm* which scales with computation, although we cannot
rule out multi-step heuristics.

Consequently, we believe this RNN is an excellent model organism for understanding planning in sophisticated neural
networks.
Its behavior is complex enough to be interesting and its size (1.29M parameters) is small enough to be
tractable to reverse engineer. This work shows that supervised probes and sparse autoencoders [Huben et al., , Bricken et al., , Gao et al., ] are not enough to fully do so, but it is likely that
other mechanistic interpretability techniques such as circuit analysis [Olah et al., , Geiger et al., , Conmy et al., ] will reveal more about the algorithm which produces plans.

Figure: Figure 2: The $\text{DRC}(3,3)$ is limited to $10\times 10$ grids by a fixed-dimension MLP. With model surgery, the ConvLSTM backbone generalizes to much more challenging settings than the $10\times 10$ training set grids. Left: XSokoban-31, which the $\text{DRC}(3,3)$ solves with 0 thinking steps. Right: XSokoban-29, which the $\text{DRC}(3,3)$ solves but only after 128 thinking steps. More examples in Appendix B.
Refer to caption: x3.png

## 2 Setting up the test subject

We train an agent closely following the setup from Guez et al. [], using the IMPALA [Espeholt et al., ] reinforcement learning (RL) algorithm with ’s Deep Repeating ConvLSTM (DRC)
recurrent architecture. We also train a ResNet baseline. We open-source both the trained networks
(^1^11[https://github.com/AlignmentResearch/learned-planner](https://github.com/AlignmentResearch/learned-planner))
and training code
(^2^22[https://github.com/AlignmentResearch/train-learned-planner](https://github.com/AlignmentResearch/train-learned-planner))
. For more details about architectures and training, see Appendix A.

#### DRC ⁢ ( D , N ) DRC 𝐷 𝑁 \text{DRC}(D,N) DRC ( italic_D , italic_N ) architecture.

Most of this paper focuses on the behavior and representations of a $\text{DRC}(3,3)$ neural network
[Guez et al., ]. The core of this network is a $D$-layer ConvLSTM [Shi et al., ], which is
repeatedly applied $N$ times per environment step. The output of the last ConvLSTM layer (the $D$th layer) is fed to the
input at the next tick, effectively giving the network $D\cdot N$ layers of sequential computation to output the next
action. We take $D=3,N=3$.

A linear combination of the mean- and max-pooled activations of the ConvLSTM is also fed to the next step, thus letting
it communicate quickly across the receptive field (known as pool-and-inject). An encoder block consisting of
two $4\times 4$ convolutions processes the input, which is also fed to each ConvLSTM layer. At the end, an MLP with 256
hidden units transforms the flattened ConvLSTM outputs into the policy (actor) and value function (critic) heads.

#### Dataset.

Sokoban is a grid puzzle game with walls, floors, movable boxes, and target tiles. The player’s
goal is to push all boxes onto target tiles while navigating walls. We use the Boxoban dataset [Guez et al., ],
consisting of $10\times 10$ procedurally generated Sokoban levels, each with 4 boxes and targets. The edge tiles are
always walls, so the playable area is $8\times 8$. Boxoban separates levels into train, validation and test sets, with
three difficulty levels: unfiltered, medium, and hard. Guez et al. [] generated these sets by filtering
levels that cannot be solved by progressively better-trained DRC networks, so easier sets occasionally contain difficult
levels. In this paper, we use the unfiltered-train (900k levels) set to train networks. To evaluate them, we use the unfiltered-test
(1k levels)(^3^33We use unfiltered-test rather than unfiltered-validation so the numbers are directly comparable to Guez et al. [].), medium-validation (50k levels), and hard ($\sim$3.4k levels) sets, which do not
overlap. To test $\text{DRC}(3,3)$ generalization to different sizes, we use the levels collected by Þorsteinsson []
(see Appendix B).

#### Environment.

The observations are $10\times 10$ RGB images, normalized by dividing each component by 255.
Each type of tile is represented by a pixel of a different color [Schrader, ], an example is in
Figure 12. The player can move in cardinal directions (Up, Down, Left, Right). The reward is
-0.1 per step, 1 for placing a box on a target, -1 for taking it off, and 10 for finishing the level by placing all of
the boxes. The time limit for evaluation is 120 steps, though large levels in Section 5.1 use
1000 steps.

## 3 Behavioral evidence of search

### 3.1 Effects of thinking time: non-myopic network

First, we examine the effect of thinking time: introducing no-ops at the beginning of an episode, while letting the RNN
process its hidden state. As in Guez et al. [], adding 6 extra thinking steps lets the $\text{DRC}(3,3)$
solve an extra $4.7\%$ of levels, slightly decreasing when going up to 16 steps
(Figure 5, left).

We find that the extra thinking time disproportionately helps with more difficult levels, as measured by longer optimal solutions (Figure 3, right). We also checked whether thinking time correlates with the number of nodes expanded by an A* search, but it does not (Figure 14).

Thinking time helps varying amounts depending on training time and level difficulties: Figure 1 (bottom left) shows that most of the planning effect steeply appears during the first 70M steps of learning. Over the remaining training time, the RNN benefits more from thinking steps for the hard levels, but less for the medium levels; indicating perhaps that it learns better heuristics for them.

We also find that thinking time disproportionately helps with levels in which the agent can greedily put a box to target early, but where doing so is detrimental to solving the level. Figure 3 (left, b) shows this: with 0 thinking steps, the first box is pushed to target on average 4 steps too early; compared to the solution reached at 6 thinking steps. Some of this is driven by the forced thinking steps preventing the $\text{DRC}(3,3)$ from taking a catastrophic action in the first few moves. However, the average time to first box is more than 6, even for 0 thinking steps, so there must also be many cases in which the 6 extra thinking steps give the NN time to find a better solution, rather than just preventing it from greedily pushing the first box. It must also be less greedy for boxes 2-3.

Figure: Figure 3: Left: Average time step to place each box $B_{i}$ on target for different numbers of thinking steps. (a) Averages across all levels where the box $B_{i}$ is placed on target by DRC with $N$ thinking steps. (b) Averages for all levels solved by $6$ thinking steps but not solved by 0 0 thinking steps. More thinking steps makes the DRC avoid greedy strategies in favor of long-term return. Right: Average optimal solution length of levels grouped by the number of thinking steps at which the level is first solved. Levels that take longer to solve tend to be harder. NS stands for “not solved”.
Refer to caption: x5.png

### 3.2 Agent “paces” to get more computation

On occasion, the DRC exhibits a curious behavior: the agent “paces” in a cycle, returning to the same location, without touching any box. The per-step penalty makes this behavior naively sub-optimal, because it does not advance the puzzle state.

Since the $\text{DRC}(3,3)$’s performance improves with thinking steps, could the agent be using pacing to gain computation time to improve its plan and so avoid irreversible actions which may block the solution? The following evidence suggests the agent has developed this meta-strategy:

#### The agent paces mostly at the start of episodes.

In some levels, a single sub-optimal step can lock the level, which makes calculating a valid plan in early steps beneficial. Figure 4 (left) shows most cycles start in the first 5 steps of the episode, and (middle) forcing the model to think for six steps eliminates about 75% of these early cycles.

#### Thinking steps almost entirely substitute pacing

We run the $\text{DRC}(3,3)$ on medium-validation levels and record when and where cycles occur. We merge all cycles which overlap with other cycles, for a total of $13\,702$ cycles.

We then make the DRC take $N$ thinking steps just before it would start a cycle of length $N$. As shown in Figure 4 (right), these thinking steps replace the cycles: comparing the cycle and the substitution, the DRC follows the same trajectory for at least $60\%$ of the levels and for a minimum of $30$ steps(^4^44For context, the median solution length for train-unfiltered is exactly 30 steps.).

If our hypothesis that cycles are utilized for thinking holds true, we would expect to see uninterrupted periods of action after each cycle. In the $N$ steps after an $N$-step cycle concludes, $98.8\%$ of trajectories have no cycles. When substituting them for NOOPs, this becomes $82.4\%$, lower but still much higher than $0\%$; indicating that thinking steps largely remove the NN’s perceived need for cycles.

#### Comparing DRC ⁢ ( 3 , 3 ) DRC 3 3 \text{DRC}(3,3) DRC ( 3 , 3 ) to ResNet

It is possible that the DRC is pacing because it is not optimized enough.
The ResNet is much worse than the $\text{DRC}(3,3)$ at Sokoban, but it is optimized enough to almost never perform cycles: since it has no memory, any cycles would make it stuck. So the amount of training is definitely enough to remove cycles if they have a big impact, and we believe it is likely to also remove cycles if they have a moderate negative impact.

It could also be that the DRC is computing a multi-step heuristic (e.g. one that depends on the convolutional
receptive field). Against this, the pool-and-inject layer in the DRC makes the receptive field less important for communication, and the ResNet (which performs worse) has a much larger receptive field per forward pass. Together, this is weak evidence that receptive-field heuristics are not the main mechanism.

Heuristics may still play a role. For example, due to the -0.1 per-step penalty and discount rate, it is rational for the training process to make the RNN place boxes greedily (hence the increasing time-to-box in Figure 3, middle). It could be that the planning effect disappearing for medium levels (Figure 1, left) is due to better heuristics about when to be greedy and when to take time to think.

Figure: Figure 4: Left: Histogram of cycle start times on the medium-difficulty validation levels. Middle: Total number of cycles the agent takes in the first 5 steps across all episodes in medium-difficulty validation levels with $N$ initial thinking steps. Right: on all levels, we replace $N$-length cycles with $N$ thinking steps and check the proportion of trajectories which are equal between these two treatments, for $x$ steps after the cycle.
Refer to caption: x6.png

## 4 The DRC ⁢ ( 3 , 3 ) DRC 3 3 \text{DRC}(3,3) DRC ( 3 , 3 ) causally represents its plan

In order to improve the plan with search, as the previous section suggests, the agent must have some internal representation of the plan. The absence of such representation would force us to discard the theory that the agent is doing search.

This section presents probes which predict the future actions of the $\text{DRC}(3,3)$ and other various future features of the environment. Of these, the one which predicts the future direction *boxes* will move in has a strong *causal effect* in the actions of the RNN. Following the standard of evidence in Li et al. [], we declare this conclusive evidence that the $\text{DRC}(3,3)$ represents and uses plans. The spatial structure of the probes and some of the targets are due to Bush et al. [].

We also observe the plan grow during thinking steps and, by intervening on the probes, are able to lock the RNN into a plan when there are two options of equal value.
However in most cases, if we stop intervening on the RNN activations with a suboptimal plan, the RNN comes up with a better plan online and follows it. We have some evidence for the $\text{DRC}(3,3)$ considering multiple plans: for example, in the level in Figure 1 (right) the box probe first considers moving box 4 down and right (shown in red), but then rethinks that and ends up taking different actions. However, we have been unable to find multiple *simultaneous* representations of plans, or anything that would decide between them such as the value of the plans. In all, this section provides very little evidence in favor of *search*, but conclusive evidence for a causal *plan*.

Figure: Figure 5: Success rate of the RNN, and plan quality as measured by box-direction probes (Section 4.1), all increase with thinking steps. We measure plan quality by 1) summing the length of chains of probe-predicted directions, which start in boxes, and 2) simply counting the number of squares in which the box-direction probe predicts something. The F1 score of the plan probe also increases over time, suggesting that the DRC builds the plan in that time. The right plot also includes plan quality evaluation on the three sequential ticks per environment step of the DRC(3, 3) network.
Refer to caption: x7.png

### 4.1 Probe methodology: training, intervention and targets

#### Training.

We train logistic or linear regression probes with L1 decay to predict various features of the environment from the agent’s activations. We train on two types of inputs. 1) Grid-wise inputs [Bush et al., ]: each square in the $10\times 10$ grid is a different data point, potentially with a different label depending on location and timestep. The input is the $64$-dimensional LSTM state $(h,c)$ at a square, usually concatenated for the 3 layers. 2) Global inputs: Each time step produces only one data point. The probes take as input the 64 channels of all of the squares in the 10x10 grid, resulting in a 6400-dimensional input for each layer. In both cases, the per-layer activations are concatenated for probes that are trained on multiple layers.

The dataset consists of states collected by evaluating the $\text{DRC}(3,3)$ on the hard Boxoban levels, excluding the first 5 steps of each episode because the plan is still forming. For multi-class probes, training and F1 scores are computed as one-vs-all: the presence of absence of a particular class is the probe label. We search the best learning rate and L1 decay with grid-search by evaluating the F1 on a validation split of $20\%$ of timesteps from the hard levels. We take the first 1000 medium-validation levels as our test set to report the results (Table 1).

#### Causal intervention.

Given a probe that is predictive of future behavior, we edit the agent’s activations such that the probe predicts the behavior we desire to induce. To the extent that the agent carries out the edited behavior, we can conclude that it executes a plan encoded in the linear direction of the probe.

For a linear probe on the activations $h$ with parameter vector $p$, we intervene in order to increase the logits for a particular behavior, given by $p\cdot h$ (we omit a scalar intercept for brevity). Unlike previous work in steering vectors for LLMs [Turner et al., , Rimsky et al., , Li et al., ], we found that simply adding a perturbation equal to a fixed multiple of $p$ caused the policy to degenerate and choose random actions. We instead add an adaptively-scaled multiple of $p$ to the activations, chosen to minimize the change in the activations while still causing the probe’s logit to be greater than a desired amount $\alpha$. More precisely, we find $\operatorname*{arg\,min}_{h^{\prime}}\|h^{\prime}-h\|_{2}^{2}$, with the constraint that $h^{\prime}\cdot p>\alpha$. The solution to this is $h^{\prime}=h+{\hat{p}}\max(0,\alpha-(h\cdot p))$, where ${\hat{p}}=p/\|p\|^{2}_{2}$.

#### Targets.

We choose targets that plausibly encode the steps of a plan. Most of these labels are for grid-wise probes, but the pacing and value probes are global. The direction probes are very similar to Bush et al. [].

- •
Agent-Directions probe: This probe has 5 outputs: NV (No Visit), UP, DOWN, LEFT, RIGHT. The probe takes as input the 64-dim hidden state activations at a square $(x,y)$ at timestep $t_{i}$. If the agent visits the square $(x,y)$ at a future timestep $t>t_{i}$, then the probe predicts the corresponding direction that the agent takes from that square. For squares that are visited multiple times in the future, the direction corresponding to the first future visit is taken as the target. If a square is never visited in a future timestep, the probe predicts NV. There is one single multinomial probe, for all time steps and locations $(x,y)$.
- •
Boxes-Directions probe: This probe is similar to the above probe, except that it is trained to predict the direction in which *any* of the four boxes will move at a given square in the future. The rationale for this probe is that the boxes, and constraints on their movement given by obstacles, are the main difficulty with Sokoban puzzles.
- •
Next-Box probe: Predicts $1$ on the square of the box that the agent will move next and 0 0 for every other square.
- •
Next-Target probe: Predicts $1$ on the square of the target that the agent will put a box in next, and 0 0 for every other square.
- •
Pacing probe: The global label is 1 if the agent is currently in a cycle, and 0 otherwise.
- •
Value probe: The global label is the numerical value that the critic head outputs.

**Table 1: Causal and predictive probe results. Confidence is one of the mean estimator percentiles $[2.5\%,97.5\%]$, whichever is furthest from the mean, estimated using 1000 bootstrap resamples. The average causal probe uses $24\text{k}$ data points for evaluation, and the best-case probe uses $8\text{k}$. The data points are sampled from the first 1000 medium validation levels.**
| Probe target | Predictive Probe Results (a) | Causal Probe Results (b) |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  |  | Sparse Probes | Causal Effect |  |  |  |
|  | Best F1 | F1 | $L_{0}$-norm | $\alpha$ | Average | Best-case |
| Box-directions | $86.4\pm 0.1$ | $73.1\pm 0.1$ | $63$ | 30 | $43.7\pm 0.6$ | $77.8\pm 0.9$ |
| Agent-directions | $72.3\pm 0.1$ | $61.3\pm 0.1$ | $121$ | 10 | $7.1\pm 0.3$ | $20.7\pm 0.7$ |
| Next box | $74.2\pm 0.4$ | $69.7\pm 0.5$ | $51$ | 40 | $5.5\pm 1.0$ | $15.1\pm 2.5$ |
| Next target | $54.3\pm 0.5$ | $44.0\pm 0.5$ | $32$ | 30 | $4.6\pm 0.8$ | $13.2\pm 2.0$ |
| Pacing | $31.0\pm 1.8$ | $31.0\pm 1.5$ | $5$ | — | — | — |

### 4.2 Probe evaluation and causality

#### Probe predictive power.

Most probes are quite predictive, as seen in Table 1.
The exception is the pacing probe, with $F_{1}=31.0\%$, which is not much better than the constant 1 probe, which has
$F_{1}=12.8\%$. This leads us to conclude that the $\text{DRC}(3,3)$ does not represent whether it is in a cycle or not. See
Figure 1 (right) and Figure 6 for visualizations of agent-direction and box-direction probes(^5^55The supplementary material provides visualization videos of all the probes across several levels.).

For the value function probe, we compute the fraction of variance explained $R^{2}$. If we train a global probe, $R^{2}$
is very high: $97.7\%$–$99.7\%$ depending on the layer. However, grid-wise probes obtain much worse but still passable
results: $41.0\%$–$79.2\%$ depending on layer. We visually checked whether the grid-wise probe reads off the values of
different plans, but could not find any such pattern.

Almost all the performance of the global probe is recovered by training on the mean-pooled inputs: $95.2\%$–$99.5\%$.
It is likely that the global and mean-pooled value probes are indirectly counting the number of squares the agent will
step on, which almost fully determines the value, and we know is possible due to the future-direction probes.

#### Only the Box-Directions probe is strongly causal.

We measure how causal the probes are by intervening at every step on the hidden state $h,c$ at the layer the probe was trained on, following the procedure in section 4.1 and checking
whether the agent follows that action. We report the result in Table 1(b) by performing a grid-search over the strength of intervention $\alpha$ and picking the best value. Despite our adaptive scaling of the intervention, a high value of $\alpha$ can still disturb the agent such that it starts performing random actions whereas a low value of $\alpha$ may not be sufficient to cause the agent to follow the intended behavior. Most of the probes are not
causal: only the “boxes directions” probe (used in Figure 1) is causal while the “agent-direction” probe mildly causal.

Even then, we have observed that it can only change boxes directions when they *would not deadlock* the level in
a naive way. That is, if you try to make the model push a box onto a wall, such that the box would not be able to reach a
target anymore, it does not. For this reason, we introduce the *Best-case causal effect* in
Table 1(b): try all three directions that are not the actual direction that the box follows, and
count the probe as “causal” if it works for *any* of them (e.g. because they do not naively deadlock the level).

#### The box-direction probe’s coverage increases during thinking time.

Figure 5 (middle) show that the amount of squares in which the probe predicts a direction, be they lone or linked to a box, grows over thinking steps. Its description of what the box movements will be also greatly improves after a few steps of thinking Figure 5 (right).

### 4.3 Bistable and unstable plans in toy environments

The causal intervention results from Table 1 show that the box-directions probe is much more causal than the agent-directions probe. We further study this in an isolated setting. Figure 6 (left) shows a level with only a single box next to a target, where the agent has only two equally good paths it can follow. Without intervening, the agent first takes the path on the right by moving three steps, while the agent-direction probe grows towards the box. After the first three moves, the chain connects. The agent then takes a U-turn to follow the path on left. This is very unlike the optimal behavior, which is to commit to any path and follow it.

We can get the agent to commit to one path early by intervening with the agent-direction probe on any of the two paths *on the first step only*, as shown in Figure 6 (middle). When we construct a similar level where there are two equally good paths for a box to follow, we find that the agent-direction probe is no longer causal; but the boxes direction probe *does* get the agent to commit. Figure 6 (right) shows an almost empty level. When we intervene with the box-directions probe on the first four steps, the agent follows the laid out path in those steps, but then quickly switches to a more optimal path with fewer turns. The laid out path (not taken, shown in red) is not followed after stopping the intervention.

This analysis provides strong evidence that the network primarily builds upon the boxes directions represented in the activations and relies on the agent-direction when the boxes directions do not inform the subsequent actions.

Figure: Figure 6: Left: A custom level where the agent has two equally good paths to follow. The arrows show the prediction of the agent-direction probe with opacity proportional to the number of times an arrow was predicted across all the steps. The green and red arrows are correct and incorrect predictions, respectively. The agent behaves suboptimally by returning to the start and going left after first going right for three steps. Middle: The agent takes the path on the right after intervening with the corresponding arrows using the agent-direction probe on the first step. The same happens on the left if we intervene on that path without the suboptimal steps right of the undisturbed agent. Right: An empty level with box-directions probe intervening on the first four steps which are followed correctly by the agent on those four steps. When the intervention is removed on the fifth step, the agent computes the simpler path in green and doesn’t follow the path laid out earlier (in red).
Refer to caption: x8.png

## 5 Model surgery exposes latent generalization capabilities

A common feature of search algorithms is that they generalize very well to unusual problems in their domain. This was the main thrust of the argument of Guez et al. []: they exhibited many ways in which the $\text{DRC}(3,3)$ generalized to levels outside of the training distribution to argue that it does search.

### 5.1 Generalizing beyond 10 × 10 10 10 10\times 10 10 × 10 inputs and training examples

This section goes beyond the capabilities of $\text{DRC}(3,3)$ as a black-box unit. By looking for interpretable features in the activations of the ConvLSTM core, we found that layer 3, at the last tick, has four channels which represent the next action to take. Since the network up until the ConvLSTM core is completely convolutional, we can evaluate it on inputs of *any* size – unlike $\text{DRC}(3,3)$ as a whole, which only works on $10\times 10$ inputs, and is trained only on levels with four boxes.

#### Spatial aggregation.

We start with the Layer 3 grid-wise next-action probes from Table 3. For each spatial location, these predict the action the $\text{DRC}(3,3)$ will take in this step. We attempt three different spatial aggregation methods: mean-pooling, max-pooling, and the proportion of locations with positive probe readouts for a particular action. We learn their relative weights and a baseline for each action by using Adam with learning rate $\ell=10^{-3}$, annealed linearly to 0 for 10000 steps. Table 2 lists the learned parameters, which place most weight on the *mean-pooling* aggregation rule. The training set consists of 3000 levels, 1000 from each of the Boxoban unfiltered, medium, and hard training sets.

**Table 2: Weight and bias for transforming direction probes to predictions**
|  | Mean | Max | Positive proportion |
| --- | --- | --- | --- |
| Weight | 1.2086 | -0.0582 | 0.2070 |

#### Results.

The aggregated features obtain 83% accuracy at predicting the MLP’s action output on the training set, which is less than expected given the $90\%+$ F1 scores of probes in Table 3. With only seven parameters, we did not bother collecting an i.i.d. test set and went straight to the generalization tests.

The adjusted ConvLSTM is able to solve many levels that are out-of-distribution: larger than $10\times 10$ in both dimensions, and with more than four boxes, e.g. those in Figure 2. Out of our test set collected from Þorsteinsson [], it can solve $59/480$ levels with both dimensions $>10$, $30/203$ levels with only one dimension $>10$ and $180/482$ levels with both dimensions $\leq 10$. Appendix B has a breakdown by level collection.

### 5.2 Looking for interpretable features with Sparse Autoencoders

The way we found out actions were represented in the ConvLSTM was with sparse autoencoders (SAEs). We
inspected an SAE [Huben et al., , Bricken et al., ] trained on the last layer with the top-k activation
[Gao et al., , Tamkin et al., ], $k=8$. However, we then noticed that the action
features were also represented directly in the ConvLSTM channels.

Table 3 compares raw channels, SAE features, and linear probes for how well they match the
interpretation of “the feature fires when the $\text{DRC}(3,3)$ takes a particular action”. We find that the raw channels
and linear probes are similarly good, and both are considerably better than the SAE.

Thus, for these features, SAEs actually worsened interpretability. We also found three other interpretable features in the
SAE, which were also represented directly in the channels (Appendix E).

**Table 3: Action features scores across channels, probes, and SAE features**
| Feature | Channel | SAE Feature | Probe |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Number | Prec | Rec | F1 | Number | Prec | Rec | F1 | Prec | Rec | F1 |
| Up | L3C29 | 95.7 | 88.1 | 91.7 | L3F270 | 93.9 | 76.2 | 84.1 | 97.5 | 86.5 | 91.7 |
| Down | L3C8 | 98.4 | 80.8 | 88.8 | L3F187 | 98.0 | 79.1 | 87.6 | 97.6 | 86.9 | 91.9 |
| Left | L3C27 | 85.5 | 84.6 | 85.1 | L3F244 | 96.1 | 63.2 | 76.2 | 83.5 | 86.6 | 85.0 |
| Right | L3C3 | 97.0 | 86.9 | 91.7 | L3F385 | 94.6 | 78.5 | 85.8 | 97.6 | 87.4 | 92.2 |

## 6 Related work

#### Interpreting agents and planning.

Several works have attempted to find the mechanism by which a simple neural network does planning in mazes
[Mini et al., , Knutson et al., , Brinkmann et al., ],
gridworlds [Bloom and Colognese, ], and graph search [Ivanitskiy et al., ].
Men et al. [] investigates LLM reasoning in a simple block-stacking task. We believe
the DRC we present is a clearer example of an agent than what these works focus on, and should be similarly possible to
interpret.

Concurrent work [Bush et al., ] interprets another DRC Sokoban agent and can also predict future actions. That team
invented grid-wise probes and the Agent-Directions and Box-Directions targets.

Other works interpret superhuman game-playing agents. Jenner et al. [] find evidence of
lookahead in Leela Chess Zero [Team, ], without predicting future actions.
McGrath et al. [], Schut et al. [] find interpretable concepts in AlphaZero
[Silver et al., ].

#### Probing for world models.

Planning usually requires a model of the world.
Li et al. [], Nanda et al. [], Karvonen [] examine agents trained to predict possible moves in games, and are able to probe for the board state. Wijmans et al. [] also probe the state of navigation neural networks. Gurnee and Tegmark [] find a representation of positions in Earth for location-associated tokens in an LLM.

#### Goal misgeneralization and mesa-optimizers for alignment.

From the alignment perspective, AIs optimizing monomaniacally for a goal have long been a concern (e.g. Yudkowsky, , Omohundro, ; see the preface of Russell, ).
In a machine learning paradigm [Hubinger et al., ], the goal of the training system is not
necessarily optimized; instead, the NN may optimize for a related or different goal [Di Langosco et al., , Shah et al., ] or for no goal at all.

#### Chain-of-thought faithfulness.

Large language models use chain of thought, but are they faithful to it, or do they think about their future actions in other ways [Lanham et al., , Pfau et al., ]? One could hope that LLMs perform all long-term reasoning in
plain English, allowing unintended human consequences to be easily monitored, as in Scheurer et al. [].

#### Ethical treatment of AIs.

Do AIs deserve moral consideration? Schwitzgebel and Garza [] argue that very human-like AIs are conceivable and clearly deserve rights. Tomasik [] suggests that most AIs deserve at least a little consideration, like biological organisms of any species [Singer, ]. But what does it mean to treat an AI *ethically*? Daswani and Leike [] argue that the way to measure pleasure and pain in a reinforcement learner is not by its absolute amount of return, but rather by the temporal difference (TD) error: the difference between its expectations and the actual return it obtained. If the internals of the NN have a potentially different objective [Hubinger et al., , Di Langosco et al., ], then the TD error should come from a place *other* than the critic head. Research into learned search algorithms is an early step toward finding the learned-reward internal TD error, if it exists. This could be a higher-assurance complement to simply asking the AI [Perez and Long, ].

#### Fully reverse engineering small networks.

Many works reverse engineer all of, or most of, a small NN that
does an algorithmic task
[Nanda et al., , Chughtai et al., , Zhong et al., , Quirke and Barez, ].

#### Neural network architectures that reason.

Many papers try to enhance NN thinking by altering the training setup or architecture.
Schwarzschild et al. [, ], Bansal et al. [] improve and
evaluate a specific kind of RNN training that produces algorithmic thinking, though
Knutson et al. [] argue that they do not generalize enough and thus do not
implement the correct algorithm. Other works endow RNNs with variable amounts of computation per step, with [Chung et al., ] or without
[Graves, ] an explicit world model.

## 7 Conclusion

We replicated a small RL-based Sokoban agent, whose performance benefits from increased compute in the form of externally imposed no-op steps [Guez et al., ]. We find that the agent learns to explicitly exploit this capability by pacing, which substitutes for the external no-op thinking steps. Extra computation helps with longer levels and those that require non-myopic thinking.

We show that the agent causally represents *plans* (sequences of actions to take). The most fundamental plan representation we find is the box-direction probes, which can predict and control where the agent pushes a box. By building up a chain of box probe steps over time, the agent is able to figure out how to place boxes to solve the level. Secondarily, we find agent-direction probes that take lower precedence than the box-direction probes, but can control the agent moves in isolated setups. We additionally find that the box probe chains get longer and longer with more thinking steps (Figure 5).

Finally, we find that extracting a decision from the plan takes place almost entirely within the ConvLSTM, with the MLP doing little more than decoding it. To demonstrate this, we show that the agent can generalize to levels bigger than what it had ever seen during training by manually replacing its MLP decoder based on our mechanistic understanding.

We hope the model organism we trained and open-source will catalyze research into NNs that learn to plan, and to understand their inner objectives.

### Acknowledgements

We would like to thank ChengCheng Tan and Siao Si Loo for editing this paper, and the rest of the FAR team for general
support. We also thank Alex F. Spies and anonymous reviewers for improving the paper with their comments and edits. We
are also very grateful for the many questions about their RL setup that Arthur Guez and Timothy Lillicrap answered, as
well as inspiration from Stephen Chung and Scott Goodfriend.

### Reproducibility Statement

We provide the code and the trained models in the supplementary materials. We open-source our code and trained model(s).
Additionally, sections 2, 4.1 and appendices A,
D, and E describe the specific experimental details and settings we used
to carry out our experiments.

## Appendix A Training the test subject

#### DRC ⁢ ( D , N ) DRC 𝐷 𝑁 \text{DRC}(D,N) DRC ( italic_D , italic_N ) architecture.

Guez et al. [] introduced the Deep Repeating ConvLSTM (DRC), whose core consists of $D$ convolutional
LSTM layers with 32 channels and $3\times 3$ filters, each applied $N$ times per time step. Our $\text{DRC}(3,3)$ – or just DRC for brevity – has $1.29$M parameters. Before the LSTM core, two convolutional layers (without nonlinearity) encode the
observation with $4\times 4$ filters.

The LSTM core uses $3\times 3$ convolutional filters, and a nonstandard tanh on the output gate [Jozefowicz et al., ]. Unlike the original ConvLSTM [Shi et al., ], the input to each layer of a DRC
consists of several concatenated components:

- •
The encoded observation is fed into each layer.
- •
To allow spatial information to travel fast in the ConvLSTM layers, we apply *pool-and-inject* by max- and mean-pooling the previous step’s hidden state. We linearly combine these values channel-wise before feeding them as input to the next step.
- •
To avoid convolution edge effects from disrupting the LSTM dynamics, we feed in a $12\times 12$ channel
with zeros on the inside and ones on the boundary. Unlike the other inputs, this one is not zero-padded, maintaining the output size.

#### ResNet architecture.

This is a convolutional residual neural network, also from
Guez et al. []. It serves as a non-recurrent baseline that can only think during the forward pass (no ability to think for extra steps) but is nevertheless good at the game. The ResNet consists of 9 blocks, each with $4\times 4$ convolutional
filters. The first two blocks have 32 channels, and the others have 64. Each block consists of a convolution, followed by two (relu, conv) sub-blocks, each of which splits off and is added back to the trunk. The ResNet has $3.07$M parameters.

#### Value and policy heads.

After the convolutions, an affine layer projects the flattened spatial output
into 256 hidden units. We then apply a ReLU and two different affine layers: one for the actor (policy) and one for the
critic (value function).

#### RL training.

We train each network for 2.003 billion environment steps(^6^66A rounding error caused this to exceed $2\text{B}$ (Section A.2).) using IMPALA
[Espeholt et al., , Huang et al., ]. For each training iteration, we collect 20 transitions on 256 actors using the network
parameters from the previous iteration, and simultaneously take a gradient step. We use a discount rate of $\gamma=0.97$ and
V-trace $\lambda=0.5$. The value and entropy loss coefficients are $0.25$ and $0.01$. We use the Adam
optimizer with a learning rate of $4\cdot 10^{-4}$, which linearly anneals to $4\cdot 10^{-6}$ at the end of training. We clip the gradient norm to $2.5\cdot 10^{-4}$. Our hyperparameters
are mostly the same as Guez et al. []; see Section A.1.

#### A* solver.

We used the A* search algorithm to obtain optimal solutions to each Sokoban puzzle. The heuristic was the sum of the Manhattan distances of each box to its nearest target. Solving a single level on one CPU
takes anywhere from a few seconds to 15 minutes.(^7^77The A* solutions may be of independent interest, so we make them available at [https://huggingface.co/datasets/AlignmentResearch/boxoban-astar-solutions/](https://huggingface.co/datasets/AlignmentResearch/boxoban-astar-solutions/).)

### A.1 Training hyperparameters

All networks were trained with the same hyperparameters, which were tuned on a combination of the ResNet and the $\text{DRC}(3,3)$. These are almost exactly the same as Guez et al. [], allowing for taking the *mean* of the per-step loss instead of the sum.

#### Time limits.

During training, we want to prevent strong time correlations between the returns, so the
gradient steps are not correlated over time. For this reason, the time limit for each episode is uniformly random
between 91 and 120 time steps.

#### Loss.

The value and entropy coefficients are 0.25 and 0.01 respectively. It is very important to *not* normalize the advantages for the policy gradient step.

#### Gradient clipping and epsilon

The original IMPALA implementation, as well as Huang et al. [], *sum* the per-step losses. We instead average them for more predictability across batch sizes, so we had to scale down some parameters by a factor of $1/640$: Adam $\epsilon$, gradient norm for clipping, and L2 regularization).

#### Weight initialization.

We initialize the network with the Flax [Heek et al., ] default: normal weights truncated at 2 standard deviations and scaled to have standard deviation $\sqrt{1/\text{fan_in}}$. Biases are initialized to 0. The forget gate of LSTMs has 1 added to it [Jozefowicz et al., ]. We initialize the value and policy head weights with orthogonal vectors of norm 1.
Surprisingly, this makes the variance of these unnormalized residual networks decently close to 1.

#### Adam optimizer.

As our batch size is medium-sized, we pick $\beta_{1}=0.9$, $\beta_{2}=0.99$. The denominator epsilon is $\epsilon=1.5625\cdot 10^{-7}$. Learning rate anneals from $4\cdot 10^{-4}$ at the beginning to $4\cdot 10^{-6}$ at 2,002,944,000 steps.

#### L2 regularization.

In the training loss, we regularize the policy logits with L2 regularization with coefficient $1.5625\times 10^{-6}$. We regularize the actor and critic heads’ weights with L2 at coefficient $1.5625\times 10^{-8}$. We believe this has essentially no effect, but we left it in to more closely match Guez et al. [].

#### Software.

We base our IMPALA implementation on Cleanba [Huang et al., ]. We implemented Sokoban in C++ using Envpool [Weng et al., ] for faster training, based on gym-sokoban [Schrader, ].

### A.2 Number of training steps

In the body of the paper we state the networks train for $2.003\text{B}$ steps. The exact number is $2\,002\,944\,000$ steps. Our code and hyperparameters require that the number of environment steps be divisible by $5\,120=256\text{ environments}\times 20\text{ steps collected}$, because that is the number of steps in one iteration of data collection.

However, $2\text{B}$ is divisible by $5\,120$, so there is no need to add a remainder. We noticed this mistake once the networks already have trained. It is not worth retraining the networks from scratch to fix this mistake.

At some point in development, we settled on $80\,025\,600$ to approximate $80\text{M}$ while being divisible by $256\times 20$ and $192\times 20$. Perhaps due to error, this mutated into $1\,001\,472\,000$ as an approximation to $1\text{B}$, which directly leads to the number we used.

### A.3 Learning curve comparison

It is difficult to fully replicate the results by Guez et al. []. Chung et al. [] propose an improved method for RL in planning-heavy domains. They employ the IMPALA $\text{DRC}(3,3)$ as a baseline and plot its performance in Chung et al. [, Figure 5]. They plot two separate curves for $\text{DRC}(3,3)$: that from Guez et al. [], and a decent replicated baseline. The baseline is considerably slower to learn and peaks at lower performance.

We did not innovate in RL, so were able to spend more time on the replication.
We compare our replication to Guez et al. [] in Figure 7, which shows that the learning curves for $\text{DRC}(3,3)$ and ResNet are compatible, but not the one for DRC(1,1). Our implementation also appears much less stable, with large error bars and large oscillations over time. We leave addressing that to future work.

The success rate in Figure 7 is computed over 1024 random levels, unlike the main body of the paper. Table 5 reports test and validation performance for the DRC and ResNet seeds which we picked for the paper body.

The parameter counts (Table 4) are very different from what Guez et al. [] report. In private communication with the authors, we confirmed that our architecture has a comparable number of parameters, and some of the originally reported numbers are a typographical error.

**Table 4: Parameter counts for each architecture.**
| Architecture | Parameter count |  |
| --- | --- | --- |
| $\text{DRC}(3,3)$ | 1,285,125 | (1.29M) |
| $\text{DRC}(1,1)$ | 987,525 | (0.99M) |
| ResNet: | 3,068,421 | (3.07M) |

Figure: Figure 7: Success rate for Test-unfiltered and Validation-medium levels vs. environment steps of training. Each architecture has 5 random seeds, the solid line is the pointwise median and the shaded area spans from the minimum to the maximum. The dotted lines are data for the performance of architectures extracted from the [Guez et al., ] PDF file. The values are slightly different from what Figures 5 and 3 report because they are calculated on a random sample of 1024 levels (24 levels are repeated for test-unfiltered, which only has 1000 levels).
Refer to caption: x11.png

**Table 5: Success rate and return of DRC and ResNet on the unfiltered test set at various training environment steps.**
| Training Env | Test Unfiltered | Valid Medium |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Steps | ResNet | DRC(3, 3) | ResNet | DRC(3, 3) |  |  |  |  |
|  | Success | Return | Success | Return | Success | Return | Success | Return |
| 100M | 87.8 | 8.13 | 95.4 | 9.58 | 18.6 | -6.59 | 47.9 | -0.98 |
| 500M | 93.1 | 9.24 | 97.9 | 10.21 | 39.7 | -2.64 | 66.6 | 2.62 |
| 1B | 95.4 | 9.75 | 99.2 | 10.47 | 50.0 | -0.64 | 70.4 | 3.40 |
| 2B | 97.9 | 10.29 | 99.3 | 10.52 | 59.4 | 1.16 | 76.6 | 4.52 |

## Appendix B Generalizing the DRC ⁢ ( 3 , 3 ) DRC 3 3 \text{DRC}(3,3) DRC ( 3 , 3 ) to larger levels

We license the levels by Þorsteinsson [] as GPLv3, and make them available in the supplementary material. The
performance of the $\text{DRC}(3,3)$ on each set of levels is in Table 6. The sets with higher
scores are also those which humans find easier: for example, “Dimitri & Yorick” was made for children by Jacques
Duthen and consists entirely of relatively small levels (largest is $12\times 10$) with at most 5 boxes. The sets where
the $\text{DRC}(3,3)$ solves nothing are also difficult for the authors of this paper.

We tried giving the $\text{DRC}(3,3)$ from 2 to 128 extra thinking steps for these levels, incrementing in powers of two. For
most sets, we find some benefit to 2-4 thinking steps, but no more. The sole exception is XSokoban, which contains one level (Figure 2, right) which requires 128 steps of thinking to solve.

We encourage the reader to go to the website by Þorsteinsson [] and try solving the levels.

Figure: (a) Microban, level 105
Refer to caption: x12.png

Figure: (a) Sokoban Jr. 2, level 17
Refer to caption: x13.png

**Table 6: Performance of the $\text{DRC}(3,3)$ on each set of levels by [Þorsteinsson, ]. The “Max solved” columns represent the proportion of levels solved at the number of steps in the “at steps” column, which is the highest solved proportion for each number of thinking steps tried.**
|  | All levels | Levels larger than $10\times 10$ |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Level collection | # | solved | Max slv. | max at | # | solved | Max slv. | max at |
| Dimitri & Yorick | 61 | 83.6% | 88.5% | 16 | 0 | — | — | — |
| Sokoban Jr. 1 | 60 | 80.0% | 81.7% | 2 | 19 | 73.7% | 73.7% | 0 |
| Howard’s 3rd set | 40 | 60.0% | 62.5% | 4 | 1 | 0.0% | 0.0% | 0 |
| Simple sokoban | 61 | 54.1% | 60.7% | 16 | 51 | 47.1% | 54.9% | 16 |
| Sokoban Jr. 2 | 54 | 48.1% | 50.0% | 2 | 40 | 45.0% | 47.5% | 2 |
| Microban | 155 | 23.9% | 24.5% | 4 | 17 | 5.9% | 11.8% | 2 |
| Deluxe | 55 | 21.8% | 23.6% | 2 | 1 | 0.0% | 0.0% | 0 |
| Sokogen 990602 | 78 | 20.5% | 23.1% | 8 | 0 | — | — | — |
| Yoshio Automatic | 52 | 13.5% | 17.3% | 2 | 0 | — | — | — |
| Sasquatch III | 16 | 6.2% | 6.2% | 0 | 8 | 0.0% | 0.0% | 0 |
| Howard’s 1st set | 100 | 5.0% | 6.0% | 2 | 54 | 0.0% | 0.0% | 0 |
| Howard’s 2nd set | 40 | 5.0% | 10.0% | 4 | 22 | 0.0% | 0.0% | 0 |
| Microcosmos | 40 | 5.0% | 7.5% | 4 | 0 | — | — | — |
| Still more levels | 35 | 2.9% | 2.9% | 0 | 34 | 2.9% | 2.9% | 0 |
| Sasquatch IV | 36 | 2.8% | 2.8% | 0 | 20 | 0.0% | 0.0% | 0 |
| Xsokoban | 40 | 2.5% | 5.0% | 128 | 39 | 2.6% | 5.1% | 128 |
| Sasquatch | 49 | 2.0% | 4.1% | 4 | 39 | 0.0% | 2.6% | 4 |
| David Holland 1 | 10 | 0.0% | 0.0% | 0 | 5 | 0.0% | 0.0% | 0 |
| David Holland 2 | 10 | 0.0% | 0.0% | 0 | 9 | 0.0% | 0.0% | 0 |
| Howard’s 4th set | 32 | 0.0% | 0.0% | 0 | 30 | 0.0% | 0.0% | 0 |
| Mas Sasquatch | 50 | 0.0% | 0.0% | 0 | 43 | 0.0% | 0.0% | 0 |
| Nabokosmos | 40 | 0.0% | 0.0% | 0 | 0 | — | — | — |
| Sokoban | 50 | 0.0% | 0.0% | 0 | 48 | 0.0% | 0.0% | 0 |

## Appendix C Case Studies

Figure: (a) file 0, level 18 unsolved $\overset{\text{think}}{\rightarrow}$ solved
Refer to caption: x16.png

#### Case Study: Thinking makes some levels solvable Figure 10 (a).

For example, thinking lets the DRC solve Figure 10(a). In the no-thinking condition, the DRC first pushes box $C$ one square to the right. It then goes back to push $A$ to $a$, but it has now become impossible to push box $B$ onto $b$. In contrast, after thinking, the DRC pushes $A$ to $a$ first, which lets it solve the level.

#### Case Study: Thinking speeds up solving Figure 10 (b).

Figure 10(b) illustrates a scenario where the agent happens. In the no-thinking condition, the DRC takes many steps back and forth before pushing any boxes. First it goes down to $y$, up to $c$, then down onto $z$, back up to $y$ and to $z$ again. It then proceeds to solve the rest of the puzzle: push box $A$ onto $a$, prepare box $B$ on $x$ and box $C$ where $B$ originally was, push in boxes $B,C$ and finally $D$. In the thinking condition, the DRC makes a beeline for $A$ and then plays the same solution.

#### Case Study: Thinking slows down Figure 10 (c).

Figure 10(c) illustrates a scenario where thinking time results in a slower solution. In the no thinking condition, the DRC starts by pushing box C into position y, then pushes boxes A, then B into place. On the way back down, the DRC pushes C onto c and finally D onto d. In contrast, in the thinking condition, the DRC goes the other way and starts by pushing B onto b. The subsequent solution (A, C, then D) is the same, but the DRC has wasted time trekking back from B to A.

## Appendix D Probe Training

#### Probe architecture

We train multiple different linear probes on the hidden states $h$ and cell states $c$ activations of the network. The probes are trained either on the states of the individual layers or by concatenating the states of all the three layers. We collect the activations of the model by letting it play through the hard levels. The training set for the probes are constructed by taking a random sample of cached levels and including all the timesteps in the level excluding the first five timesteps. We exclude these intial steps as we believe they can be noisy as the network comes up with the stategy to solve the levels.

The probes are trained with logistic regression with L1 decay using the Scikit-Learn library. We do a grid-search on the
learning rate and the L1 weight decay to select the probe that has the highest F1-score on a validation set.

For multi-class probe targets, each potential output class $l$ is treated like a different data point. That is: for each
label $l$, we consider a prediction positive if the highest probe logit is for label $l$, and negative otherwise. We
consider a data point positive iff the true label is label $l$. Then, we compute the confusion matrix and the F1 score
from the $n\cdot l$ data points.

## Appendix E Details about Sparse Autoencoders (SAEs)

Figure: Figure 11: Visualization of some interpretable features from the SAE of last layer. These features also appear monosemantically in the channels. The precision, recall, and F1 score for the features are reported in Table 10.
Refer to caption: x19.png

In order to search for more monosemantic and interpretable features in the network, we train sparse-autoencoders (SAEs) [Huben et al., , Bricken et al., ] on the individual squares in the $h$ hidden state of the network consisting of 32 neurons. Thus, we get a $10\times 10$ visualization for each SAE feature as shown in Figure 11. We use the top-k activation function [Gao et al., , Tamkin et al., ] for the SAE as it is currently the state-of-the-art method to train SAEs and is easier to work with as it directly sets the $L_{0}$ norm of the SAE activations. The hyperparameter search space for training the SAE is provided in Table 7. We train separate SAEs for each layer with the specified hyperparameters and pick the one that achieves greater than $90\%$ explained variance while having interpretable features assessed through manual visual inspection. We release these trained SAEs and probes in the same huggingface repo as our trained DRC networks. (^8^88[https://huggingface.co/AlignmentResearch/learned-planner](https://huggingface.co/AlignmentResearch/learned-planner))

We provide examples of some of the interpretable features from an SAE trained on the last layer with $k=8$ in Table 8 with corresponding visualization in Figure 11. For the “Target”, “Unsolved”, and “Solved” concepts (Table 10), we observed cases where the SAE feature is offset from the ground truth by 1 square in either horizontal or vertical directions. The level setup includes a permanent one-square outer-edge wall, so this offset never results in an out-of-bounds issue. We evaluated these potential “Offset” variants for the Target, Unsolved, and Solved concepts.

All of the features in the SAE that we find to be interpretable are already embedded in individual channels. These
channels are either as monosemantic as the SAE features, or in some cases more monosemantic as measured by the F1 score.
Table 3 reports the precision, recall, and F1 scores for the action features as evaluated through
channels, SAE neurons, and linear probe trained against the ground truth action predictions.
Table 10 reports the same scores for additional interpretable features. The SAE action
features have lower F1 scores than channels by a margin of $5.9\%$ on average. On the other hand, the linear probes
trained across all channels of the hidden state have similar F1 scores as the specified channels indicating that the
channels are monosemantic and cannot be improved upon by combining various channels using a linear probe.

**Table 7: Hyperparameter search space for training SAE**
| Hyperparameter | Search Space |
| --- | --- |
| $k$ | $\{4,8,12,16\}$ |
| learning rate | $\{1e-5,5e-5,1e-4,5e-4,1e-3\}$ |
| expansion factor | $\{16,32,64\}$ |

**Table 8: SAE Feature Concepts**
| Concept | Description |
| --- | --- |
| Target | The 4 target squares (static) |
| Unsolved | Targets and boxes that aren’t solved |
| Solved | Solved target squares with a box on them |
| Agent Up | The agent will move Up next step |
| Agent Down | The agent will move Down next step |
| Agent Left | The agent will move Left next step |
| Agent Right | The agent will move Right next step |

**Table 9: Breakdown of levels by category at 6 thinking steps.**
| Level categorization | Percentage |
| --- | --- |
| Solved, previously unsolved | 6.87 |
| Unsolved, previously solved | 2.23 |
| Solved, with better returns | 18.98 |
| Solved, with the same returns | 50.16 |
| Solved, with worse returns | 5.26 |
| Unsolved, with same or better returns | 15.14 |
| Unsolved, with worse returns | 1.36 |

**Table 10: Scores for SAE and Channel features**
| Concept | Offset $(dy,dx)$ | Channel | SAE Feature |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | Number | Prec | Rec | F1 | Number | Prec | Rec | F1 |
| Target | (1, 0) | L3C17 | 97.8 | 97.7 | 97.8 | L3F278 | 97.8 | 98.1 | 98.0 |
| Unsolved targets and boxes | (0, 0) | L3C7 | 94.9 | 90.8 | 92.8 | L3F212 | 95.3 | 86.6 | 90.7 |
| Solved targets | (0, 0) | -L3C7 | 91.6 | 94.6 | 93.0 | L3F179 | 91.7 | 91.5 | 91.6 |

## Appendix F Additional quantitative behavior figures and tables

Figure: Figure 12: Left: A Sokoban level from the hard set. Right: the same level as the NN sees it, one pixel per tile. Walls are black, boxes are brown, targets are pink and the robot is green.
Refer to caption: x20.png

Figure: Figure 13: Number of thinking steps required to solve the level vs. number of nodes A* needs to expand to solve it. The trend is increasing at the end but very unclear, indicating different heuristics used by the DRC and A*.
Refer to caption: x22.png

Figure: (a) Cycle length distribution
Refer to caption: x24.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Could the authors clarify the relation to Guez et al. 2019 by explaining the overlapping parts between this paper and Guez et al. 2019, and new results.
2. For self-containness, could the authors provides a preliminary about DRC since it is the backbone of the current work.
3. Could the authors provides some illutrative examples  (rendering) of "pacing"?
4. How does the network deal with different size of input (other than 10x10) in the CNN part (Sec. 5.1)?  I am confused about how the convLSTM was adjust in Sec 5.1. Maybe a visualized diagram could help. 
5. What was the learning objective in spatial aggregation ("We learn their relative weights ....... for 10000 steps.")?
6. DRC was compared only with ResNet, how about plain LSTM (as in Guez et al. 2019)? And I am curious about how a Transformer network performs in this case. 
7. Many of the plots seems to miss error-bars, is it because the error bars are too small? Could the authors clarify the error bars (and number of repeats) not only in Tables but also in plotted curves?
8. In the Related Work section, could the authors make a more concrete explanation about "reasoning NN architectures"?
9. Is it necessary to include "ethical treatment of AIs" in the main texts? 
10. Could the authors clarify "perhaps due to errors" ? (line 790)
soundness: 3
strengths: 1. The current manuscript is highly informative a number of interesting analysis, many of which are novel to my knowledge (beyond those discussed in Guez et al. 2019).
2. Some of the results are thought-involking, which I have not thought about, such as the "pacing" mechanism.
3. The being touched task (planning in Sokoban) is quite interesting (while simple enough for humans to interpret) and potentially important for understanding planning.
summary: The paper dives into understanding the inner mechanisms of planning by deep reinforcement learning (RL) agents in the Sokoban game, especially using a deep repeated convLSTM (DRC) (Guez et al. 2019) network structure. This work proposes several interesting findings including the effect of thinking steps, "pacing", probing the interpretable information, and generalizability to different sizes of Sokoban. All these results unveil a path to open the black box of a (recurrent) neural networks for model-free planning. I believe the major contribution of the paper lies on, like the authors stated, understanding the inner objective of the planning network.
weaknesses: While the current study presents a number of interesting results. I believe the manuscript could be improved in several aspects.

1. The paper could benefit from a better struture and more scientific style of writting. The current version, while informative, is not self-contained and often makes me feel difficult to understand the implementation details of each analyze method. The authors should assume minimal pre-knowledge of the methods used in this paper, for a more general audience.
2. The paper does spend enough efforts to convince the audience that the presented results are not cherry-picking. Although I tend to believe the presented findings are mostly general, the paper could be improved by e.g., discussing the variety among random seeds in training, and providing more error bars in the plots.

See the following questions for detailed comments.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 1
questions: Besides my comments above some specific points on clarity:
* You never explain what probes are and how their work, if their reader is not already very familiar iwth them they won't understand the work.
* Lines 153-156 is difficult to follow the thinking process here, for instance you mention there that adding extra thinking steps solve 4.7% extra levels, but in line 157 you mention that the extra thinking  helps greatly in dificult problems. Are only a very small percentage of the results these har problems?
* It is not clear how the ResNet would perform cycles, the few it does
* Line 200-202: here is a good example of a point that a research question is made but the paper only offers evidence that "suggests" not that confirm. In that and the lines below i would be good to put down the reasoning that would give us certainty that the agent is pacying, and giving an intuition of what is pacing here.
*  The later two paragrapghs in section 4.0 would be better placed as a summary at the end of the section pointing out to where the claims are supported.
* Line 278 explain what is the standard of evidence in Li et al.
soundness: 1
strengths: The paper explores an interesting research direction on how a neural network within an emobied agent abstract their problem at hand, projects a plan and makes the decision making. It is an interesting problem that trhough a sound study can be very useful and benefitial for the community
summary: This paper makes an study on how a neural network within an RL agent learns to play Sokoban. Specifically, this work studies how and where the network abstracts the problem at hand, plans and where decision making takes place.
weaknesses: While the line of research is of interest, the current paper has important shortcomings that require significant changes. 

Specifically, there are three major weaknesses in this paper:

* Clarity: As I will detail below the paper is very difficult to understand and it would greatly benefit from a rewritting. Specifically, I would suggest authors to do a simple exercise at the beggining what are the hypothesis that are being tested, how this thesis is going to be tested and what results would be interpreted as a validation or violation of the hypothesis. Additionally, there are several concepts that are never introduced in the paper or that are introduced first and explained better. Last, a diagram about the architecture studied and where the probes (i.e., the system to track the hypothesis) are placed within the architecture would greatly help readers

* Generalization of the results: the whole study focuses exlusively on the environment of Sokoban and studies only two neural networks, casting doubts about how general the conclussions extracted are.

* Related literature. There are several works that study which drivers help best to RL agents to generalize in similar problems, both within the environment, the nn architecture or the perception field. The present paper doesn't follow any of the recomendations of those works,and they can be influential in the conclussions extracted, e.g. [1] demonstrated that for an RL agent to generalise and not memorize puzzle solving, should have an egocentric perspective, how does using a egocentric perspectice  or not effect planning?. [3,4,5] highlighted changes in the ability of rl agents to generalize their plans when using different architectures, do we observe changing on abstractions when using any of those configurations?


[1] Hill, Felix, et al. "Environmental drivers of systematicity and generalization in a situated agent." International Conference on Learning Representations.
[2] Hill, F., Tieleman, O., von Glehn, T., Wong, N., Merzic, H., & Clark, S. Grounded Language Learning Fast and Slow. In ICLR 2021.
[3] León, Borja G., Murray Shanahan, and Francesco Belardinelli. "Agent, do you see it now? systematic generalisation in deep reinforcement learning." ICLR Workshop on Agent Learning in Open-Endedness.
[4] León, Borja G., Murray Shanahan, and Francesco Belardinelli. "In a Nutshell, the Human Asked for This: Latent Goals for Following Temporal Specifications." ICLR 2022.
[5] Lake, Brenden M., and Marco Baroni. "Human-like systematic generalization through a meta-learning neural network." Nature 623.7985 (2023): 115-121.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: 1) Why does the causal probe perform so poorly compared to their predictive power?  
2) What is special about pacing as opposed to other forms of moving around?  
3) What is being aggregated over in Section 5.1.?
soundness: 2
strengths: - **Nice task.** I think the use of Sokoban is interesting as it is a hard planning problem that has many failure modes for greedy models which perform actions without careful regard and deliberation of future actions. I think research that interprets how models perform planning is especially relevant now with current developments happening in the reasoning abilities of large (language) models.  
- **A variety of interpretability techniques.** The authors make use of a variety of techniques of mechanistic interpretability.  
- **Planning in a model free setting.** I think it is interesting that the authors attempt to explicitly study the ability to plan in a model-free agent, where historically planning has been in the camp of model-based RL. DRC represents a very flexible model class with minimal inductive biases. This flexibility leaves open if the model learns to plan in the traditional sense. It is an interesting question if model representations and computation might represent plans explicitly.
summary: The paper studies the behaviour of a Deep Repeated ConvLSTM (DRC) network architecture that learns to play Sokoban. The paper specifically documents the phenomenon of “pacing” whereby the network moves in cycles before performing goal-directed actions. The authors interpret these behaviours as an indication that the model gives itself additional computational steps before the execution of a strategy. The authors also engage in an interpretability analysis of the models via the training of linear probes on LSTM states. The authors find that probes (logistic regressions) trained on these states are able to decode a variety of task relevant variables such as the future movements of boxes and agent movements for locations in the grid world. The performance of these probes varies over the different probe targets but are generally predictive of future actions. One specific type of probe, encoding the direction of boxes, is also successfully leveraged to intervene in the model activations to change (steer) model behaviour. Finally, the authors modify the DRC decoder to allow the model to somewhat successfully generalise to new puzzles of differing sizes and with variable number of boxes.

### Recommendation

I think the paper should be rejected because (1) the claims are not well supported, (2) it is unclear and unpolished in many places, and (3) I think quite disconnected from the literature it claims to be connected to. See more detailed criticisms above.
weaknesses: Despite these strengths I think the paper has many weaknesses and does not do a great job of elucidating if the model is planning. In addition, I think the paper is very unclear in many places. 

- **Framing of results.** The paper is abundant with sections in which the authors make quite substantial and large claims. I.e. in the introduction claims that “linear probes show that activations represent a plan that causes actions” and that “the behaviour of the RNN suggests that it performs search”. However, Section 4 openly admits that they find little evidence to suggest that the model is performing search. I think the claims that the paper makes and the presented evidence are quite disconnected and need substantial revamping. I also think that a contribution statement would make clear what really is being contributed.  
- **Strength of evidence.** I do not think some of the author's evidence is strong enough to support their claims. In Section 4 on linear probing the predictive power of the next box and next target probes are much lower than direction probes. I think intuitively should these not be a much better measure of planning depth as they are more relevant to long term plans? The results on generalization also appear somewhat limited as the performance is not as good as the intro and abstract might make you believe.   
- **Novelty.** Section 5 shows some generalisation ability of DRCs but this has already been shown in a variety of settings by Guez et al. (2019) and I think it is not super clear what the current paper adds.   
- **Clarity.** I think there are also places in which the elaboration of the methods used are not clear. For example, I found myself quite confused about the methods used in Section 5.1 *Spatial aggregation*. I was looking in the appendix but could not find an exact explanation of what was done here. What exactly is being aggregated? Also, what is the number column in Table 3?  
- **Interpretability results left uninterpreted.** There are several places where the authors fail to give sufficient explanations, interpretation, or context for their findings. For example, the section on linear probes finds good predictiveness but *much* worse causal probe results. Why is this the case? I felt like I was just left wondering as the reader.  I am also wondering what is special about pacing as opposed to other forms of taking inconsequential actions that are not cyclical, both should yield the same result, right? I think the authors would have to give an explanation about what is special about pacing.  
- **Unpolishedness.** There are several places where I felt that the paper was not well polished and presented. For example the caption in Figure 5 does not have labels for the subplots (i.e. left, right, centre) even though these are referred to in the main text. Similarly, in Figure 3 the x-axis labels are below the subplot caption, which makes the figure unclear. Further, often the terms NN, RNN, and DRC are used interchangeably. I think it would be good to just stick to one term referring to the model (DRC is used the majority of times). There is also the sentence “we did not bother collecting an i.i.d. test set” in Section 5, I personally think such language is not perfectly appropriate for a paper.   
- **Placing in the literature.** I think the paper does quite a poor job of placing the paper in the literature. Especially the related work section (Section 6\) I find quite strange. For example it is not clear to me how the presented work connects to the “Ethical treatment of AIs”  or to “Goal misgeneralization and mesa-optimizers for alignment”. 

### Potential Improvements
1. The claims should be appropriately fitted to the findings.  
2. The clarity of the paper should be improved with respect to the methods used in Section 5.1.  
3. The related works section should be more carefully written.   
4. Perhaps reconsider having the section on sparse autoencoders in the main text. If they do not perform well then why have them with your main results as opposed to the appendix?

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See questions in the weaknesses about the unclear points of the paper.
soundness: 1
strengths: The claims of the paper are ambitious.
summary: The paper investigates an RL agent that can pace itself by taking extra computational steps in complex scenarios. The paper claims that by adjusting the network's hidden states, they extended its capabilities, allowing it to solve much larger and more challenging levels, and have open-sourced the compact, 1.29M-parameter model for further study in learned planning.
weaknesses: Some claims seem unrealistic, for instance in Figure 2, it is written "Left: XSokoban-31, which the DRC(3, 3) solves with 0 thinking steps", I'm not convinced that this claim is correct. Are the trained models already shared such that this could be checked? (beside the promise "We will open-source our code and trained model(s)").

Some words do not seem to have a clear meaning, e.g. in the abstract "perform model surgery" and "makes it an excellent model organism".

Some other parts are unclear: 
- For instance, at the end of the introduction, it is written "The evidence highly suggests that the RNN implements an online search algorithm". What evidence exactly is there for such a "high suggestion"?
- Beginning of Section 3.1 "the extra thinking time disproportionately helps with more difficult levels": what does disproportionately mean in this case?

Meta-review: The paper examines how a Deep Repeated ConvLSTM (DRC) neural network learns to play the game Sokoban, focusing on the emergence of planning behavior in this model-free agent. The authors claim the DRC exhibits "pacing" behavior, taking extra steps in complex scenarios to allow for more computation time. They further assert that linear probes can predict the network's future actions and that manipulating the hidden state through these probes allows for control over the agent's subsequent actions. 

Strengths: The paper explores the under-researched area of planning in model-free agents. Specifically, it utilizes a variety of interpretability techniques to analyze the DRC network's internal representations and decision-making processes. Moreover, the use of Sokoban as the task domain is interesting, as it requires the agent to reason about sequences of actions and their consequences.  However, the original submission had several weaknesses: 

- Lack of Clarity and Rigor: The paper uses vague language, poorly defined hypotheses, and an unstructured presentation that hinders comprehension. This lack of clarity extends to the methodology, which is not described in sufficient detail to allow for proper evaluation of the results.

- Insufficient Evidence and Limited Generalizability: The core claim of planning behavior is not convincingly supported by the evidence presented. The connection between the observed "pacing" behavior and actual planning is tenuous, and the limited experimental scope to Sokoban raises concerns about the generalizability of the findings.

- The paper fails to adequately contextualize its findings within the broader literature on planning in artificial intelligence. The related work section is inadequate and does not effectively position the research within the existing knowledge base.

Based on the discussion period, I'm leaning towards rejecting the paper despite author's attempts to address the above weaknesses, as they result in substantial modifications that seem to be appropriate for a revised submission to another conference.

Decision: Reject