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}
code: def __init__(self, gamma: float):
      super().__init__()
      self.gamma = gamma
      self.huber_loss = nn.SmoothL1Loss(reduction="none")
  
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 1 to create a new QFuncLoss object from inherited base classes: Module with example: obj = QFuncLoss(gamma)