class_docstr: Trainer for MDN-RNN
class_docstr_tok: ["Trainer","for","MDN-RNN"]
code: def train_step_gen(
      self, training_batch: rlt.MemoryNetworkInput, batch_idx: int
  ):
      (
          seq_len,
          batch_size,
          state_dim,
      ) = training_batch.state.float_features.shape
  
      losses = self.get_loss(training_batch, state_dim)
  
      detached_losses = {
          k: loss.cpu().detach().item() for k, loss in losses.items()
      }
      self.reporter.log(
          loss=detached_losses["loss"],
          gmm=detached_losses["gmm"],
          bce=detached_losses["bce"],
          mse=detached_losses["mse"],
      )
      if self.all_batches_processed % 10 == 0:
          logger.info(
              f'loss={detached_losses["loss"]}, gmm={detached_losses["loss"]}, bce={detached_losses["bce"]}, mse={detached_losses["mse"]}'
          )
      loss = losses["loss"]
      # TODO: Must setup (or mock) trainer and a LoggerConnector to call self.log()!
      if self.trainer is not None and self.trainer.logger is not None:
          self.log(
              "td_loss",
              loss,
              prog_bar=True,
              batch_size=training_batch.batch_size(),
          )
      yield loss
  
file_path: ReAgent/reagent/training/world_model/mdnrnn_trainer.py
imports: import logging
  from typing import Optional
  import reagent.core.types
  import torch
  import torch.nn.functional
  from reagent.core.parameters import MDNRNNTrainerParameters
  from reagent.models.mdn_rnn import gmm_loss
  from reagent.models.world_model import MemoryNetwork
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
label_desc: Use image node_id 3 for calling the MDNRNNTrainer obj's underlying member method code with example usage: obj.train_step_gen(training_batch, batch_idx) without return types