class_docstr: ## Trainer
class_docstr_tok: ["#","#","Trainer"]
docstr: ### Train the model
docstr_tok: ["#","#","#","Train","the","model"]
code: def train(self, beta: float):
      """
      ### Train the model
      """
      for _ in range(self.train_epochs):
          # Sample from priority replay buffer
          samples = self.replay_buffer.sample(
              self.mini_batch_size, beta
          )
          # Get the predicted Q-value
          q_value = self.model(obs_to_torch(samples["obs"]))
  
          # Get the Q-values of the next state for [Double Q-learning](index.html).
          # Gradients shouldn't propagate for these
          with torch.no_grad():
              # Get $\textcolor{cyan}Q(s';\textcolor{cyan}{\theta_i})$
              double_q_value = self.model(
                  obs_to_torch(samples["next_obs"])
              )
              # Get $\textcolor{orange}Q(s';\textcolor{orange}{\theta_i^{-}})$
              target_q_value = self.target_model(
                  obs_to_torch(samples["next_obs"])
              )
  
          # Compute Temporal Difference (TD) errors, $\delta$, and the loss, $\mathcal{L}(\theta)$.
          td_errors, loss = self.loss_func(
              q_value,
              q_value.new_tensor(samples["action"]),
              double_q_value,
              target_q_value,
              q_value.new_tensor(samples["done"]),
              q_value.new_tensor(samples["reward"]),
              q_value.new_tensor(samples["weights"]),
          )
  
          # Calculate priorities for replay buffer $p_i = |\delta_i| + \epsilon$
          new_priorities = np.abs(td_errors.cpu().numpy()) + 1e-6
          # Update replay buffer priorities
          self.replay_buffer.update_priorities(
              samples["indexes"], new_priorities
          )
  
          # Set learning rate
          for pg in self.optimizer.param_groups:
              pg["lr"] = self.learning_rate()
          # Zero out the previously calculated gradients
          self.optimizer.zero_grad()
          # Calculate gradients
          loss.backward()
          # Clip gradients
          torch.nn.utils.clip_grad_norm_(
              self.model.parameters(), max_norm=0.5
          )
          # Update parameters based on gradients
          self.optimizer.step()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/rl/dqn/experiment.py
imports: import numpy
  import torch
  from labml import tracker, experiment, logger, monit
  from labml.internal.configs.dynamic_hyperparam import FloatDynamicHyperParam
  from labml_helpers.schedule import Piecewise
  from labml_nn.rl.dqn import QFuncLoss
  from labml_nn.rl.dqn.model import Model
  from labml_nn.rl.dqn.replay_buffer import ReplayBuffer
  from labml_nn.rl.game import Worker
label_desc: Use image node_id 4 for calling the Trainer obj's underlying member method code with example usage: obj.train(beta) without return types