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 get_loss(self, batch: rlt.MemoryNetworkInput):
      state_first_step = CompressModelTrainer.extract_state_first_step(
          batch
      )
      # shape: batch_size, num_action
      compress_model_output = self.compress_model_network(
          state_first_step
      )
  
      target = get_Q(
          self.seq2reward_network,
          state_first_step.float_features,
          self.all_permut,
      )
      assert (
          compress_model_output.size() == target.size()
      ), f"{compress_model_output.size()}!={target.size()}"
      mse = F.mse_loss(compress_model_output, target)
  
      with torch.no_grad():
          target_action = torch.max(target, dim=1).indices
          model_action = torch.max(compress_model_output, dim=1).indices
          accuracy = torch.mean((target_action == model_action).float())
  
      return mse, accuracy
  
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 6 for calling the CompressModelTrainer obj's underlying member method code with example usage: obj.get_loss(batch) and returns: mse,  accuracy