Sequential prediction in a gridworld: an RNN exploration

Overview

A simplified reproduction of the spatial representation results from the first stage of Prof. Levenstein's paper "Sequential predictive learning is a unifying theory for hippocampal representation and replay", exploring whether a trained RNN in a gridworld environment is capable of developing spatially tuned hidden units. This project was developed independently using Levenstein et al. as the starting point, to reproduce and explore the emergence of spatially tuned representations.

Where to start

The main result of this project is contained in experiment_MultiTrajRNN_spatialtuning_analysis.ipynb, which tests whether spatially tuned hidden representations emerge in the trained RNN and whether these representations remain consistent across independent trajectories.

The other experiment notebooks provide supporting comparisons:

1. experiment_SimpleMLP.ipynb : First implementation, using a simple MLP model
2. experiment_FirstRNN.ipynb : Implemented the RNN model
3. experiment_MultiTrajRNN_Restructured.ipynb : Includes multiple trajectories, RNN-MLP comparison for both standard and turn-heavy action probabilities
4. experiment_MultiTrajRNN_spatialtuning_analysis.ipynb : Analyzes spatial tuning and consistency across independent trajectories.

Additional exploratory and supporting notebooks can be found in the Experiments/ directory.

What I implemented

• Gridworld and agent:

A 14X21 T-shaped environment with visual cues in the form of floor tiles of different colors. The agent can start from a set or randomized position and direction, and can observe a 5X5 cutout of the floor in front of it. Floor tile colors are currently represented by different integer values.

• Trajectory generation:

The agent movement is composed of four fundamental actions: forward, left, right and stop. The action selection at each step is random but can be controlled via the selection of action probabilities and randomization seed. The agent trajectory is generated by chaining these randomized actions by an user selected number of steps.

• Predictive RNN:

A single RNN of 128 hidden units is trained to predict the agent's next observation given its current observation and action selection.

• Spatial Tuning Map:

the activations of the hidden units are collected over the course of the trajectory, then classified based on each agent position and averaged to generate the spatial tuning map.

Experiments and observations

• Initial MLP-RNN comparison:

As a first test run to explore PyTorch's capabilities, the agent prediction was first generated through an MLP, followed by the RNN implementation. The first tests did not feature multi-trajectory generation, instead relying on only one trajectory alone. In addition, the gridworld was smaller at 15X10.

The limited number of observations from one single trajectory caused the models to easily overfit, as can be observed in experiment_FirstRNN.ipynb, where the lowest validator error was at 4.2, increasing significantly afterwards. View the initial RNN experiment

Interestingly, the MLP fares better with the limited data, as can be observed in experiment_SimpleMLP.ipynb. Both the training and validator loss are lower in the MLP. View the initial MLP experiment

• Second MLP-RNN Comparison: multi-trajectory data, gridworld size increase and turn-heavy action probabilities

To test whether insufficient training data was the reason for RNN's reduced performance, The ability to generate multiple trajectories at once was introduced. In addition, the ability to spawn the agent at a random location and direction was added. The added trajectory diversity and random starts coincided with better validation performance and less apparent overfitting. With 20 trajectories of 500 steps each, the model is able to get a full view of the entire gridworld, which allows it to train much more effectively.

The key distinction between the MLP and the RNN is the fact that the RNN is capable of integrating previous observations into its current state through its hidden units. This accumulated history allows the agent to infer its location and orientation in a way that would not be available through observation alone.

One of the most telling ways to see if the agent is capable of inferring its location and orientation is by observing the model's loss for turn actions. Since turning causes the entire observation surface to change, its correct prediction is much more reliant on having good information on its current location. Two tests were performed to support this point: first, a second batch of trajectories was run and a new model constructed, but with an increased probability of turn actions, and no stop actions. While the standard action probability run had both RNN and MLP performing similarly, the turn-biased trajectories widened the RNN-MLP performance gap, lowering the RNN validator loss relative to the MLP. Second, by breaking loss down based on action type, we can differentiate each action's loss contribution. It is apparent that while the forward and stop actions perform better through the MLP, the turn actions perform much better through the RNN. View the standard vs. turn heavy action probabilities comparison experiment

	Trainer	Validator
RNN, standard	0.733	0.925
RNN, turn heavy	0.898	1.013
MLP, standard	0.846	0.984
MLP, turn heavy	1.429	1.529

Table 1: trainer and validator mean squared error (MSE), standard action probabilities vs. turn heavy action probabilities. The MLP performs significantly worse once the probability of turning is much higher.

	Forward	Left	Right	Stop
trainer, RNN	0.316	1.509	1.481	0.933
trainer, MLP	0.276	2.037	2.006	0.727
validator, RNN	0.419	1.885	2.019	1.135
validator, MLP	0.343	2.425	2.418	0.884

Table 2: trainer and validator MSE by action type for standard action probability trajectory runs.

• Hidden unit activation grid and detection of spatial tuning

The central question of this project is whether some RNN hidden units develop location-dependent activity. For each unit, A spatial activation map was constructed by averaging its activation over all visits to each grid position. Two independent batches of trajectories were then generated and a map computed from each batch using the same trained RNN. The selected maps show similar regions of elevated activity across batches, suggesting that these patterns are reasonably robust to the particular trajectories sampled.

Next, we calculated the explained variance $%EV$ for each neuron, which, similar to Prof. Levenstein's paper, consists of calculating the variance of activations of each neuron, comparing it to the variance of its activation minus their average activation at each position, and testing whether subtracting its average activation removes most of its activation variance. Out of the 128 neurons, 11 of them had an $%EV$ above $50%$, which is the cutoff value from the paper. Comparing their activation maps shows a very apparent picture: each neuron does seem to activate according to slighly different but spatial criteria: some activate in the upper region, some activate when within the colored tile floors, etc.

How to run

From the repository root, install the packages in requirements.txt and start Jupyter in the Experiments/ folder. Open the spatial tuning analysis notebook, select a Python kernel with those packages installed, and run all cells from a fresh kernel. Make sure to run the .ipynb notebooks from the Experiments/ folder, since the code makes its parent folder, containing the relevant dependencies, as the root.

Full Experiment Settings

The gridworld is generated manually, by first setting an initial outer boundary, followed by defining the agent movement area, then marking the floor tiles according to different numerical values. Regions out of the gridworld are defined with a value of -1. It is recommended to make sure the gridworld padding is set to equal or greater to the agent's grid view. The agent movement area is set through the add_rectangle_region() function. There are multiple possible shapes to choose to mark the floor tiles, including squares, triangles, and rectangles. The results shown earlier were generated with the settings below. These are the values used for the reported experiment.

Gridworld creation

Parameter	What it controls	Value used
HEIGHT	Height of the Gridworld environment (including padding)	29
WIDTH	Width of the Gridworld environment (including padding)	36

Agent creation

Parameter	What it controls	Value used
VIEW_SIZE	Agent's NxN view of the grid floor	5
START_X	Agent's starting X position (vertical)	7 (Overriden by RANDOM_START)
START_Y	Agent's starting Y position (horizontal)	7 (Overriden by RANDOM_START)
START_DIRECTION	Agent's start direction (0, 1, 2, 3 correspond to N, E, S, W, respectively)	2 (Overriden by RANDOM_START)
RANDOM_START	Let the agent's starting position be random	True
ACTION_PROBS	Probability of forward, left, right, and stop movement	0.6, 0.15, 0.15, 0.1
ACTION_PROBS_TURN_HEAVY	Probability of forward, left, right, and stop movement (No stop, turn heavy)	0.3, 0.35, 0.35, 0

Trajectory creation

Parameter	What it controls	Value used
TRAINER_START_SEED	Seed to determine trainer agent's random start	150
VALIDATOR_START_SEED	Seed to determine validator agent's random start	350
RANDOM_START	Let the agent's starting position be random	True
TRAINER_MOVE_SEED	Seed to determine trainer agent's random movement	100
VALIDATOR_MOVE_SEED	Seed to determine validator agent's random movement	300
N_STEPS	Number of steps in each trainer trajectory	500
N_TRAJ	Number of trajectories for each trainer batch	20

Model and training

Parameter	What it controls	Value used
HIDDEN_SIZE	Unit count of the RNN layer	128
TRAINING_EPOCHS	Passes through the training data	3000

Spatial tuning analysis

Parameter	What it controls	Value used
STUNING_START_SEED_1	1st spatial tuning analysis agent random starting position seed	5
STUNING_MOVE_SEED_1	1st spatial tuning analysis agent movement seed	10
STUNING_START_SEED_2	2nd spatial tuning analysis agent random starting position seed	6
STUNING_MOVE_SEED_2	2nd spatial tuning analysis agent movement seed	11

Reference

Sequential predictive learning is a unifying theory for hippocampal representation and replay Daniel Levenstein, Aleksei Efremov, Roy Henha Eyono, Adrien Peyrache, Blake Richards bioRxiv 2024.04.28.591528; doi: https://doi.org/10.1101/2024.04.28.591528