class_docstr: ## Train the model
  
  We want to find optimal action-value function.
  
  egin{align}
      Q^*(s,a) &= \max_\pi \mathbb{E} \Big[
          r_t + \gamma r_{t + 1} + \gamma^2 r_{t + 2} + ... | s_t = s, a_t = a, \pi
      \Big]
  \
      Q^*(s,a) &= \mathop{\mathbb{E}}_{s' \sim \large{arepsilon}} \Big[
          r + \gamma \max_{a'} Q^* (s', a') | s, a
      \Big]
  \end{align}
  
  ### Target network 🎯
  In order to improve stability we use experience replay that randomly sample
  from previous experience $U(D)$. We also use a Q network
  with a separate set of parameters $ extcolor{orange}{       heta_i^{-}}$ to calculate the target.
  $   extcolor{orange}{       heta_i^{-}}$ is updated periodically.
  This is according to paper
  [Human Level Control Through Deep Reinforcement Learning](https://deepmind.com/research/dqn/).
  
  So the loss function is,
  $$
  \mathcal{L}_i(      heta_i) = \mathop{\mathbb{E}}_{(s,a,r,s') \sim U(D)}
  igg[
      \Big(
          r + \gamma \max_{a'} Q(s', a';      extcolor{orange}{       heta_i^{-}}) - Q(s,a;   heta_i)
      \Big) ^ 2
  igg]
  $$
  
  ### Double $Q$-Learning
  The max operator in the above calculation uses same network for both
  selecting the best action and for evaluating the value.
  That is,
  $$
  \max_{a'} Q(s', a';         heta) =         extcolor{cyan}{Q}
  \Big(
      s', \mathop{\operatorname{argmax}}_{a'}
              extcolor{cyan}{Q}(s', a';       extcolor{cyan}{ heta});         extcolor{cyan}{ heta}
  \Big)
  $$
  We use [double Q-learning](https://arxiv.org/abs/1509.06461), where
  the $\operatorname{argmax}$ is taken from $ extcolor{cyan}{ heta_i}$ and
  the value is taken from $   extcolor{orange}{       heta_i^{-}}$.
  
  And the loss function becomes,
  
  egin{align}
      \mathcal{L}_i(  heta_i) = \mathop{\mathbb{E}}_{(s,a,r,s') \sim U(D)}
      \Bigg[
          igg(
              &r + \gamma     extcolor{orange}{Q}
              \Big(
                  s',
                  \mathop{\operatorname{argmax}}_{a'}
                              extcolor{cyan}{Q}(s', a';       extcolor{cyan}{ heta_i});       extcolor{orange}{       heta_i^{-}}
              \Big)
              \
              - &Q(s,a;       heta_i)
          igg) ^ 2
      \Bigg]
  \end{align}
docstr: * `q` - $Q(s;   heta_i)$
  * `action` - $a$
  * `double_q` - $        extcolor{cyan}Q(s';     extcolor{cyan}{ heta_i})$
  * `target_q` - $        extcolor{orange}Q(s';   extcolor{orange}{       heta_i^{-}})$
  * `done` - whether the game ended after taking the action
  * `reward` - $r$
  * `weights` - weights of the samples from prioritized experienced replay
code: def forward(
      self,
      q: torch.Tensor,
      action: torch.Tensor,
      double_q: torch.Tensor,
      target_q: torch.Tensor,
      done: torch.Tensor,
      reward: torch.Tensor,
      weights: torch.Tensor,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      * `q` - $Q(s;\theta_i)$
      * `action` - $a$
      * `double_q` - $\textcolor{cyan}Q(s';\textcolor{cyan}{\theta_i})$
      * `target_q` - $\textcolor{orange}Q(s';\textcolor{orange}{\theta_i^{-}})$
      * `done` - whether the game ended after taking the action
      * `reward` - $r$
      * `weights` - weights of the samples from prioritized experienced replay
      """
  
      # $Q(s,a;\theta_i)$
      q_sampled_action = q.gather(
          -1, action.to(torch.long).unsqueeze(-1)
      ).squeeze(-1)
      tracker.add("q_sampled_action", q_sampled_action)
  
      # Gradients shouldn't propagate gradients
      # $$r + \gamma \textcolor{orange}{Q}
      #                 \Big(s',
      #                     \mathop{\operatorname{argmax}}_{a'}
      #                         \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i}); \textcolor{orange}{\theta_i^{-}}
      #                 \Big)$$
      with torch.no_grad():
          # Get the best action at state $s'$
          # $$\mathop{\operatorname{argmax}}_{a'}
          #  \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i})$$
          best_next_action = torch.argmax(double_q, -1)
          # Get the q value from the target network for the best action at state $s'$
          # $$\textcolor{orange}{Q}
          # \Big(s',\mathop{\operatorname{argmax}}_{a'}
          # \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i}); \textcolor{orange}{\theta_i^{-}}
          # \Big)$$
          best_next_q_value = target_q.gather(
              -1, best_next_action.unsqueeze(-1)
          ).squeeze(-1)
  
          # Calculate the desired Q value.
          # We multiply by `(1 - done)` to zero out
          # the next state Q values if the game ended.
          #
          # $$r + \gamma \textcolor{orange}{Q}
          #                 \Big(s',
          #                     \mathop{\operatorname{argmax}}_{a'}
          #                         \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i}); \textcolor{orange}{\theta_i^{-}}
          #                 \Big)$$
          q_update = reward + self.gamma * best_next_q_value * (
              1 - done
          )
          tracker.add("q_update", q_update)
  
          # Temporal difference error $\delta$ is used to weigh samples in replay buffer
          td_error = q_sampled_action - q_update
          tracker.add("td_error", td_error)
  
      # We take [Huber loss](https://en.wikipedia.org/wiki/Huber_loss) instead of
      # mean squared error loss because it is less sensitive to outliers
      losses = self.huber_loss(q_sampled_action, q_update)
      # Get weighted means
      loss = torch.mean(weights * losses)
      tracker.add("loss", loss)
  
      return td_error, loss
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/rl/dqn/__init__.py
imports: from typing import Tuple
  import torch
  from torch import nn
  from labml import tracker
  from labml_helpers.module import Module
  from labml_nn.rl.dqn.replay_buffer import ReplayBuffer
label_desc: Use image node_id 2 for calling the QFuncLoss obj's underlying member method code with example usage: obj.forward(q, action, double_q, target_q, done, reward, weights) and returns: td_error,  loss