Planning in a recurrent neural network that plays Sokoban

We are delighted that you found our work ambitious, novel and surprising. We are happy to assist you in replicating any of our claims.

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)").

Yes, we submitted one trained model (the one we examine) along with code as supplementary material to OpenReview. We should have stated that instead of “will open-source”, apologies. For your ease of verification, we have updated the code and model weights format so it only requires Pytorch (not Jax) and contains all the dependencies, anonymized and in a runnable format.

Instructions for running the NN that solves XSokoban-31: Get a fresh Python 3.11 environment and go to the uncompressed supplementary-material folder. Only CPU is needed. Then install everything:

pip install -e MambaLens -e farconf -e gym-sokoban -e stable-baselines3 -e 'learned-planners[torch,dev-local]'

And run the script: python learned-planners/plot/play_bigger_levels.py . This should run in under a minute and will create a level_22_031.mp4 , where you can watch the DRC(3, 3) solve the X-sokoban-31 level. You can change the level/s to be generated by passing options to the script (run python learned-planners/plot/play_bigger_levels.py --help). We encourage you to read the code to check it does what we wrote in the paper.

Please do not hesitate to write back if you have any problems running the script. We have tested on a fresh virtualenv in a Mac, and also inside the python:3.11 Docker container ( use docker run -v $(pwd):/workspace -it python:3.11 /bin/bash to get an environment and then run the commands above ).

We originally also shared pickled probes and sparse autoencoder weights (SAEs) under the learned-planners subdirectory of the submitted supplementary material. You can also watch the generated videos, e.g. at videos/bigger-levels/xsokoban-31-0.mp4 . We have also added the anonymised training code train-learned-planners, but we have not tested running it so it is mainly for you to read.

We can answer any further questions you have and can assist you in reproducing the claims.

It is worth noting that the levels depicted in Figure 2 are cherry-picked from the set of levels intended for humans. We have also cherry-picked examples of levels the DRC(3, 3) does not solve or for which thinking steps hurt in Figure 8, as well as a (non-cherry-picked) full breakdown in Table 6. As Table 6 indicates, the DRC(3, 3) can only solve these 2 levels (5.0% of 40) from XSokoban at its best, though there are many large levels from the other sets that it solves. On reflection we did not emphasize this enough, and will change “more examples in appendix B” to “For many example levels which the DRC does not solve, as well as a breakdown of solve rate by level collection, see appendix B”.

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".

We borrow the terms from the ML interpretability community. Model organism originally comes from biology (https://en.wikipedia.org/wiki/Model_organism) and is used prominently in e.g. Hubinger et al. (2024). In biology (interpretability), a model organism refers to a simpler organism (neural network) that can be studied with the expectation that discoveries made in the model organism will provide insight into the workings of other more complex organisms (neural networks).

Model surgery is also employed in some interpretability works, e.g. Wang et al. 2024 and Li et al. 2023. It refers to changing the capabilities of the model in a targeted way, typically one that requires some knowledge about the algorithm it implements. It is possible we should call this ‘model editing’ (doesn’t have to be targeted) or ‘model stitching’. In any case we will refer to the earlier works. Thank you for pointing out that this is not a common term.

Thank you for taking the time to review our paper.

Framing of results. The specific claim we make in this paper is that we have strong evidence that the model has a plan, and weak evidence that it’s doing search. This matches the sentences you highlight: the ability to read/write the plan with probes is strong evidence that there is a plan, and behavioral evidence is weak.

We agree that a clearer contribution statement was necessary. We have updated the paper by summarizing the results from Guez et al. (2019) and added a contributions subsection in the introduction. We have also reformatted the paper to clearly specify the 4 different hypothesis we have in section 3 and explained before each section which hypothesis we are testing there.

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?

Note that we use linear models as probes. The comparatively low F1 score for next target/box probe shows that the information for the next target is not exactly linearly encoded in the hidden states. One reason could be that the agent directions indirectly encodes the next target information which should be extractable using a non-linear probe. We have added a comment clarifying this in the paper.

With that said, the NN does not necessarily need to encode the next boxes/targets in a local, linear way.

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.

Guez et al. (2019) showed the generalization to levels with more boxes but the same grid size of, $10 \times 10$. However, our results are more striking since we check generalization to much larger levels with many boxes as shown in Figure 2. In fact, we do this by removing the MLP layer and using the next action information linearly encoded in the last layer’s hidden state channels. We show that the underlying algorithm mechanism that's learned is generalisable even when the network naively isn't

the section on linear probes finds good predictiveness but much worse causal probe results. Why is this the case?

Our probe results show that many of the probes are highly predictive but only the box directions probe is highly causal to modify the agent’s next action. This shows that the future box directions computation is part of the model’s algorithm to predict action. A probe trained to predict a property $P$ that has high predictive score but low causal score shows that the information about $P$ can be read from the activations but $P$ is not part of the model’s algorithm to predict action. It could be that $P$ is used by the network to do something else (e.g., predict the value of the state) or $P$ correlates with the activations. For example, we found that the agent directions probe is always encoded (highly predictive) but is only causal when the box directions representations are not informative for predicting the next action, as shown in the case study of section 4.3. This results in lower causal scores for the agent directions probe. We have made this interpretation of the causal probes results more clear in the paper.

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”.

Regarding ethical treatment of AIs related work, we believe it is important and related and we have explained it better now. However, several reviewers had raised that this is a strange argument, so we have happily moved it to the appendix.

What is special about pacing as opposed to other forms of moving around?

The pacing in cycle behavior is a completely redundant behavior that an optimal policy (such as that computed using A*) would not perform. However, this behavior was implicitly learned by RL to provide DRC more compute during inference to come up with better plans as shown in the paper. This is clearly distinct from other forms of moving around such as moving in a straight line or pushing a box which are goal-directed and are based on the plans that DRC forms early on.

We have added new evidence in lines 399-416 of the updated paper to show that cyclic-pacing behavior is more special. You can find a summary of these new evidence in the new common response.

What is being aggregated over in Section 5.1.?

The action probes are applied across channels of the last layer’s hidden state $h$ for individual squares in the grid, which results in a $10 \times 10$ probe output for each action. We aggregate this output using the spatial-aggregation methods in section 5.1. We have clarified this in the updated paper.

1. We have provided a summary of the overlapping results from Guez et al. (2019) and the novel contributions we make through our paper in the common response. We are updating the introduction section in the paper to reflect this.
2. We describe the DRC architecture in slightly more detail in appendix A, but it’s pretty compressed. We will include a figure of the architecture in our updated paper.
3. You can find a video of the pacing behavior at the path videos/pacing-cycles.mp4 in the uncompressed supplementary material. We are also adding a figure demonstrating pacing behavior in the paper.
4. The DRC architecture is based on a convolutional backbone with “same” padding such that it can take in any $M \times N \times 3$ image as input and process them into a $M \times N \times 32$ hidden states $h$ and $c$. Normally, this is flattened and fed into an MLP layer. In our model surgery section, we remove the MLP and spatially-aggregate the hidden state using mean-pooling, max-pooling, and proportion of positive squares readouts with action probes. This gives us 7 inputs as shown in Table 2 to convert hidden state to actions. We will add a visualization of the architecture in the paper with and without the model surgery.
5. We use the cross-entropy loss function to optimize the 7 parameters against the ground-truth actions taken by the network using the MLP layer. We include this in our updated paper.
6. We compare with the ResNet to demonstrate that the recurrent architecture is significantly more performative than bigger non-recurrent architectures. We agree with the reviewer that it is interesting to see how a transformer network performs in this environment and whether it will demonstrate a similar planning mechanism. We are planning to do that in the next project, the aim of this paper is to analyze the workings of the DRC architecture that Guez. et al (2019) found and described the behavior of.
7. Thank you for pointing this out, this is our mistake. Some figures have error bars: e.g. Figure 5(mid, right) and Figure 7. But for most figures we forgot, we focused on error bars for the tables. There are two kinds of error bars that make sense for our paper to have: error across seeds, and across dataset samples. Most of this paper is about interpreting a single trained NN (i.e. a single seed) so it only makes sense to add dataset error bars to figures: 3, 4, 5. Figure 1 would benefit from seed-based error bars because we’re making statements about training.
8. We’ll expand the section on “reasoning architectures” in the related work. It’s not a concrete category, we meant to highlight works that come up with NN architectures inspired by making them explicitly reason, or giving it a model of the world.
9. We believe it is important. However, several reviewers have raised that this is a strange argument, so we are happy to move it to the appendix.
10. We made bugs when setting the number of steps. It was important that the number of steps be divisible by 256x20, and the number on that line is divisible by that, but so is 2Billion (though not 1Billion). So basically: we wrote the wrong number before training, and after training we decided the number of steps was close enough to 2 billion that we would not bother running the training again.

Thank you for pointing out how the paper is unclear. We’ve definitely been thinking of the hypotheses and evidence for/against, but did not make that explicit enough. Hopefully the common response explains a bit better which hypotheses and evidence we are considering. (hypotheses: does this NN have plans? Does this NN do search?)

Thank you also for the very thoughtful review – we’ve come up with better experiments about pacing that we will definitely incorporate into the next version of the paper.

the whole study focuses exclusively on the environment of Sokoban and studies only two neural networks

In fact we only really study one network, the DRC. The ResNet is used as a baseline for performance and little else.

Mechanistic interpretability requires manual work and is not easily repeatable for several architectures. But now that we have a methodology, we could re-run our spatial probe analysis in other grid environments (e.g. gridworld mazes) – would that be satisfactory for generalization?

There are several works that study which drivers help best to RL agents to generalize in similar problems

Thank you for pointing out related works, we will definitely include them! Studying the different abstractions of the various architectures would be very interesting, if a bit difficult with the probing methodology.

The focus of this paper is not on making better RL architectures or making RL architectures that reason more, but instead on studying RL architectures that are as close to standard as possible, and whether they plan. For this reason, we stayed as close to the normal RL pipeline as possible, though it is a bit unfortunate that we have this strange DRC architecture.

Question Responses:

1. Thank you very much for mentioning this, we will explain in the paper: A probe is a simple machine learning model that predicts something based on the internal activations of a NN. It is meant to show where (if present) the NN represents certain concepts.
2. Thinking steps helps in solving an extra 4.7% levels on the medium difficulty validation set. On the hard difficulty set, the network solves extra 7.1% levels with thinking steps (see FIgure 1 Bottom left). Figure 3 (right) also shows increasing thinking steps helps in solving levels with longer optimal solutions. Thank you for this question. We will clarify this on line 157 in the updated paper.
3. That’s correct. Since ResNet doesn’t have a recurrent state, any cycles that it may perform will always render the level unsolvable as it will keep moving in cycles. But it’s also disadvantageous for the DRC to do cycles!
4. Yes, good point. Here is our definition of pacing:
   1. The agent takes more actions than are necessary to get to the next “state-altering” action. In this case: the path taken to the next box-push is longer than optimal
   2. Because of these extra actions, the NN now ‘knows’ what actions to take, whereas before it did not. This suggests a very nice kind of experiment, on top of what we have: compare how accurate the probed plan is before and after the ‘on-policy’ actions of the NN, and compare them to just bee-lining to the next box. Thank you very much for prompting us to think of that!
5. Thank you for suggesting we put this summary at the end of the section. Is it because it’s difficult to understand without having read the section? Should we put them under a “Summary of evidence” header? We thought it would be better to tell the reader what to expect and what claims to look for, in advance, to help with skimming.
6. The standard of evidence is that the probes should be causal as well as predictive. We will add the explanation of the standard of evidence in the updated paper.