code: def _test_mdnrnn_simulate_world(self, use_gpu=False):
      num_epochs = 300
      num_episodes = 400
      batch_size = 200
      action_dim = 2
      seq_len = 5
      state_dim = 2
      simulated_num_gaussians = 2
      mdrnn_num_gaussians = 2
      simulated_num_hidden_layers = 1
      simulated_num_hiddens = 3
      mdnrnn_num_hidden_layers = 1
      mdnrnn_num_hiddens = 10
      adam_lr = 0.01
  
      replay_buffer = MDNRNNMemoryPool(
          max_replay_memory_size=num_episodes
      )
      swm = SimulatedWorldModel(
          action_dim=action_dim,
          state_dim=state_dim,
          num_gaussians=simulated_num_gaussians,
          lstm_num_hidden_layers=simulated_num_hidden_layers,
          lstm_num_hiddens=simulated_num_hiddens,
      )
  
      possible_actions = torch.eye(action_dim)
      for _ in range(num_episodes):
          cur_state_mem = torch.zeros((seq_len, state_dim))
          next_state_mem = torch.zeros((seq_len, state_dim))
          action_mem = torch.zeros((seq_len, action_dim))
          reward_mem = torch.zeros(seq_len)
          not_terminal_mem = torch.zeros(seq_len)
          next_mus_mem = torch.zeros(
              (seq_len, simulated_num_gaussians, state_dim)
          )
  
          swm.init_hidden(batch_size=1)
          next_state = torch.randn((1, 1, state_dim))
          for s in range(seq_len):
              cur_state = next_state
              action = possible_actions[
                  np.random.randint(action_dim)
              ].view(1, 1, action_dim)
              next_mus, reward = swm(action, cur_state)
  
              not_terminal = 1
              if s == seq_len - 1:
                  not_terminal = 0
  
              # randomly draw for next state
              next_pi = (
                  torch.ones(simulated_num_gaussians)
                  / simulated_num_gaussians
              )
              index = Categorical(next_pi).sample((1,)).long().item()
              next_state = next_mus[0, 0, index].view(1, 1, state_dim)
  
              cur_state_mem[s] = cur_state.detach()
              action_mem[s] = action
              reward_mem[s] = reward.detach()
              not_terminal_mem[s] = not_terminal
              next_state_mem[s] = next_state.detach()
              next_mus_mem[s] = next_mus.detach()
  
          replay_buffer.insert_into_memory(
              cur_state_mem,
              action_mem,
              next_state_mem,
              reward_mem,
              not_terminal_mem,
          )
  
      num_batch = num_episodes // batch_size
      mdnrnn_params = MDNRNNTrainerParameters(
          hidden_size=mdnrnn_num_hiddens,
          num_hidden_layers=mdnrnn_num_hidden_layers,
          learning_rate=adam_lr,
          num_gaussians=mdrnn_num_gaussians,
      )
      mdnrnn_net = MemoryNetwork(
          state_dim=state_dim,
          action_dim=action_dim,
          num_hiddens=mdnrnn_params.hidden_size,
          num_hidden_layers=mdnrnn_params.num_hidden_layers,
          num_gaussians=mdnrnn_params.num_gaussians,
      )
      if use_gpu:
          mdnrnn_net = mdnrnn_net.cuda()
      trainer = MDNRNNTrainer(
          memory_network=mdnrnn_net,
          params=mdnrnn_params,
          cum_loss_hist=num_batch,
      )
      reporter = WorldModelReporter(report_interval=1)
      trainer.set_reporter(reporter)
  
      optimizer = trainer.configure_optimizers()[0]
      for e in range(num_epochs):
          for i in range(num_batch):
              training_batch = replay_buffer.sample_memories(
                  batch_size, use_gpu=use_gpu
              )
              optimizer.zero_grad()
              loss = next(trainer.train_step_gen(training_batch, i))
              loss.backward()
              optimizer.step()
  
              logger.info(
                  "{}-th epoch, {}-th minibatch: \n"
                  "loss={}, bce={}, gmm={}, mse={} \n"
                  "cum loss={}, cum bce={}, cum gmm={}, cum mse={}\n".format(
                      e,
                      i,
                      reporter.loss.values[-1],
                      reporter.bce.values[-1],
                      reporter.gmm.values[-1],
                      reporter.mse.values[-1],
                      np.mean(reporter.loss.values[-100:]),
                      np.mean(reporter.bce.values[-100:]),
                      np.mean(reporter.gmm.values[-100:]),
                      np.mean(reporter.mse.values[-100:]),
                  )
              )
  
              if (
                  np.mean(reporter.loss.values[-100:]) < 0
                  and np.mean(reporter.gmm.values[-100:]) < -3.0
                  and np.mean(reporter.bce.values[-100:]) < 0.6
                  and np.mean(reporter.mse.values[-100:]) < 0.2
              ):
                  return
  
      raise RuntimeError(
          "losses not reduced significantly during training"
      )
  
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 4 for calling the TestMDNRNN obj's underlying member method code with example usage: obj._test_mdnrnn_simulate_world(use_gpu) without return types