Q-learning and DQN

Deep dive

How it works

Parts:

• online Q network — maps a state to one value per action; the only part that is trained
• target network — frozen copy of the online net, synced every C updates, computes the bootstrap target
• replay ring buffer — stores (s, a, r, s′, done), sampled uniformly
• ε-greedy policy — explores with probability ε, else takes the online argmax
• Huber loss — squared error for small errors, linear for large ones

Tabular update with step size α:

Q(s,a) ← Q(s,a) + α·(r + γ·(1 − done)·max_a′ Q(s′,a′) − Q(s,a))

Operations (DQN loop):

1. Act ε-greedily with the online network. Store (s, a, r, s′, done) in the replay ring buffer.
2. Every 4 environment steps, sample a uniform minibatch of 32.
3. Targets y = r + γ·(1 − done)·Q_target(s′, argmax_a Q_online(s′, a)) (Double DQN), with no gradient through y.
4. Loss = Huber(Q_online(s, a) − y), only on the taken action's output. Gradient step.
5. Every C updates copy online → target (or Polyak-average with τ of order 0.005 per step).
6. Anneal ε linearly, then keep a small floor (0.01-0.1).

env ──(s,a,r,s′,done)──► replay ring (1M) ──uniform batch──► online Q ──loss──► grads
                                                    └──────► target Q (copy every C) ─► y

The deadly triad (function approximation + bootstrapping + off-policy data) can diverge. Replay, target networks and Huber loss are the practical mitigations, not guarantees.

Code

/// DQN machinery on a linear Q-function over one-hot states: chain of N states, actions {left, right},
/// reward +1 on reaching the right end (terminal), -0.01 per step otherwise.
const N: usize = 12;
const A: usize = 2;

struct Rng(u64);
impl Rng { fn f(&mut self) -> f64 { self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17;
                                     (self.0 >> 11) as f64 / (1u64 << 53) as f64 } }

#[derive(Clone, Copy)]
struct Tr { s: usize, a: usize, r: f32, s2: usize, done: bool }

struct Replay { buf: Vec<Tr>, cap: usize, next: usize }
impl Replay {
    fn push(&mut self, t: Tr) {                               // ring buffer: overwrite the oldest
        if self.buf.len() < self.cap { self.buf.push(t) } else { self.buf[self.next] = t }
        self.next = (self.next + 1) % self.cap;
    }
}

fn step(s: usize, a: usize) -> (usize, f32, bool) {
    let s2 = if a == 1 { s + 1 } else { s.saturating_sub(1) };
    if s2 == N - 1 { (s2, 1.0, true) } else { (s2, -0.01, false) }
}

fn argmax(q: &[f32]) -> usize { if q[1] > q[0] { 1 } else { 0 } }

/// Huber gradient w.r.t. the prediction.
fn huber_grad(err: f32, delta: f32) -> f32 { err.clamp(-delta, delta) }

fn main() {
    let (gamma, lr, batch, sync_every, train_every) = (0.99f32, 0.1f32, 32, 200, 4);
    let mut online = vec![0.0f32; N * A];                      // Q[s][a] = W[s*A + a] for one-hot s
    let mut target = online.clone();
    let mut rb = Replay { buf: Vec::new(), cap: 10_000, next: 0 };
    let mut rng = Rng(2024);
    let (mut s, mut t, mut updates, mut ep_len, mut lens) = (0usize, 0usize, 0usize, 0, Vec::new());
    while t < 30_000 {
        let eps = (1.0 - t as f64 / 10_000.0).max(0.05);    // linear anneal, then a floor
        let a = if rng.f() < eps { (rng.f() * A as f64) as usize } else { argmax(&online[s * A..s * A + A]) };
        let (s2, r, done) = step(s, a);
        let timeout = ep_len + 1 >= 200;                       // truncation is NOT termination
        rb.push(Tr { s, a, r, s2, done });
        s = if done || timeout { lens.push(ep_len + 1); ep_len = 0; 0 } else { ep_len += 1; s2 };
        t += 1;
        if rb.buf.len() < 500 || t % train_every != 0 { continue; }   // warm-up, update frequency
        let mut grad = vec![0.0f32; N * A];
        for _ in 0..batch {
            let tr = rb.buf[(rng.f() * rb.buf.len() as f64) as usize];
            let a_star = argmax(&online[tr.s2 * A..tr.s2 * A + A]);  // Double DQN: online picks,
            let boot = target[tr.s2 * A + a_star];                   // target evaluates
            let y = tr.r + if tr.done { 0.0 } else { gamma * boot };
            let err = online[tr.s * A + tr.a] - y;
            grad[tr.s * A + tr.a] += huber_grad(err, 1.0) / batch as f32; // only the taken action
        }
        for (w, g) in online.iter_mut().zip(&grad) { *w -= lr * g; }
        updates += 1;
        if updates % sync_every == 0 { target.copy_from_slice(&online); }
    }
    let greedy: Vec<usize> = (0..N - 1).map(|s| argmax(&online[s * A..s * A + A])).collect();
    let recent = &lens[lens.len() - 20..];
    println!("greedy policy (1 = right): {greedy:?}");
    println!("mean episode length, last 20: {:.1} (optimal {})", recent.iter().sum::<usize>() as f32 / 20.0, N - 1);
    println!("Q(0, right) = {:.3}, ideal ≈ {:.3}", online[1],
             (0..N - 2).map(|k| -0.01 * gamma.powi(k as i32)).sum::<f32>() + gamma.powi(N as i32 - 2));
}

Complexity and constants

• Per update: forward online on s (and s′ for Double DQN), forward target on s′, backward on s. About 3 forwards + 1 backward per sample.
• Atari replay of 1M transitions: an 84×84 uint8 frame is 7 KB, so about 7 GB if each frame is stored once and stacks of 4 are rebuilt by index. Storing stacked observations for s and s′ separately costs 8x more.
• DQN trained for 50M frames per game; later papers often count 200M emulator frames (50M agent steps under frame skip 4). Sample efficiency is low compared to model-based methods.
• Replay ratio (gradient updates per environment step): DQN uses 1 update of 32 samples per 4 steps, so each transition is replayed about 8 times on average.

Efficiency

• Cost: about 3 forwards and 1 backward per sampled transition. An Atari replay of 1M transitions is about 7 GB with each 84×84 uint8 frame stored once. DQN used 1 update of 32 samples per 4 steps, so each transition is replayed about 8 times. Bound: environment steps (50M frames per game) and replay memory.
• Levers:
  • store each frame once, rebuild stacks by index — 8x less replay memory than stacked s and s′ (above) — pixel inputs
  • uint8 storage, normalize on the GPU — 4x less than fp32 (arithmetic) — pixel inputs
  • raise the replay ratio or use n-step returns — more learning per environment step (Hessel et al. 2018) — expensive environments
  • prioritized replay — samples high-TD-error transitions more often (Schaul et al. 2016) — sparse rewards
  • sync the target network every C steps, not every step — the target forward needs no gradient (above) — always
• Naive → standard → state of the art:
  • online Q-learning on each transition, no replay · 1 use per sample, unstable
  • DQN with replay and target network · 1 update of 32 per 4 steps
  • Rainbow components (double, dueling, prioritized, n-step, distributional) · better data efficiency per frame (Hessel et al. 2018)

Implementation notes

• Terminal vs truncated: set done only on true termination. On a time-limit cut, still bootstrap from s′; otherwise the agent learns that time runs out at random states.
• Normalize or clip rewards (DQN clipped to [−1, 1]) so one learning rate fits all games, or use value rescaling.
• Later variants switched to Adam with a large ε (Rainbow: lr 6.25e-5, ε 1.5e-4) and often clip the gradient norm.
• Prioritized replay: sample in proportion to |δ|^α (α = 0.6) via a sum-tree, correct with importance weights (N·P)^−β, β annealed 0.4 → 1.
• Store observations as uint8 and convert to float on the GPU. The buffer is the main memory cost.
• n-step returns (n = 3 in Rainbow) propagate reward faster, at the price of off-policy bias without correction.

Pitfalls

• Bootstrapping through terminal states (missing (1 − done)): values blow up toward r/(1 − γ).
• Gradient through the target: the loss chases itself. Detach y.
• Target network synced too often (or no target network): oscillation and divergence.
• ε annealed to 0: the agent stops exploring and gets stuck. Keep a floor.
• Comparing to papers with different frame-counting conventions (frames vs agent steps under frame skip).
• Overestimation from the max with noisy Q values. Use Double DQN by default.

Variants and alternatives

In production

Atari and game-playing benchmarks (DQN, Rainbow, Agent57), Dopamine and CleanRL reference implementations, recommendation and ad-bidding systems with discrete actions, Stable-Baselines3 DQN.

Implementations

Sources

• Watkins and Dayan, Q-learning, Machine Learning 1992. paper
• Mnih et al., Human-level Control through Deep Reinforcement Learning, Nature 2015. paper
• van Hasselt, Guez, Silver, Deep Reinforcement Learning with Double Q-learning, AAAI 2016. paper
• Wang et al., Dueling Network Architectures for Deep Reinforcement Learning, ICML 2016. paper
• Schaul, Quan, Antonoglou, Silver, Prioritized Experience Replay, ICLR 2016. paper
• Hessel et al., Rainbow: Combining Improvements in Deep Reinforcement Learning, AAAI 2018. paper
• Bellemare, Dabney, Munos, A Distributional Perspective on Reinforcement Learning, ICML 2017. paper
• Sutton and Barto, Reinforcement Learning: An Introduction, 2nd ed., MIT Press 2018. book
• Boyan and Littman, Packet Routing in Dynamically Changing Networks: A Reinforcement Learning Approach, NeurIPS 1993.
• Crites and Barto, Improving Elevator Performance Using Reinforcement Learning, NeurIPS 1995.