class_docstr: Trainer for MDN-RNN
class_docstr_tok: ["Trainer","for","MDN-RNN"]
docstr: Compute losses:
      GMMLoss(next_state, GMMPredicted) / (STATE_DIM + 2)
      + MSE(reward, predicted_reward)
      + BCE(not_terminal, logit_not_terminal)
  
  The STATE_DIM + 2 factor is here to counteract the fact that the GMMLoss scales
      approximately linearly with STATE_DIM, dim of states. All losses
      are averaged both on the batch and the sequence dimensions (the two first
      dimensions).
  
  :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
      - not-terminal: (SEQ_LEN, BATCH_SIZE) torch tensor
      - next_state: (SEQ_LEN, BATCH_SIZE, STATE_DIM) torch tensor
  
  :param state_dim: the dimension of states. If provided, use it to normalize
      gmm loss
  
  :returns: dictionary of losses, containing the gmm, the mse, the bce and
      the averaged loss.
docstr_tok: ["Compute","losses",":","GMMLoss","(","next_state",",","GMMPredicted",")","\/","(","STATE_DIM","+","2",")","+","MSE","(","reward",",","predicted_reward",")","+","BCE","(","not_terminal",",","logit_not_terminal",")","The","STATE_DIM","+","2","factor","is","here","to","counteract","the","fact","that","the","GMMLoss","scales","approximately","linearly","with","STATE_DIM",",","dim","of","states",".","All","losses","are","averaged","both","on","the","batch","and","the","sequence","dimensions","(","the","two","first","dimensions",")",".",":","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","-","not-terminal",":","(","SEQ_LEN",",","BATCH_SIZE",")","torch","tensor","-","next_state",":","(","SEQ_LEN",",","BATCH_SIZE",",","STATE_DIM",")","torch","tensor",":","param","state_dim",":","the","dimension","of","states",".","If","provided",",","use","it","to","normalize","gmm","loss",":","returns",":","dictionary","of","losses",",","containing","the","gmm",",","the","mse",",","the","bce","and","the","averaged","loss","."]
code: def get_loss(
      self,
      training_batch: rlt.MemoryNetworkInput,
      state_dim: Optional[int] = None,
  ):
      """
      Compute losses:
          GMMLoss(next_state, GMMPredicted) / (STATE_DIM + 2)
          + MSE(reward, predicted_reward)
          + BCE(not_terminal, logit_not_terminal)
  
      The STATE_DIM + 2 factor is here to counteract the fact that the GMMLoss scales
          approximately linearly with STATE_DIM, dim of states. All losses
          are averaged both on the batch and the sequence dimensions (the two first
          dimensions).
  
      :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
          - not-terminal: (SEQ_LEN, BATCH_SIZE) torch tensor
          - next_state: (SEQ_LEN, BATCH_SIZE, STATE_DIM) torch tensor
  
      :param state_dim: the dimension of states. If provided, use it to normalize
          gmm loss
  
      :returns: dictionary of losses, containing the gmm, the mse, the bce and
          the averaged loss.
      """
      assert isinstance(training_batch, rlt.MemoryNetworkInput)
      # mdnrnn's input should have seq_len as the first dimension
  
      mdnrnn_output = self.memory_network(
          training_batch.state, training_batch.action
      )
      # mus, sigmas: [seq_len, batch_size, num_gaussian, state_dim]
      mus, sigmas, logpi, rs, nts = (
          mdnrnn_output.mus,
          mdnrnn_output.sigmas,
          mdnrnn_output.logpi,
          mdnrnn_output.reward,
          mdnrnn_output.not_terminal,
      )
  
      next_state = training_batch.next_state.float_features
      not_terminal = training_batch.not_terminal.float()
      reward = training_batch.reward
      if self.params.fit_only_one_next_step:
          (
              next_state,
              not_terminal,
              reward,
              mus,
              sigmas,
              logpi,
              nts,
              rs,
          ) = tuple(
              map(
                  lambda x: x[-1:],
                  (
                      next_state,
                      not_terminal,
                      reward,
                      mus,
                      sigmas,
                      logpi,
                      nts,
                      rs,
                  ),
              )
          )
  
      gmm = (
          gmm_loss(next_state, mus, sigmas, logpi)
          * self.params.next_state_loss_weight
      )
      bce = (
          F.binary_cross_entropy_with_logits(nts, not_terminal)
          * self.params.not_terminal_loss_weight
      )
      mse = F.mse_loss(rs, reward) * self.params.reward_loss_weight
      if state_dim is not None:
          loss = gmm / (state_dim + 2) + bce + mse
      else:
          loss = gmm + bce + mse
      return {"gmm": gmm, "bce": bce, "mse": mse, "loss": 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 6 for calling the MDNRNNTrainer obj's underlying member method code with example usage: obj.get_loss(training_batch, state_dim) and returns: dict