class_docstr: Twin Delayed Deep Deterministic Policy Gradient algorithm trainer
  as described in https://arxiv.org/pdf/1802.09477
class_docstr_tok: ["Twin","Delayed","Deep","Deterministic","Policy","Gradient","algorithm","trainer","as","described","in","https",":","\/\/arxiv.org\/pdf\/1802.09477"]
docstr: IMPORTANT: the input action here is assumed to be preprocessed to match the
  range of the output of the actor.
docstr_tok: ["IMPORTANT",":","the","input","action","here","is","assumed","to","be","preprocessed","to","match","the","range","of","the","output","of","the","actor","."]
code: def train_step_gen(
      self, training_batch: rlt.PolicyNetworkInput, batch_idx: int
  ):
      """
      IMPORTANT: the input action here is assumed to be preprocessed to match the
      range of the output of the actor.
      """
      assert isinstance(training_batch, rlt.PolicyNetworkInput)
  
      state = training_batch.state
      action = training_batch.action
      next_state = training_batch.next_state
      reward = training_batch.reward
      not_terminal = training_batch.not_terminal
  
      # Generate target = r + y * min (Q1(s',pi(s')), Q2(s',pi(s')))
      with torch.no_grad():
          next_actor = self.actor_network_target(next_state).action
          noise = torch.randn_like(next_actor) * self.noise_variance
          next_actor = (
              next_actor + noise.clamp(*self.noise_clip_range)
          ).clamp(*CONTINUOUS_TRAINING_ACTION_RANGE)
          next_state_actor = (next_state, rlt.FeatureData(next_actor))
          next_q_value = self.q1_network_target(*next_state_actor)
  
          if self.q2_network is not None:
              next_q_value = torch.min(
                  next_q_value,
                  self.q2_network_target(*next_state_actor),
              )
  
          target_q_value = (
              reward + self.gamma * next_q_value * not_terminal.float()
          )
  
      # Optimize Q1 and Q2
      q1_value = self.q1_network(state, action)
      q1_loss = F.mse_loss(q1_value, target_q_value)
      if batch_idx % self.trainer.log_every_n_steps == 0:
          self.reporter.log(
              q1_loss=q1_loss,
              q1_value=q1_value,
              next_q_value=next_q_value,
              target_q_value=target_q_value,
          )
      self.log(
          "td_loss",
          q1_loss,
          prog_bar=True,
          batch_size=training_batch.batch_size(),
      )
      yield q1_loss
  
      if self.q2_network:
          q2_value = self.q2_network(state, action)
          q2_loss = F.mse_loss(q2_value, target_q_value)
          if batch_idx % self.trainer.log_every_n_steps == 0:
              self.reporter.log(
                  q2_loss=q2_loss,
                  q2_value=q2_value,
              )
          yield q2_loss
  
      # Only update actor and target networks after a fixed number of Q updates
      if batch_idx % self.delayed_policy_update == 0:
          actor_action = self.actor_network(state).action
          actor_q1_value = self.q1_network(
              state, rlt.FeatureData(actor_action)
          )
          actor_loss = -(actor_q1_value.mean())
          if batch_idx % self.trainer.log_every_n_steps == 0:
              self.reporter.log(
                  actor_loss=actor_loss,
                  actor_q1_value=actor_q1_value,
              )
          yield actor_loss
  
          # Use the soft update rule to update the target networks
          result = self.soft_update_result()
          yield result
  
      else:
          # Yielding None prevents the actor and target networks from updating
          yield None
          yield None
  
file_path: ReAgent/reagent/training/td3_trainer.py
imports: import copy
  import logging
  import reagent.core.types
  import torch
  import torch.nn.functional
  from reagent.core.configuration import resolve_defaults
  from reagent.core.dataclasses import field
  from reagent.core.parameters import CONTINUOUS_TRAINING_ACTION_RANGE, RLParameters
  from reagent.optimizer import Optimizer__Union, SoftUpdate
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.rl_trainer_pytorch import RLTrainerMixin
label_desc: Use image node_id 3 for calling the TD3Trainer obj's underlying member method code with example usage: obj.train_step_gen(training_batch, batch_idx) without return types