code: def test_gmm_loss(self):
      # seq_len x batch_size x gaussian_size x feature_size
      # 1 x 1 x 2 x 2
      mus = torch.Tensor([[[[0.0, 0.0], [6.0, 6.0]]]])
      sigmas = torch.Tensor([[[[2.0, 2.0], [2.0, 2.0]]]])
      # seq_len x batch_size x gaussian_size
      pi = torch.Tensor([[[0.5, 0.5]]])
      logpi = torch.log(pi)
  
      # seq_len x batch_size x feature_size
      batch = torch.Tensor([[[3.0, 3.0]]])
      gl = gmm_loss(batch, mus, sigmas, logpi)
  
      # first component, first dimension
      n11 = Normal(mus[0, 0, 0, 0], sigmas[0, 0, 0, 0])
      # first component, second dimension
      n12 = Normal(mus[0, 0, 0, 1], sigmas[0, 0, 0, 1])
      p1 = (
          pi[0, 0, 0]
          * torch.exp(n11.log_prob(batch[0, 0, 0]))
          * torch.exp(n12.log_prob(batch[0, 0, 1]))
      )
      # second component, first dimension
      n21 = Normal(mus[0, 0, 1, 0], sigmas[0, 0, 1, 0])
      # second component, second dimension
      n22 = Normal(mus[0, 0, 1, 1], sigmas[0, 0, 1, 1])
      p2 = (
          pi[0, 0, 1]
          * torch.exp(n21.log_prob(batch[0, 0, 0]))
          * torch.exp(n22.log_prob(batch[0, 0, 1]))
      )
  
      logger.info(
          "gmm loss={}, p1={}, p2={}, p1+p2={}, -log(p1+p2)={}".format(
              gl, p1, p2, p1 + p2, -(torch.log(p1 + p2))
          )
      )
      assert -(torch.log(p1 + p2)) == gl
  
file_path: ReAgent/reagent/test/world_model/test_mdnrnn.py
imports: import logging
  import unittest
  import numpy
  import torch
  from reagent.core.parameters import MDNRNNTrainerParameters
  from reagent.models.mdn_rnn import gmm_loss, MDNRNNMemoryPool
  from reagent.models.world_model import MemoryNetwork
  from reagent.reporting.world_model_reporter import WorldModelReporter
  from reagent.test.world_model.simulated_world_model import SimulatedWorldModel
  from reagent.training.world_model.mdnrnn_trainer import MDNRNNTrainer
  from torch.distributions.categorical import Categorical
  from torch.distributions.normal import Normal
label_desc: Use image node_id 1 for calling the TestMDNRNN obj's underlying member method code with example usage: obj.test_gmm_loss() without return types