class_docstr: Trainer for Seq2Reward
class_docstr_tok: ["Trainer","for","Seq2Reward"]
docstr: Compute losses:
      MSE(predicted_acc_reward, target_acc_reward)
  
  :param training_batch:
      training_batch has these fields:
      - state: (SEQ_LEN, BATCH_SIZE, STATE_DIM) torch tensor
      - action: (SEQ_LEN, BATCH_SIZE, ACTION_DIM) torch tensor
      - reward: (SEQ_LEN, BATCH_SIZE) torch tensor
  
  :returns:
      mse loss on reward
docstr_tok: ["Compute","losses",":","MSE","(","predicted_acc_reward",",","target_acc_reward",")",":","param","training_batch",":","training_batch","has","these","fields",":","-","state",":","(","SEQ_LEN",",","BATCH_SIZE",",","STATE_DIM",")","torch","tensor","-","action",":","(","SEQ_LEN",",","BATCH_SIZE",",","ACTION_DIM",")","torch","tensor","-","reward",":","(","SEQ_LEN",",","BATCH_SIZE",")","torch","tensor",":","returns",":","mse","loss","on","reward"]
code: def get_mse_loss(self, training_batch: rlt.MemoryNetworkInput):
      """
      Compute losses:
          MSE(predicted_acc_reward, target_acc_reward)
  
      :param training_batch:
          training_batch has these fields:
          - state: (SEQ_LEN, BATCH_SIZE, STATE_DIM) torch tensor
          - action: (SEQ_LEN, BATCH_SIZE, ACTION_DIM) torch tensor
          - reward: (SEQ_LEN, BATCH_SIZE) torch tensor
  
      :returns:
          mse loss on reward
      """
      # pyre-fixme[16]: Optional type has no attribute `flatten`.
      valid_step = training_batch.valid_step.flatten()
  
      seq2reward_output = self.seq2reward_network(
          training_batch.state,
          training_batch.action,
          valid_step,
      )
      predicted_acc_reward = seq2reward_output.acc_reward
  
      seq_len, batch_size = training_batch.reward.size()
      gamma = self.params.gamma
      gamma_mask = (
          torch.Tensor(
              [
                  [gamma**i for i in range(seq_len)]
                  for _ in range(batch_size)
              ]
          )
          .transpose(0, 1)
          .to(training_batch.reward.device)
      )
  
      target_acc_rewards = torch.cumsum(
          training_batch.reward * gamma_mask, dim=0
      )
      target_acc_reward = target_acc_rewards[
          valid_step - 1, torch.arange(batch_size)
      ].unsqueeze(1)
  
      # make sure the prediction and target tensors have the same size
      # the size should both be (BATCH_SIZE, 1) in this case.
      assert (
          predicted_acc_reward.size() == target_acc_reward.size()
      ), f"{predicted_acc_reward.size()}!={target_acc_reward.size()}"
      return self.mse_loss(predicted_acc_reward, target_acc_reward)
  
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 5 for calling the Seq2RewardTrainer obj's underlying member method code with example usage: obj.get_mse_loss(training_batch) and returns: self