class_docstr: Trainer for Seq2Reward
class_docstr_tok: ["Trainer","for","Seq2Reward"]
docstr: Compute cross-entropy losses of step predictions
  
  :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:
      step_entropy_loss on step prediction
docstr_tok: ["Compute","cross-entropy","losses","of","step","predictions",":","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",":","step_entropy_loss","on","step","prediction"]
code: def get_step_entropy_loss(
      self, training_batch: rlt.MemoryNetworkInput
  ):
      """
      Compute cross-entropy losses of step predictions
  
      :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:
          step_entropy_loss on step prediction
      """
      # pyre-fixme[16]: Optional type has no attribute `flatten`.
      valid_step = training_batch.valid_step.flatten()
  
      first_step_state = training_batch.state.float_features[0]
      valid_step_output = self.step_predict_network(first_step_state)
  
      # step loss's target is zero-based indexed, so subtract 1 from valid_step
      return self.step_loss(valid_step_output, valid_step - 1)
  
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 6 for calling the Seq2RewardTrainer obj's underlying member method code with example usage: obj.get_step_entropy_loss(training_batch) and returns: self