Frozen Vision Encoders for Deep RL in VizDoom

Overview

This project trains a Double DQN agent to play first-person shooter scenarios in VizDoom, an RL research platform built on the Doom engine. The agent observes raw RGB frames, encodes them with a frozen pretrained vision transformer, and learns Q-values over a small discrete action space.

The design rests on one idea: the agent should not have to learn to see and to act at once. Large self-supervised vision models (AIMv2, V-JEPA 2) already turn raw pixels into good features. Freeze one, train a small MLP Q-network on top, and training comes out fast, stable, and light on memory.

Key Findings

• AIMv2 works — mean evaluation return of ~6 on DefendTheCenter after 500k steps (6.5 and 5.6 across the two seeds).
• V-JEPA 2 diverges — returns never leave zero on either pooling variant, and Q-values run away: past 60 mean / 400 max before the mean-pool run died at 83k steps, past 3,000 mean / 6,000 max over the last-pool run's full 500k. The likely reason: the choppy frame-skip-4 clips look nothing like the video it was pretrained on.
• Slower target updates help — $\tau = 0.001$ gives slightly higher returns than $\tau = 0.005$ (6.5 vs 6.1) and keeps Q-estimates in range (peak mean Q of 8 vs 23).
• Dueling is inconclusive — those runs crashed at 66k and 178k steps; up to there the standard Q-network stayed ahead.

Encoder Comparison: AIMv2 vs V-JEPA 2

AIMv2-Large beats V-JEPA 2-Large outright, and neither V-JEPA 2 pooling variant ever leaves zero return. The mean-pool run crashed at 83k steps with Q-values already running away (mean above 60, max above 400). The last-pool run survived all 500k steps and diverged completely: mean Q past 3,000, max past 6,000, returns drifting below zero. The likely cause is the frame rate. V-JEPA 2 was pretrained on video; with frame skip 4, the eight-frame clips it sees here have frames far apart in time, and the smooth motion it expects is gone.

Vision Encoders

The code supports two pretrained encoders:

Both come from HuggingFace Transformers and stay frozen throughout. AIMv2 encodes each frame on its own. V-JEPA 2 reads a sliding window of the last 8 frames; the pipeline averages its spatiotemporal patch tokens over time, then over space, leaving one 1024-dimensional vector per observation.

Feature Pipeline

Raw encoder outputs are 1024-dimensional and unwhitened. Before training, a calibration phase collects observations from a random policy, encodes them in batches, and fits a PCA whitening transform: 1024 dimensions down to $d = 256$, decorrelated, variance normalized. The whitened features then pass through LayerNorm and L2 normalization, and the replay buffer stores the result — 256-dimensional unit vectors.

Q-Network

The Q-network is a 2-layer MLP with hidden size 256, one output per action. An optional dueling variant splits state value from action advantage:

Training

The agent uses Double DQN: the online network picks the action, the target network scores it. After each gradient step, Polyak averaging nudges the target network toward the online one ($\theta^- \leftarrow \tau \, \theta + (1 - \tau) \, \theta^-$). Epsilon-greedy exploration, annealed linearly, fills a replay buffer of 100k transitions.

Environment

Training uses VizDoom's DefendTheCenter scenario through a Gymnasium wrapper: the agent stands at the center of a circular room, ammunition limited, and shoots enemies as they close in.

Results

All experiments use DefendTheCenter with frame skip 4, a 100k-transition replay buffer, batch size 256, and 500k training steps. Evaluation runs 20 greedy episodes every 5k steps.

Training Curves

Ablation: Target Update Rate

Slower target updates ($\tau = 0.001$) give slightly higher, steadier evaluation returns than the default ($\tau = 0.005$).

Standard vs Dueling Q-Network

The dueling runs never reached 500k steps — they crashed at 66k and 178k. Up to those points the standard Q-network with $\tau = 0.005$ stayed ahead of both dueling variants (4.9 vs 3.9 at 178k) — though the dueling curves were still rising when they stopped.

Hyperparameters