Skip to content

Repository files navigation

SNN Reverse Replay for Reinforcement Learning

A spiking neural network (SNN) implementation of hippocampal reverse replay for reinforcement learning. This project models place cells and action cells in a grid world environment, demonstrating how reverse replay events can accelerate learning through eligibility trace-based plasticity.

Overview

This implementation simulates:

  • Place Cells (100 neurons in a 10x10 grid): Represent spatial locations with Gaussian place fields
  • Action Cells (72 neurons): Encode movement directions (or 4 cardinal directions in simplified mode)
  • Reverse Replay: Hippocampal-inspired replay events triggered at reward delivery, propagating activity backward through the experienced trajectory
  • Eligibility Traces: Temporal credit assignment mechanism linking actions to delayed rewards
  • Spike-Timing Dependent Plasticity: Weight updates driven by reward prediction errors

The model supports both rate-based and spiking (Poisson/LIF) place cell representations.

Project Structure

snn-reverse-replay/
├── main.py                 # Main entry point - runs the grid world simulation
├── network.py              # Rate-based neural network implementation
├── spiking_network.py      # Spiking neural network (Poisson place cells, LIF action cells)
├── replay_agent.py         # Rate-based agent with reverse replay
├── spiking_replay_agent.py # Spiking agent with reverse replay
├── grid_world.py           # 10x10 grid environment with visualization
├── functions.py            # Configuration utilities (reads options.json)
├── options.json            # Experiment configuration
├── live_plotting.py        # Real-time visualization of network activity
├── plot_raster.py          # Raster plot visualization for spiking activity
├── trial_times_plotting.py # Post-hoc analysis plotting
├── main.sh                 # Local execution script with live plotting
├── HPC.sh                  # HPC cluster submission script (SLURM)
└── all_results.sh          # Batch execution for multiple parameter sweeps

Installation

Requirements

  • Python 3.8+
  • NumPy
  • Matplotlib
  • Pygame

Install Dependencies

pip install numpy matplotlib pygame

Usage

Basic Execution

# Run with rate-based network
python main.py

# Run with spiking network
python main.py --spiking

# Run with custom options file
python main.py -o custom_options.json

# Set random seed
python main.py -s 42

# Specify experiment number
python main.py -e 5

# Run with live plotting
./main.sh

Command Line Options

Option Description Default
-o, --options FILE Options JSON file options.json
-s, --seed SEED Random seed None
-e, --experiment-number NUM Experiment identifier 0
--spiking Use spiking neural network False
-h, --help Show help message -

Configuration (options.json)

{
    "render": false,           // Enable pygame visualization
    "num_trials": 30,          // Number of trials per experiment
    "grid_size": [10, 10],     // Environment dimensions
    "num_pcs": 100,            // Number of place cells (must be perfect square)
    "num_acs": 72,             // Number of action cells (72 directions or 4 cardinal)
    "tau_e": [5, 1, 0.2, 0.04],// Eligibility trace time constants (s)
    "eta": [10, 1, 0.1, 0.01, 0.001], // Learning rates
    "plots": true              // Enable live plotting
}

Running Parameter Sweeps

The all_results.sh script runs experiments across all combinations of tau_e and eta values:

chmod +x all_results.sh
./all_results.sh

Results are saved to data/results/trial_times_tau_e{X}_eta{Y}.csv.

HPC Execution

For cluster deployment (SLURM):

sbatch HPC.sh

Model Architecture

Place Cell Network

  • 100 place cells arranged in a 10x10 topological grid
  • Gaussian place fields centered on each grid location
  • Lateral connections to 8 neighbors with STP (short-term plasticity)
  • Intrinsic plasticity for homeostatic rate regulation

Action Cell Network

  • 72 action cells (5-degree angular resolution) or 4 cardinal directions
  • Sigmoidal activation from weighted place cell input
  • LIF dynamics (spiking mode) with refractory period

Learning Mechanism

  1. Eligibility Trace: e_ij(t) tracks co-activation of action cell i and place cell j
  2. Reverse Replay: Upon reward, place cells reactivate in reverse order
  3. Weight Update: Δw_ij = η * RPE * (1/σ²) * e_ij

Output

Data Files

During execution, the following files are saved to data/:

File Contents
weights.npy Place cell → action cell weights
rates_data.npy Place cell firing rates
place_data.npy Place cell input currents (I_place)
action_cells_vals.npy Action cell activity
eligibility_trace.npy Eligibility trace matrix
intrinsic_e.npy Intrinsic excitability values
spikes.npy Spike indicators (spiking mode)
spikes_raster.npy Raster plot data (spiking mode)

Results

Trial completion times are saved to:

data/results/trial_times_tau_e{X}_eta{Y}.csv

Each row contains: [experiment_number, trial_1_time, trial_2_time, ...]

Visualization

Live Plotting

Run live_plotting.py alongside the simulation to view:

  • Network rates heatmap
  • Intrinsic plasticity values
  • Weight vector field
  • Eligibility trace vector field
python live_plotting.py &
python main.py

Raster Plot

For spiking simulations:

python plot_raster.py

Key Parameters

Parameter Symbol Typical Range Description
Eligibility trace time constant τ_e 0.01 - 5 s Memory duration for credit assignment
Learning rate η 0.001 - 10 Weight update magnitude
Place field width d 0.5 - 1.0 Spatial tuning of place cells
Action noise σ 0.1 - 0.3 Exploration noise
STP time constants τ_STD, τ_STF 1.0 - 1.5 s Short-term depression/facilitation

Citing

If you use this code in your research, please cite:

@software{snn_reverse_replay,
  title = {SNN Reverse Replay for Reinforcement Learning},
  year = {2026},
  url = {https://github.com/daveydhruti/snn-reverse-replay}
}

@mscthesis{davey2025hippocampal,
  title = {SNN Hippocampal Reverse Replay for Reinforcement Learning},
  author = {Davey, Dhruti},
  year = {2025},
  url = {https://daveydhruti.gitlab.io/documents/2025/hippocampal_replay.pdf}
}

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgements

  • The project was was on the rate-based model, robotic_RL_replay by M. T. Whelan, et. al.
  • This work was developed as part of my dissertation research

About

SNN implementation of a model for hippocampal reverse replay for reinforcement learning.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages