ICM Transfer Learning — Maze → Heist

Hypothesis: An agent trained with Intrinsic Curiosity (ICM) on Maze transfers better to Heist than a vanilla PPO agent.

File Structure

icm_project/
├── train.py            # PPO + optional ICM training on Maze
├── transfer.py         # Zero-shot and finetune evaluation on Heist
├── models/
│   ├── __init__.py
│   └── icm.py          # Intrinsic Curiosity Module (Pathak et al. 2017)
├── requirements.txt
└── README.md

---

Install

Python 3.8 is required. Procgen does not support 3.9+ cleanly.

# 1. Create a clean environment
conda create -n icm python=3.8 -y
conda activate icm

# 2. Install dependencies
pip install -r requirements.txt

# 3. Verify Procgen works
python -c "import procgen; print('procgen ok')"

---

Run

Step 1 — Train Model 1: Vanilla PPO on Maze (baseline)

python train.py --env_id maze

Checkpoint saved to: checkpoints/maze__vanilla__1__TIMESTAMP/final.pt

---

Step 2 — Train Model 2: PPO + ICM on Maze (curiosity)

python train.py --env_id maze --use_icm

Checkpoint saved to: checkpoints/maze__icm__1__TIMESTAMP/final.pt

---

Step 3 — Zero-Shot Transfer to Heist

Run both checkpoints through Heist with no training:

# Vanilla baseline
python transfer.py \
    --checkpoint checkpoints/maze__vanilla__1__TIMESTAMP/final.pt \
    --mode zero_shot

# ICM agent
python transfer.py \
    --checkpoint checkpoints/maze__icm__1__TIMESTAMP/final.pt \
    --mode zero_shot

Compare the mean rewards. The ICM agent should score higher.

---

Step 4 — Finetune Transfer (N-layer ablation)

Tests George's suggestion: freeze backbone, train only last N layers on Heist.

# Finetune actor + critic heads only (n=1)
python transfer.py \
    --checkpoint checkpoints/maze__icm__1__TIMESTAMP/final.pt \
    --mode finetune --n_layers 1

# Finetune heads + final linear layer (n=2)
python transfer.py \
    --checkpoint checkpoints/maze__icm__1__TIMESTAMP/final.pt \
    --mode finetune --n_layers 2

---

Optional: GPU + WandB logging

python train.py --env_id maze --use_icm --cuda --track --wandb_project_name icm-transfer

python train.py --env_id maze --cuda --track --wandb_project_name without_icm-transfer

Optional: TensorBoard (always logged)

tensorboard --logdir runs/

---

Experiment Matrix

Run	Command	Purpose
Model 1	train.py --env_id maze	Vanilla PPO baseline
Model 2	train.py --env_id maze --use_icm	Curiosity agent
Transfer A	transfer.py --mode zero_shot (Model 1 ckpt)	Baseline transfer
Transfer B	transfer.py --mode zero_shot (Model 2 ckpt)	Curiosity transfer
Ablation 1	transfer.py --mode finetune --n_layers 1	N-layer test
Ablation 2	transfer.py --mode finetune --n_layers 2	N-layer test

---

Key Metrics to Compare

• Mean reward on Heist (zero-shot) — primary result
• Learning curve on Heist (finetune) — sample efficiency
• Intrinsic reward over time (logged to tensorboard) — did ICM actually explore?
• ICM forward loss — is the model learning the dynamics?

---

Common Errors

ModuleNotFoundError: procgen → Make sure you're in the icm conda env: conda activate icm

gym.error.UnregisteredEnv → You need gym==0.21.0 not gymnasium. Check pip show gym.

CUDA out of memory → Reduce --num_envs 32 or run without --cuda

assert x.shape[1:] == self.get_output_shape() → The ConvSequence shape check failed. This means obs shape changed. Procgen should always be 64x64x3 — check your wrapper.