DQN: value learning that beats the floor (and then hits a ceiling)

Code: rl_ablations/models/dqn, built on the rl-factory engine.

This is the second post in the series (rl-factory came first — the framework everything plugs into). It’s the first agent that actually learns. Before this, the only baseline was a random agent: on World 1-1 it reaches a mean max_x ≈ 665 and completes the level 0% of the time (the flagpole is at x ≈ 3160, so random Mario gets about a fifth of the way). That’s the floor. DQN’s job is to beat it.

If you remember one sentence: DQN learns a function Q(s, a) — “how good is pressing button a in state s?” — and acts by picking the button with the highest Q; the entire trick is making that estimate stable enough to train against itself.

Part 0 — where DQN sits in the factory

From the previous post: the one real axis of variation is on-policy vs. off-policy, and DQN is the off-policy family. It inherits OffPolicyStrategy unchanged — async-ish collection into a replay buffer, learning from sampled minibatches. So this post is not about a training loop (that’s shared); it’s about one Learner — DQNLearner — and the network it trains.

Part 1 — the Q-network

The network is the Nature-DQN conv stack: a stacked (4, 84, 84) frame (four grayscale frames, so the agent can see motion) in, one Q-value per action out.

self.conv = nn.Sequential(
    nn.Conv2d(in_channels, 32, kernel_size=8, stride=4), nn.ReLU(),  # 84x84 -> 20x20
    nn.Conv2d(32, 64, kernel_size=4, stride=2), nn.ReLU(),           # 20x20 -> 9x9
    nn.Conv2d(64, 64, kernel_size=3, stride=1), nn.ReLU(),           # 9x9   -> 7x7
    nn.Flatten(),
)
self.head = nn.Sequential(
    nn.Linear(64 * 7 * 7, 512), nn.ReLU(),
    nn.Linear(512, n_actions),   # one Q-value per action; no activation
)

No softmax, no value in [0,1] — the outputs are raw expected returns, one per button combination. Acting is argmax over them (with ε-greedy exploration: a small fraction of the time, press a random button so the agent keeps discovering).

Part 2 — the three ideas that make it trainable

Q-learning’s target is bootstrapped: you regress Q(s, a) toward r + γ·maxₐ' Q(s', a') — an estimate that uses the network’s own output. Do that naïvely and it diverges: you’re chasing a target that moves every time you take a step. DQN’s contribution is three stabilizers, all visible in DQNLearner:

1. Experience replay. Store transitions in a buffer and learn from random minibatches, not the latest step. This breaks the correlation between consecutive frames (which would otherwise make each gradient step redundant with the last) and lets each transition train the network many times.

2. A frozen target network. Keep a second copy of the Q-network, self._target, and compute the bootstrap target from it, not the live network. It’s a stationary goalpost:

self._target = copy.deepcopy(self._q)   # frozen goalpost
...
with torch.no_grad():
    q_next = self._target(next_states).max(dim=1).values
    td_target = rewards + config.gamma * q_next * (1.0 - dones)   # 0 bootstrap on terminal

Every target_sync_every steps the target is refreshed from the live network. Between syncs, the thing you’re regressing toward holds still — which is the difference between training and diverging.

3. Huber loss + reward clipping. The TD error is squashed with smooth_l1_loss (Huber) and rewards are clipped to [-1, 1], so a single surprising transition can’t produce a giant gradient. Mario’s raw reward has big spikes; this keeps the update well-behaved.

The learner self-gates the whole schedule — warm up the buffer (learning_starts), then a gradient step every train_every, and a target sync on its own cadence:

def learn(self):
    self._steps += 1
    metrics = self._gradient_step() if self._is_train_step() else None
    if self._is_sync_step():
        self._sync_target()
    return metrics

That’s DQN in one method: replay-sample, regress toward the frozen-target TD value, periodically move the goalpost.

Part 3 — does it beat the floor?

Yes — comfortably. Against the random baseline’s max_x ≈ 665, DQN climbs to a mean episode return around 2774 on Mario 1-1. It learns to run right, jump the first pits and Goombas, and make real progress into the level. Here it is against the PPO agent we build next:

Read the blue curve first: DQN rises off the floor fast, then plateaus around ~2800 by roughly 8M steps and stops improving. It’s a real agent — but it’s stuck.

Part 4 — the ceiling, and why we move on

DQN plateaus for reasons that are structural, not just tuning:

• It’s value-based on a big discrete action space. Every improvement has to route through an argmax over Q-values; the policy is implicit and changes in lurches.
• Off-policy replay mixes stale experience. The buffer is full of transitions from an older, worse policy, which is efficient but noisy for a fast-moving control problem like Mario.
• Sparse, delayed reward is hard for TD. Progress in Mario often needs a committed sequence of jumps before any reward arrives; bootstrapped value estimates propagate that credit slowly.

The orange curve above is the fix: PPO, an on-policy policy-gradient method that optimizes the policy directly and keeps climbing past DQN’s plateau to ~4500. That’s the next post — and, in factory terms, it’s the other strategy family. Same Trainer, same env, a different half of the split.