class_docstr: Trainer for Seq2Reward
class_docstr_tok: ["Trainer","for","Seq2Reward"]
code: def train_step_gen(
      self, training_batch: rlt.MemoryNetworkInput, batch_idx: int
  ):
      mse_loss = self.get_mse_loss(training_batch)
      detached_mse_loss = mse_loss.cpu().detach().item()
      yield mse_loss
  
      step_entropy_loss = self.get_step_entropy_loss(training_batch)
      detached_step_entropy_loss = (
          step_entropy_loss.cpu().detach().item()
      )
  
      if self.view_q_value:
          state_first_step = training_batch.state.float_features[0]
          q_values = (
              get_Q(
                  self.seq2reward_network,
                  state_first_step,
                  self.all_permut,
              )
              .cpu()
              .mean(0)
              .tolist()
          )
      else:
          q_values = [0] * len(self.params.action_names)
  
      step_probability = (
          get_step_prediction(self.step_predict_network, training_batch)
          .cpu()
          .mean(dim=0)
          .numpy()
      )
      logger.info(
          f"Seq2Reward trainer output: mse_loss={detached_mse_loss}, "
          f"step_entropy_loss={detached_step_entropy_loss}, q_values={q_values}, "
          f"step_probability={step_probability}"
      )
      self.reporter.log(
          mse_loss=detached_mse_loss,
          step_entropy_loss=detached_step_entropy_loss,
          q_values=[q_values],
      )
  
      yield step_entropy_loss
  
imports: import logging
  import reagent.core.types
  import torch
  import torch.nn
  import torch.nn.functional
  from reagent.core.parameters import Seq2RewardTrainerParameters
  from reagent.models.fully_connected_network import FullyConnectedNetwork
  from reagent.models.seq2reward_model import Seq2RewardNetwork
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.utils import gen_permutations
label_desc: Use image node_id 3 for calling the Seq2RewardTrainer obj's underlying member method code with example usage: obj.train_step_gen(training_batch, batch_idx) without return types