class_docstr: Trainer for fitting Seq2Reward planning outcomes to a neural network-based policy
class_docstr_tok: ["Trainer","for","fitting","Seq2Reward","planning","outcomes","to","a","neural","network-based","policy"]
code: def train_step_gen(
      self, training_batch: rlt.MemoryNetworkInput, batch_idx: int
  ):
      loss, accuracy = self.get_loss(training_batch)
      detached_loss = loss.cpu().detach().item()
      accuracy = accuracy.item()
      logger.info(
          f"Seq2Reward Compress trainer MSE/Accuracy: {detached_loss}, {accuracy}"
      )
      self.reporter.log(mse_loss=detached_loss, accuracy=accuracy)
      yield loss
  
imports: import logging
  import reagent.core.types
  import torch
  import torch.nn.functional
  from reagent.core.parameters import Seq2RewardTrainerParameters
  from reagent.core.types import FeatureData
  from reagent.models.fully_connected_network import FloatFeatureFullyConnected
  from reagent.models.seq2reward_model import Seq2RewardNetwork
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.utils import gen_permutations
  from reagent.training.world_model.seq2reward_trainer import get_Q
label_desc: Use image node_id 3 for calling the CompressModelTrainer obj's underlying member method code with example usage: obj.train_step_gen(training_batch, batch_idx) without return types