19. Q-Learning TD Target With a Frozen Target Network

Problem

Vanilla Q-learning is unstable when the same network both produces the current Q-estimate and bootstraps its own target: the target moves every step you take toward it, so training chases a shifting goalpost. DQN fixed this by keeping a second, frozen “target” copy of the network and computing the Bellman target from that network's Q-values instead of the online network's.

Given the online network's Q-values (still attached to the graph) and the target network's Q-values at the next state (already detached, as they'd be after a with torch.no_grad() forward pass), compute the TD target and the resulting loss:
y_t = r_t + gamma * (1 - done_t) * max_a' Q_target(s_{t+1}, a')
The (1 - done_t) mask is the part that's easy to forget: if the episode ended at this transition, there is no next state to bootstrap from, so the target collapses to just the immediate reward.

Source: src/19_q_learning_td_target.py

def q_learning_td_loss(
    q_values: torch.Tensor, actions: torch.Tensor, rewards: torch.Tensor,
    dones: torch.Tensor, next_q_target: torch.Tensor, gamma: float,
) -> torch.Tensor: ...

>>> q = torch.randn(2, 3, requires_grad=True)
>>> actions = torch.tensor([0, 2])
>>> rewards = torch.tensor([1.0, 0.0])
>>> dones = torch.tensor([0.0, 1.0])
>>> next_q = torch.randn(2, 3)
>>> q_learning_td_loss(q, actions, rewards, dones, next_q, gamma=0.99).shape
torch.Size([])

Step-by-Step Approach

1. Select the online Q-value actually taken at each timestep: gather along the action dimension using actions, going from (batch, num_actions) down to (batch,).
2. Wrap the target computation in torch.no_grad() — the target must never receive gradient from this loss, since it represents a frozen network's output by construction.
3. Take the max over actions of the target network's next-state Q-values: next_q_target.max(dim=1).values, giving the greedy bootstrap value per transition.
4. Mask the bootstrap term by (1 - dones): a terminal transition has no valid next state, so its target must reduce to the reward alone.
5. Assemble the Bellman target: rewards + gamma * (1 - dones) * max_next_q.
6. Compute the mean-squared-error between the selected online Q-value and the (gradient-free) target — this is the loss a training loop calls .backward() on.

The two-network split only matters at one place in the code: the target computation must be wrapped so gradient never flows through it, while the online Q-value that's compared against it must stay fully attached to the graph — get that asymmetry backwards and the target network provides no stabilization at all.

Reference solution

import torch
import torch.nn.functional as F


def q_learning_td_loss(
    q_values: torch.Tensor, actions: torch.Tensor, rewards: torch.Tensor,
    dones: torch.Tensor, next_q_target: torch.Tensor, gamma: float,
) -> torch.Tensor:
    # gather the online Q-value for the action actually taken: (batch, num_actions) -> (batch,)
    selected_q = q_values.gather(1, actions.unsqueeze(1)).squeeze(1)

    # the target must never receive gradient -- it comes from a frozen
    # network by construction, and even if next_q_target already has
    # requires_grad=False, no_grad here documents that intent explicitly
    with torch.no_grad():
        max_next_q = next_q_target.max(dim=1).values  # (batch,)
        target = rewards + gamma * (1.0 - dones) * max_next_q

    return F.mse_loss(selected_q, target)

Key Functions & Tricks

• Tensor.gather(1, actions.unsqueeze(1)).squeeze(1) — pulls out each row's Q-value for the specific action taken, without a Python-level loop over the batch
• torch.no_grad() — detaches the target computation from autograd, matching how a real frozen target network's forward pass would behave
• Tensor.max(dim=1).values — the greedy max over the target network's next-state Q-values, i.e. the standard (non-Double) DQN bootstrap
• F.mse_loss — the squared TD-error, averaged over the batch, matching DQN's original loss
• (1.0 - dones) — zeroes the bootstrap term exactly at terminal transitions, a one-line but easy-to-omit correctness detail

How to Recognize This Pattern