class_docstr: Mixture Density Network - Recurrent Neural Network
docstr: Forward pass of MDN-RNN
  
  :param actions: (SEQ_LEN, BATCH_SIZE, ACTION_DIM) torch tensor
  :param states: (SEQ_LEN, BATCH_SIZE, STATE_DIM) torch tensor
  
  :returns: parameters of the GMM prediction for the next state,
  gaussian prediction of the reward and logit prediction of
  non-terminality. And the RNN's outputs.
      - mus: (SEQ_LEN, BATCH_SIZE, NUM_GAUSSIANS, STATE_DIM) torch tensor
      - sigmas: (SEQ_LEN, BATCH_SIZE, NUM_GAUSSIANS, STATE_DIM) torch tensor
      - logpi: (SEQ_LEN, BATCH_SIZE, NUM_GAUSSIANS) torch tensor
      - reward: (SEQ_LEN, BATCH_SIZE) torch tensor
      - not_terminal: (SEQ_LEN, BATCH_SIZE) torch tensor
      - last_step_hidden_and_cell: TUPLE(
          (NUM_LAYERS, BATCH_SIZE, HIDDEN_SIZE),
          (NUM_LAYERS, BATCH_SIZE, HIDDEN_SIZE)
      ) torch tensor
      - all_steps_hidden: (SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE) torch tensor
docstr_tok: ["Forward","pass","of","MDN-RNN",":","param","actions",":","(","SEQ_LEN",",","BATCH_SIZE",",","ACTION_DIM",")","torch","tensor",":","param","states",":","(","SEQ_LEN",",","BATCH_SIZE",",","STATE_DIM",")","torch","tensor",":","returns",":","parameters","of","the","GMM","prediction","for","the","next","state",",","gaussian","prediction","of","the","reward","and","logit","prediction","of","non-terminality",".","And","the","RNN","'s","outputs",".","-","mus",":","(","SEQ_LEN",",","BATCH_SIZE",",","NUM_GAUSSIANS",",","STATE_DIM",")","torch","tensor","-","sigmas",":","(","SEQ_LEN",",","BATCH_SIZE",",","NUM_GAUSSIANS",",","STATE_DIM",")","torch","tensor","-","logpi",":","(","SEQ_LEN",",","BATCH_SIZE",",","NUM_GAUSSIANS",")","torch","tensor","-","reward",":","(","SEQ_LEN",",","BATCH_SIZE",")","torch","tensor","-","not_terminal",":","(","SEQ_LEN",",","BATCH_SIZE",")","torch","tensor","-","last_step_hidden_and_cell",":","TUPLE","(","(","NUM_LAYERS",",","BATCH_SIZE",",","HIDDEN_SIZE",")",",","(","NUM_LAYERS",",","BATCH_SIZE",",","HIDDEN_SIZE",")",")","torch","tensor","-","all_steps_hidden",":","(","SEQ_LEN",",","BATCH_SIZE",",","HIDDEN_SIZE",")","torch","tensor"]
code: def forward(
      self, actions: torch.Tensor, states: torch.Tensor, hidden=None
  ):
      """Forward pass of MDN-RNN
  
      :param actions: (SEQ_LEN, BATCH_SIZE, ACTION_DIM) torch tensor
      :param states: (SEQ_LEN, BATCH_SIZE, STATE_DIM) torch tensor
  
      :returns: parameters of the GMM prediction for the next state,
      gaussian prediction of the reward and logit prediction of
      non-terminality. And the RNN's outputs.
          - mus: (SEQ_LEN, BATCH_SIZE, NUM_GAUSSIANS, STATE_DIM) torch tensor
          - sigmas: (SEQ_LEN, BATCH_SIZE, NUM_GAUSSIANS, STATE_DIM) torch tensor
          - logpi: (SEQ_LEN, BATCH_SIZE, NUM_GAUSSIANS) torch tensor
          - reward: (SEQ_LEN, BATCH_SIZE) torch tensor
          - not_terminal: (SEQ_LEN, BATCH_SIZE) torch tensor
          - last_step_hidden_and_cell: TUPLE(
              (NUM_LAYERS, BATCH_SIZE, HIDDEN_SIZE),
              (NUM_LAYERS, BATCH_SIZE, HIDDEN_SIZE)
          ) torch tensor
          - all_steps_hidden: (SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE) torch tensor
      """
      device = next(self.parameters()).device
      actions = actions.to(device)
      states = states.to(device)
      seq_len, batch_size = actions.size(0), actions.size(1)
      ins = torch.cat([actions, states], dim=-1)
      all_steps_hidden, last_step_hidden_and_cell = self.rnn(
          ins, hidden
      )
      gmm_outs = self.gmm_linear(all_steps_hidden)
  
      stride = self.num_gaussians * self.state_dim
  
      mus = gmm_outs[:, :, :stride].view(
          seq_len, batch_size, self.num_gaussians, self.state_dim
      )
      sigmas = torch.exp(
          gmm_outs[:, :, stride : 2 * stride].view(
              seq_len, batch_size, self.num_gaussians, self.state_dim
          )
      )
      logpi = f.log_softmax(
          gmm_outs[
              :, :, 2 * stride : 2 * stride + self.num_gaussians
          ].view(seq_len, batch_size, self.num_gaussians),
          dim=-1,
      )
      reward = gmm_outs[:, :, -2]
      not_terminal = gmm_outs[:, :, -1]
  
      return (
          mus,
          sigmas,
          logpi,
          reward,
          not_terminal,
          all_steps_hidden,
          last_step_hidden_and_cell,
      )
  
imports: import logging
  from collections import deque
  from typing import NamedTuple
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from reagent.core import types
  from reagent.core.torch_utils import stack
  from torch.distributions.normal import Normal
label_desc: Use image node_id 2 for calling the MDNRNN obj's underlying member method code with example usage: obj.forward(actions, states, hidden) and returns: mus,  sigmas,  logpi,  reward,  not_terminal,  all_steps_hidden,  last_step_hidden_and_cell