Skip to content

← JAX Autodiff step 13 of 25

Easy Primitives

stop_gradient: Target Network

Jump to the task

Compare a prediction with a fixed target

A Q-value estimates the return from taking an action. In a DQN-style update, the online model's value for the action taken is compared with a target built from the reward and estimated future return:

target = reward + gamma * max(q_target)
error = target - q_online[action]
loss = error**2

Here q_online contains values for the current state, and q_target contains target-model values for the next state. gamma discounts the future contribution. This exercise models a transition where that contribution is present; it has no terminal-state mask. A real training loop must handle terminal transitions separately.

Stop the bootstrap derivative

The bootstrap term uses an estimated next-state value. Treating it as fixed during this update gives the desired online-model learning signal. lax.stop_gradient keeps its value but blocks derivatives through it:

import jax
import jax.numpy as jnp
from jax import lax

def loss(online, target_values):
    target = 1.0 + 0.5 * lax.stop_gradient(jnp.max(target_values))
    return (target - online[0]) ** 2

online = jnp.array([1.0, 2.0])
target_values = jnp.array([3.0, 4.0])
grads = jax.grad(loss, argnums=(0, 1))(online, target_values)
# loss = (3 - 1)**2 = 4
# grads[0] = [-4.0, 0.0]
# grads[1] = [0.0, 0.0]

Only the chosen online action has a nonzero derivative. Without the stop, the loss value would still be 4, but the target-values gradient would be [0, 2]. The stop affects both forward- and reverse-mode automatic differentiation.

If online and target parameters are separate arguments and you only differentiate the online argument, that alone excludes the target argument from the returned gradient. It is therefore too broad to say omitting stop_gradient always updates a separate target network. The optimizer and target-update rules also determine which parameters change.

Problem

Implement td_loss(q_online, q_target, action_taken, reward, gamma). Compute the squared error in the formula above, placing lax.stop_gradient around jnp.max(q_target) only.

Both Q-value inputs are nonempty one-dimensional real floating-point JAX arrays with shape (n_actions,). Reward and gamma are real floating-point scalars. The action is a valid zero-based integer index supplied as a float, such as 1.0; convert it with jnp.int32(action_taken) before indexing. This float representation is part of the exercise interface, not a general JAX indexing rule.

Return the scalar squared error. Do not return the signed error, take a maximum over q_online, or stop the online prediction's derivative. Reward and gamma remain outside the stop; if differentiated, they can still have nonzero derivatives.

For q_online=[3, 5, 7], q_target=[2, 4, 6], action 2, reward 2, and gamma 0.5, the target is 5 and the prediction is 7. Return 4. The output-only tests check that value; the joint-gradient example above demonstrates the stopping behavior that values alone cannot test.