Q-learning and DQN

What it is and when you reach for it

The other route to a policy, taken by policy gradients, is to parameterize behavior directly and ascend expected return. Q-learning takes the value route: learn how much discounted reward each action is worth in each state, assuming optimal play afterwards, and act greedily on those estimates. That framing buys two things the policy gradient family cannot offer. First, the learned quantity satisfies a fixed-point equation, the Bellman optimality equation, so every observed transition is a usable training signal regardless of which policy produced it; experience can be stored and replayed for years without going stale. Second, there is no sampling variance from a stochastic policy in the update itself; the noise comes only from the environment. The costs are the mirror image: the argmax needs a discrete (or discretized) action space, and combining bootstrapped targets with function approximation off-policy is famously treacherous. Q-learning (Watkins, 1989) is the tabular algorithm; DQN (Mnih et al., 2013, and the 2015 Nature paper) is the same update carried by a deep network plus the two stabilizers that made it work on Atari from raw pixels. Its on-policy cousins that blend value learning with an explicit policy live on the actor-critic page.

The math

Bellman optimality

Define Q*(s, a) as the expected discounted return from taking action a in state s and behaving optimally forever after. The definition is recursive, because "behaving optimally after" means taking the best action in the next state, and unrolling one step gives the Bellman optimality equation:

Q*(s, a) = E_{s′∼P(·|s,a)}[ r(s, a) + γ max_a′ Q*(s′, a′) ].

The max inside the expectation is the whole character of the method: it says the value of an action is its immediate reward plus the value of the best continuation, not the continuation your current behavior would actually pick. Q* is the unique fixed point of this equation, and the operator that maps any Q to its right-hand side is a γ-contraction in the max norm: applying it shrinks the distance to Q* by a factor γ. That contraction is why iterating the equation converges from any starting point, and once you have Q*, the optimal policy is simply π*(s) = argmax_a Q*(s, a). No separate policy object ever needs to exist.

The tabular update

With a known model you could iterate the Bellman operator exactly. Q-learning is the model-free, sample-based version: each observed transition (s, a, r, s′) is a one-sample estimate of the right-hand side, and the update moves the table entry a step of size α toward it:

Q(s, a) ← Q(s, a) + α [ r + γ max_a′ Q(s′, a′) − Q(s, a) ].

The bracketed quantity is the temporal-difference error: the gap between what the table currently claims and what one step of fresh evidence plus the table's own view of the future suggests. Using the table's own estimate of the next state inside the target is called bootstrapping, and it is both the source of Q-learning's efficiency and, later, the first ingredient of its instability. For terminal transitions the future term is dropped and the target is just r. Watkins and Dayan proved the tabular version converges to Q* with probability 1 provided every state-action pair is visited infinitely often and the step sizes decay appropriately; note what is absent from those conditions: nothing about which policy generates the visits.

A worked gridworld

Take a three-cell corridor, s₀ → s₁ → s₂, where s₂ is terminal, stepping right into it pays reward 1, every other reward is 0, γ = 0.9, α = 0.5, and the table starts at zero. Episode one starts at s₁ and steps right into the goal. The target is r = 1 (terminal, no future term), the TD error is 1 − 0 = 1, and the update gives Q(s₁, R) ← 0 + 0.5 · 1 = 0.5. Episode two starts at s₀. Stepping right to s₁: target = 0 + 0.9 · max(0, 0.5) = 0.45, so Q(s₀, R) ← 0 + 0.5 · 0.45 = 0.225. Stepping right again: target = 1, TD error = 1 − 0.5 = 0.5, so Q(s₁, R) ← 0.5 + 0.5 · 0.5 = 0.75. The picture to keep: value leaks backward from the reward, one state per visit, at a rate set by α and discounted by γ, and the true values here, Q*(s₁, R) = 1 and Q*(s₀, R) = 0.9, are approached geometrically as episodes repeat.

Off-policy learning and epsilon-greedy

Look again at the update: nowhere does it ask which policy chose a, and the max in the target evaluates the greedy continuation even if the agent then goes on to explore. Q-learning therefore learns about the greedy policy while behaving however it likes; it is off-policy. The behavior policy only has to keep visiting everything, and the standard choice is epsilon-greedy: with probability ε take a uniformly random action, otherwise take argmax_a Q(s, a), with ε annealed from 1.0 toward a small floor like 0.05 as the table firms up. Compare SARSA, which plugs the action actually taken next into the target instead of the max: that seemingly small change makes it on-policy, learning the value of the exploring policy itself. Off-policy learning is what makes experience replay legal at all, and it is the property the on-policy policy gradient family gives up in exchange for direct optimization of behavior.

DQN: two fixes, one deadly triad

A table cannot represent 210×160 Atari frames, so DQN replaces it with a network Q_θ(s, ·) and turns the update into a regression: minimize a Huber loss between Q_θ(s, a) and the target r + γ max_a′ Q(s′, a′). Done naively this diverges routinely, and the reason has a name. Sutton and Barto call bootstrapping, function approximation, and off-policy learning the deadly triad: any two are safe, all three together void every convergence guarantee. The mechanism is concrete. With a table, updating Q(s, a) touches one cell. With a network, updating Q_θ(s, a) moves the value of every nearby state, including s′, which sits inside the target you were regressing toward, so the target moves in the same direction as the prediction and errors can compound in a feedback loop instead of contracting. Off-policy data distribution mismatch means the states where errors grow may never be corrected by fresh visits.

DQN's two fixes each cut one loop of that feedback. The replay buffer stores transitions in a large ring and trains on uniformly sampled minibatches. This breaks the temporal correlation of consecutive frames, which otherwise makes SGD steps highly correlated, and it reshapes the training distribution toward something closer to i.i.d. over recent history, exactly the setting SGD is designed for. It is only legal because the update is off-policy. The target network computes the bootstrap target with a frozen copy θ⁻ of the parameters, refreshed every C steps (or Polyak-averaged continuously): target = r + γ max_a′ Q_θ⁻(s′, a′). Between refreshes the regression target is stationary, so each phase of training is an ordinary supervised problem, and the prediction-chases-its-own-tail loop is opened. Neither fix is a convergence proof; both are empirically what let the same architecture and hyperparameters play 49 Atari games from pixels.

Implementation, twice

First tabular Q-learning on a small inline gridworld: a 4×4 grid, start in the top-left corner, terminal goal in the bottom-right, reward −1 per step so the greedy policy is the shortest path. The environment is five lines, which is the point: the algorithm is the update, and everything else is bookkeeping. The JAX version keeps the table as an immutable array updated functionally with .at[].add() inside a jitted update, with exploration driven by NumPy outside.

import torch

GOAL = 15                              # 4x4 grid, state = 4*row + col
def grid_step(s, a):
    """Deterministic gridworld. Actions 0..3 = up, down, left, right."""
    r, c = divmod(s, 4)
    if a == 0: r = max(r - 1, 0)
    if a == 1: r = min(r + 1, 3)
    if a == 2: c = max(c - 1, 0)
    if a == 3: c = min(c + 1, 3)
    s2 = 4 * r + c
    return s2, -1.0, s2 == GOAL        # -1 per step: greedy = shortest path

Q = torch.zeros(16, 4)
alpha, gamma, eps = 0.1, 0.95, 0.1
g = torch.Generator().manual_seed(0)

for episode in range(500):
    s, done = 0, False
    while not done:
        if torch.rand((), generator=g) < eps:              # explore
            a = int(torch.randint(4, (1,), generator=g))
        else:                                              # exploit
            a = int(Q[s].argmax())
        s2, r, done = grid_step(s, a)
        target = r if done else r + gamma * float(Q[s2].max())
        Q[s, a] += alpha * (target - float(Q[s, a]))       # the whole algorithm
        s = s2

# greedy-policy state values; start corner ~ -5.29
# (6 steps to goal, each -1, discounted by gamma)
print(Q.max(dim=1).values.view(4, 4))

import numpy as np
import jax
import jax.numpy as jnp

GOAL = 15                              # 4x4 grid, state = 4*row + col
def grid_step(s, a):
    """Deterministic gridworld. Actions 0..3 = up, down, left, right."""
    r, c = divmod(s, 4)
    if a == 0: r = max(r - 1, 0)
    if a == 1: r = min(r + 1, 3)
    if a == 2: c = max(c - 1, 0)
    if a == 3: c = min(c + 1, 3)
    s2 = 4 * r + c
    return s2, -1.0, s2 == GOAL        # -1 per step: greedy = shortest path

alpha, gamma, eps = 0.1, 0.95, 0.1

@jax.jit
def q_update(Q, s, a, r, s2, done):
    """One Bellman backup; (1 - done) drops the future term at terminals."""
    target = r + (1.0 - done) * gamma * Q[s2].max()
    return Q.at[s, a].add(alpha * (target - Q[s, a]))

Q = jnp.zeros((16, 4))
rng = np.random.default_rng(0)

for episode in range(500):
    s, done = 0, False
    while not done:
        if rng.random() < eps:                             # explore
            a = int(rng.integers(4))
        else:                                              # exploit
            a = int(jnp.argmax(Q[s]))
        s2, r, terminal = grid_step(s, a)
        Q = q_update(Q, s, a, r, s2, float(terminal))
        s, done = s2, terminal

# greedy-policy state values; start corner ~ -5.29
# (6 steps to goal, each -1, discounted by gamma)
print(Q.max(axis=1).reshape(4, 4))

Then the full DQN on CartPole-v1 via gymnasium, the standard environment API: replay buffer, target network, Huber loss, and a linear epsilon schedule. Both scripts are complete and print running mean returns. In the JAX version the replay buffer lives in plain NumPy on the host, since a ring buffer is mutation-shaped and only the sampled minibatch needs to be a device array; the gradient step is one jitted pure function. The target parameters enter that function as data, not as the differentiated argument, so no gradient flows through the bootstrap target, which is the JAX equivalent of PyTorch's torch.no_grad() around the target computation.

import random
from collections import deque
import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

env = gym.make("CartPole-v1")
obs_dim = env.observation_space.shape[0]          # 4
n_act = env.action_space.n                        # 2

def qnet():
    return nn.Sequential(nn.Linear(obs_dim, 128), nn.ReLU(),
                         nn.Linear(128, 128), nn.ReLU(),
                         nn.Linear(128, n_act))

random.seed(0); torch.manual_seed(0)
q, q_target = qnet(), qnet()
q_target.load_state_dict(q.state_dict())          # start in sync
opt = torch.optim.Adam(q.parameters(), lr=1e-3)

buffer = deque(maxlen=50_000)
gamma, batch_size = 0.99, 64
eps_hi, eps_lo, eps_steps = 1.0, 0.05, 10_000
warmup, target_every, total_steps = 1_000, 500, 60_000

obs, _ = env.reset(seed=0)
ep_ret, ep_returns = 0.0, []

for t in range(total_steps):
    eps = max(eps_lo, eps_hi - (eps_hi - eps_lo) * t / eps_steps)
    if random.random() < eps:
        a = env.action_space.sample()
    else:
        with torch.no_grad():
            a = int(q(torch.as_tensor(obs, dtype=torch.float32)).argmax())
    next_obs, r, terminated, truncated, _ = env.step(a)
    # Store `terminated` only: a time-limit truncation is not a real
    # terminal state, so we must still bootstrap through it.
    buffer.append((obs, a, r, next_obs, float(terminated)))
    ep_ret += r
    obs = next_obs
    if terminated or truncated:
        ep_returns.append(ep_ret)
        obs, _ = env.reset()
        ep_ret = 0.0

    if len(buffer) >= warmup:
        batch = random.sample(buffer, batch_size)  # uniform: decorrelates
        o, a_b, r_b, o2, d = map(np.array, zip(*batch))
        o = torch.as_tensor(o, dtype=torch.float32)
        o2 = torch.as_tensor(o2, dtype=torch.float32)
        a_b, r_b, d = (torch.as_tensor(a_b),
                       torch.as_tensor(r_b, dtype=torch.float32),
                       torch.as_tensor(d, dtype=torch.float32))

        q_sa = q(o).gather(1, a_b[:, None]).squeeze(1)
        with torch.no_grad():                      # frozen bootstrap target
            tgt = r_b + gamma * (1 - d) * q_target(o2).max(dim=1).values
        loss = F.smooth_l1_loss(q_sa, tgt)         # Huber: clips TD outliers
        opt.zero_grad(); loss.backward(); opt.step()

    if t % target_every == 0:
        q_target.load_state_dict(q.state_dict())   # refresh frozen copy
    if t % 5_000 == 0 and ep_returns:
        print(f"step {t:6d}  eps {eps:.2f}  "
              f"mean return (last 20) {np.mean(ep_returns[-20:]):6.1f}")

import gymnasium as gym
import numpy as np
import jax
import jax.numpy as jnp
import optax

def init_mlp(key, sizes):
    params = []
    for m, n in zip(sizes[:-1], sizes[1:]):
        key, sub = jax.random.split(key)
        params.append({"w": jax.random.normal(sub, (m, n)) * jnp.sqrt(2.0 / m),
                       "b": jnp.zeros(n)})
    return params

def mlp_apply(params, x):
    for layer in params[:-1]:
        x = jax.nn.relu(x @ layer["w"] + layer["b"])
    return x @ params[-1]["w"] + params[-1]["b"]

class Replay:
    """Flat NumPy ring buffer. The host owns the mutable data;
    only sampled minibatches ever become device arrays."""
    def __init__(self, cap, obs_dim):
        self.obs = np.zeros((cap, obs_dim), np.float32)
        self.act = np.zeros(cap, np.int32)
        self.rew = np.zeros(cap, np.float32)
        self.nxt = np.zeros((cap, obs_dim), np.float32)
        self.done = np.zeros(cap, np.float32)
        self.cap, self.ptr, self.full = cap, 0, False
    def add(self, o, a, r, o2, d):
        i = self.ptr
        self.obs[i], self.act[i], self.rew[i] = o, a, r
        self.nxt[i], self.done[i] = o2, d
        self.ptr = (i + 1) % self.cap
        self.full = self.full or self.ptr == 0
    def sample(self, rng, n):
        idx = rng.integers(self.cap if self.full else self.ptr, size=n)
        return (self.obs[idx], self.act[idx], self.rew[idx],
                self.nxt[idx], self.done[idx])

gamma = 0.99
opt = optax.adam(1e-3)

@jax.jit
def dqn_update(params, target_params, opt_state, o, a, r, o2, d):
    def loss_fn(p):
        q_sa = jnp.take_along_axis(mlp_apply(p, o), a[:, None], 1)[:, 0]
        # target_params is not the differentiated argument, so the
        # bootstrap target is frozen: no stop_gradient needed.
        tgt = r + gamma * (1.0 - d) * mlp_apply(target_params, o2).max(axis=1)
        return jnp.mean(optax.huber_loss(q_sa, tgt))
    loss, grads = jax.value_and_grad(loss_fn)(params)
    updates, opt_state = opt.update(grads, opt_state)
    return optax.apply_updates(params, updates), opt_state, loss

@jax.jit
def greedy_action(params, obs):
    return jnp.argmax(mlp_apply(params, obs))

env, rng = gym.make("CartPole-v1"), np.random.default_rng(0)
params = init_mlp(jax.random.PRNGKey(0), [4, 128, 128, 2])
target_params = params                             # start in sync
opt_state = opt.init(params)
buf = Replay(50_000, 4)
eps_hi, eps_lo, eps_steps = 1.0, 0.05, 10_000
warmup, target_every, total_steps = 1_000, 500, 60_000

obs, _ = env.reset(seed=0)
ep_ret, ep_returns = 0.0, []

for t in range(total_steps):
    eps = max(eps_lo, eps_hi - (eps_hi - eps_lo) * t / eps_steps)
    if rng.random() < eps:
        a = int(rng.integers(2))
    else:
        a = int(greedy_action(params, jnp.asarray(obs)))
    next_obs, r, terminated, truncated, _ = env.step(a)
    # Store `terminated` only: bootstrap through time-limit truncation.
    buf.add(obs, a, r, next_obs, float(terminated))
    ep_ret += r
    obs = next_obs
    if terminated or truncated:
        ep_returns.append(ep_ret)
        obs, _ = env.reset()
        ep_ret = 0.0

    if (buf.full or buf.ptr >= warmup):
        o, a_b, r_b, o2, d = buf.sample(rng, 64)
        params, opt_state, loss = dqn_update(
            params, target_params, opt_state,
            jnp.asarray(o), jnp.asarray(a_b), jnp.asarray(r_b),
            jnp.asarray(o2), jnp.asarray(d))

    if t % target_every == 0:
        target_params = params                     # refresh frozen copy
    if t % 5_000 == 0 and ep_returns:
        print(f"step {t:6d}  eps {eps:.2f}  "
              f"mean return (last 20) {np.mean(ep_returns[-20:]):6.1f}")

Using it on a real shape of problem

The gridworld converges in well under 500 episodes; the printed value table should read approximately −5.29 in the start corner (the discounted cost of the 6-step shortest path, −(1 − 0.95⁶)/0.05), climbing to exactly −1 adjacent to the goal, a smoothed version of the −(steps remaining) shape that a −1-per-step reward implies. For the CartPole DQN, expect returns near 20 while epsilon is high, visible learning once the buffer warms up and epsilon decays (roughly steps 5,000 to 20,000), and runs that reach the 400 to 500 ceiling somewhere between 30,000 and 60,000 steps. All of it is seed and machine dependent, and DQN curves are notoriously spiky: a policy can hit 500, then briefly collapse as the buffer's distribution shifts, then recover. Evaluate greedily, with exploration off:

env = gym.make("CartPole-v1", render_mode="human")
obs, _ = env.reset(seed=0)
done, total = False, 0.0
while not done:
    with torch.no_grad():
        a = int(q(torch.as_tensor(obs, dtype=torch.float32)).argmax())
    obs, r, terminated, truncated, _ = env.step(a)
    total += r
    done = terminated or truncated
print(total)                          # ~500 once trained

env = gym.make("CartPole-v1", render_mode="human")
obs, _ = env.reset(seed=0)
done, total = False, 0.0
while not done:
    a = int(greedy_action(params, jnp.asarray(obs)))
    obs, r, terminated, truncated, _ = env.step(a)
    total += r
    done = terminated or truncated
print(total)                          # ~500 once trained

The diagnostic to watch besides return is the mean Q-value on a fixed batch of held-out states. It should rise smoothly toward the true achievable return; a Q estimate exploding past any return the environment can actually pay is the classic signature of the deadly triad reasserting itself, and the first knobs to reach for are a slower target refresh and a lower learning rate.

Applications

DQN's founding demonstration remains its calling card: one architecture, one set of hyperparameters, 49 Atari 2600 games from raw pixels, with the 2015 Nature paper reporting human-level or better play on a majority of them. That result seeded an entire lineage of value-based improvements: Double DQN against overestimation, dueling networks separating state value from action advantage, prioritized replay, and distributional Q-learning, bundled by DeepMind's Rainbow paper into the strongest classic value-based Atari agent. Where value-based methods still win today is the regime that plays to their strengths: discrete, moderate-cardinality action spaces where off-policy replay makes every logged transition reusable. That covers recommendation and slate-style decision problems framed as discrete choices, traffic-signal and resource scheduling, game AI beyond Atari, and notably offline RL, where learning must proceed entirely from a fixed log of someone else's behavior; the leading offline algorithms (CQL and its relatives) are Q-learning with a conservatism penalty. Q-functions also live inside methods nominally from other families: DDPG, TD3, and SAC train Q networks with replay buffers and target networks, DQN's machinery verbatim, and use a learned actor as a differentiable argmax for continuous actions; that hybrid is picked up on the actor-critic page. When the problem instead demands stochastic policies, continuous high-dimensional actions, or optimization of a non-Markov score, the policy gradient route and its PPO refinement take over.

Against the real libraries

CleanRL's dqn.py is the single most instructive next read: one self-contained file, structurally almost identical to the scripts above, but carrying the tuned defaults that matter (buffer sizes, a train frequency that does not update every step, evaluation protocol) and, in its sibling dqn_atari.py, the full Atari preprocessing stack: frame stacking, reward clipping, and the convolutional torso from the Nature paper. CleanRL publishes benchmark curves for every script, which turns "is my implementation right" from a feeling into a comparison. There is also a JAX variant of its DQN, useful as a cross-check for the JAX version here.

stable-baselines3 ships DQN as a supported algorithm behind a three-line API, and what it adds over a from-scratch script is exactly what production needs and pages like this omit: correct bootstrapping across gymnasium's terminated/truncated boundary handled for you, vectorized environments, polyak or hard target updates behind a flag, checkpointing, and a long issue-tracker history of edge cases already fixed. Its contrib package adds QR-DQN from the distributional line. Gymnasium is the environment substrate throughout, and its Atari environments (via ale-py) are the standard benchmark plumbing. On the JAX side, optax provides the optimizer and Huber loss used above.

The from-scratch version is enough for small state spaces, coursework, and algorithm research where you intend to modify the update itself. Verification is two-tier. The tabular algorithm admits an exact check: on the 4×4 gridworld, run value iteration to convergence with the known model and assert your learned Q matches Q* within a small tolerance (with α decayed, agreement to two or three decimals is achievable; with fixed α expect agreement to within about α). For DQN, verification is statistical: run CleanRL's dqn.py and your script on CartPole-v1 for the same step budget across several seeds and compare learning-curve envelopes; both should reliably exceed 475 mean return, and a from-scratch version that never does has a bug, not bad luck.

Traps and misconceptions

Treating time-limit truncation as death. CartPole ends at 500 steps by a timer, not by falling. If that cut is stored as a terminal (dropping the bootstrap term), the agent learns that the world ends at step 500 and the values near the limit are systematically wrong. Gymnasium splits terminated from truncated precisely so replay buffers can store only true termination, and both scripts above do; folding them together with done = terminated or truncated inside the buffer is among the most common DQN bugs in the wild.

"Off-policy means the exploration policy does not matter." The convergence theorem needs every state-action pair visited infinitely often; off-policy freedom is freedom in which exploratory policy you follow, not freedom from exploration. Decay epsilon too fast and whole regions of the state space are never corrected, leaving stale values the argmax will happily exploit into a wall.

Ignoring maximization bias. The max of noisy estimates is biased upward: with many actions whose true values are equal, max_a′ Q(s′, a′) picks whichever estimate is currently most overestimated. This compounds through bootstrapping and inflates Q-values systematically. Double DQN decouples selection from evaluation, choosing a′ with the online network but scoring it with the target network, and costs roughly one extra forward pass; it is a near-universal upgrade.

Expecting the replay buffer and target network to guarantee convergence. They are stabilizers, not proofs; the deadly triad is managed, not solved. Divergence still happens, especially with aggressive learning rates, tiny buffers, or very fast target refreshes, and the visible symptom is Q estimates growing past any return the environment can pay. Treat those three hyperparameters as the stability budget.

Reaching for DQN with continuous actions. The update needs max_a′ Q(s′, a′), an argmax over actions, which is trivial over 4 discrete choices and an optimization problem of its own over a continuous torque vector. Continuous control either discretizes coarsely or moves to the actor-as-argmax methods (DDPG, TD3, SAC) covered with actor-critic, which keep DQN's replay and target machinery but learn the maximizer.