Path Channels and Plan Extension Kernels: a Mechanistic Description of Planning in a Sokoban RNN

TL;DR

This paper reveals how a convolutional RNN trained with reinforcement learning for Sokoban encodes plans as activations in specific channels and uses learned kernels to perform bidirectional planning and backtracking.

Contribution

It provides a mechanistic understanding of how a model-free reinforcement learning RNN encodes and constructs plans in Sokoban through specialized channels and kernels.

Findings

Path channels encode future move plans.

Kernels encode transition models for position changes.

Backtracking is implemented via negative values in kernels.

Abstract

We partially reverse-engineer a convolutional recurrent neural network (RNN) trained with model-free reinforcement learning to play the box-pushing game Sokoban. We find that the RNN stores future moves (plans) as activations in particular channels of the hidden state, which we call path channels. A high activation in a particular location means that, when a box is in that location, it will get pushed in the channel's assigned direction. We examine the convolutional kernels between path channels and find that they encode the change in position resulting from each possible action, thus representing part of a learned transition model. The RNN constructs plans by starting at the boxes and goals. These kernels extend activations in path channels forwards from boxes and backwards from the goal. Negative values are placed in channels at obstacles. This causes the extension kernels to propagate the negative value in reverse, thus pruning the last few steps and letting an alternative plan emerge; a form of backtracking. Our work shows that, a precise understanding of the plan representation allows us to directly understand the bidirectional planning-like algorithm learned by model-free training in more familiar terms.