# Sequential Predictive Gridworld (RNN, k-step prediction)

A PyTorch implementation of a GRU-based RNN trained for k-step sequential prediction in a long, narrow gridworld with repeating visual cues.

## Features

- Long corridor gridworld with repeating color patterns
- Random walk agent for trajectory sampling
- GRU-based RNN for k-step sequential prediction
- PCA visualization of learned hidden representations
- Full PyTorch implementation with no gym dependencies

## Installation

```bash
python -m venv .venv && source .venv/bin/activate
pip install -e .    # or: pip install -r requirements.txt
```

## Usage

### Training

```bash
python -m src.train --epochs 10 --n_traj_train 200 --n_traj_val 30 --T 40 --k 3
```

### PCA Evaluation

```bash
python -m src.eval.pca_hidden --ckpt outputs/checkpoints/best.pt --T_eval 800 --k 3
```

### Quick Demo

```bash
# Train (quick demo)
python -m src.train --epochs 8 --n_traj_train 220 --n_traj_val 30 --T 40 --k 3 --hidden 128

# Visualize hidden states
python -m src.eval.pca_hidden --ckpt outputs/checkpoints/best.pt --T_eval 800 --k 3
```

## Expected Outputs

- `outputs/checkpoints/best.pt` - Trained model checkpoint
- `outputs/figures/hidden_pca.png` - PCA visualization of hidden states colored by x-position

## Reproducibility

All random seeds are set deterministically. Use the `--seed` flag to control randomness across runs.

## Testing

```bash
pytest -q
```
