Loss Function for DQN Training

Training a Deep Q-Network (DQN) necessitates a precise method for adjusting the main Q-network's parameters. This training takes place within the DQN architecture, which includes mechanisms such as Experience Replay and Target Networks. The central task is to refine the network's parameters, represented by $\theta$ , so that its output $Q(s, a; \theta)$ converges toward the true optimal action-value function $Q^*(s, a)$ .

Recall that in standard Q-learning, we iteratively update our estimate of $Q(s, a)$ based on the Bellman equation. We move our current estimate towards a target value derived from the observed reward and the estimated value of the next state. DQN borrows this core idea but adapts it for function approximation using neural networks.

We need a way to quantify the error between our network's current prediction for a given state-action pair $(s, a)$ and a target value. This error measure is defined by a loss function. The goal of training is to minimize this loss.

For a transition $(s, a, r, s')$ sampled from the experience replay buffer, our Q-network predicts the value $Q(s, a; \theta)$ . What should the target value be? Just like in Q-learning, the target aims to incorporate the immediate reward $r$ and the discounted value of the best action possible from the next state $s'$ . However, to improve stability (as discussed in the previous section on Target Networks), we use the target network with its parameters $\theta^-$ to estimate the value of the next state.

The target value, often denoted as $y$ , is calculated as:

y = r + \gamma \max_{a'} Q(s', a'; \theta^-)

Here, $\gamma$ is the discount factor, and $\max_{a'} Q(s', a'; \theta^-)$ represents the maximum Q-value predicted by the target network for the next state $s'$ across all possible next actions $a'$ . If $s'$ is a terminal state, the target is simply $y = r$ .

With the predicted value $Q(s, a; \theta)$ and the target value $y$ , we can now define the loss. Since we want our prediction to be as close as possible to the target, a natural choice is the Mean Squared Error (MSE) loss, commonly used in regression problems. For a single transition, the squared error is $(y - Q(s, a; \theta))^2$ .

In practice, we compute the loss over a mini-batch of transitions sampled from the experience replay buffer $D$ . The loss function $L(\theta)$ that we aim to minimize is the expectation of this squared error over sampled transitions:

L(\theta) = \mathbb{E}_{(s,a,r,s') \sim D} \left[ (y - Q(s, a; \theta))^2 \right]

Or, more practically, for a mini-batch of $N$ transitions $\{(s_i, a_i, r_i, s'_i)\}_{i=1}^N$ :

L(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \left( (r_i + \gamma \max_{a'} Q(s'_i, a'; \theta^-)) - Q(s_i, a_i; \theta) \right)^2

This loss function measures the average squared difference between the target values (calculated using the target network) and the Q-values predicted by our main network.

Training the Network

The training process involves iteratively performing the following steps:

Sample a mini-batch of transitions $(s, a, r, s')$ from the experience replay buffer $D$ .
For each transition in the mini-batch: a. Calculate the target value $y = r + \gamma \max_{a'} Q(s', a'; \theta^-)$ using the target network. Handle terminal states where $y=r$ . b. Calculate the predicted value $Q(s, a; \theta)$ using the main network.
Compute the MSE loss between the target values and predicted values across the mini-batch using the formula above.
Perform a gradient descent step on the loss $L(\theta)$ with respect to the main network parameters $\theta$ . This updates the weights of the main Q-network to reduce the error.

\theta \leftarrow \theta - \alpha \nabla_\theta L(\theta)

where $\alpha$ is the learning rate.

It's important to note that during the calculation of the gradient $\nabla_\theta L(\theta)$ , the target values $y$ are treated as fixed constants. The gradients only flow through the main network parameters $\theta$ , not the target network parameters $\theta^-$ . This decoupling, achieved by using the separate target network and experience replay, is fundamental to stabilizing the training process for DQNs. Without these mechanisms, the constantly shifting targets and correlated data samples would make convergence difficult.

Periodically, after a set number of training steps, the weights of the target network $\theta^-$ are updated to match the weights of the main network $\theta$ . This ensures the target values gradually adapt to the improving policy represented by the main network, while still providing stability over shorter timescales.

Understanding this loss function and the training procedure is central to implementing and tuning DQN agents. In the next section, we'll walk through a practical implementation.

Was this section helpful?

References

Human-level control through deep reinforcement learning, Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, Martin A. Riedmiller, 2015 Nature, Vol. 518 DOI: 10.1038/nature14236 - This foundational paper introduces the Deep Q-Network (DQN) architecture, including the use of experience replay and target networks, and defines the training objective with the Q-learning loss function.
Reinforcement Learning: An Introduction, Richard S. Sutton and Andrew G. Barto, 2018 (MIT Press) - This widely used textbook covers the fundamental concepts of reinforcement learning, including Q-learning, Bellman equations, and provides a clear explanation of how these principles apply to deep Q-networks.
Spinning Up in Deep RL, Josh Achiam, 2018 - This online resource offers a practical introduction to deep reinforcement learning, providing detailed explanations and examples of the DQN algorithm, its loss function, and training methodology.