code: import torch
  import torch.nn.functional as F
  
  from rlpyt.models.conv2d import Conv2dHeadModel
  from rlpyt.utils.collections import namedarraytuple
  from rlpyt.utils.tensor import infer_leading_dims, restore_leading_dims
  
  RnnState = namedarraytuple(
      "RnnState", ["h", "c"]
  )  # For downstream namedarraytuples to work
  
  
  class AtariLstmModel(torch.nn.Module):
      """Recurrent model for Atari agents: a convolutional network into an FC layer
      into an LSTM which outputs action probabilities and state-value estimate.
      """
  
      def __init__(
          self,
          image_shape,
          output_size,
          fc_sizes=512,  # Between conv and lstm.
          lstm_size=512,
          use_maxpool=False,
          channels=None,  # None uses default.
          kernel_sizes=None,
          strides=None,
          paddings=None,
      ):
          """Instantiate neural net module according to inputs."""
          super().__init__()
          self.conv = Conv2dHeadModel(
              image_shape=image_shape,
              channels=channels or [16, 32],
              kernel_sizes=kernel_sizes or [8, 4],
              strides=strides or [4, 2],
              paddings=paddings or [0, 1],
              use_maxpool=use_maxpool,
              hidden_sizes=fc_sizes,  # Applies nonlinearity at end.
          )
          self.lstm = torch.nn.LSTM(self.conv.output_size + output_size + 1, lstm_size)
          self.pi = torch.nn.Linear(lstm_size, output_size)
          self.value = torch.nn.Linear(lstm_size, 1)
  
      def forward(self, image, prev_action, prev_reward, init_rnn_state):
          """
          Compute action probabilities and value estimate from input state.
          Infers leading dimensions of input: can be [T,B], [B], or []; provides
          returns with same leading dims.  Convolution layers process as [T*B,
          *image_shape], with T=1,B=1 when not given.  Expects uint8 images in
          [0,255] and converts them to float32 in [0,1] (to minimize image data
          storage and transfer).  Recurrent layers processed as [T,B,H]. Used in
          both sampler and in algorithm (both via the agent).  Also returns the
          next RNN state.
          """
          img = image.type(torch.float)  # Expect torch.uint8 inputs
          img = img.mul_(1.0 / 255)  # From [0-255] to [0-1], in place.
  
          # Infer (presence of) leading dimensions: [T,B], [B], or [].
          lead_dim, T, B, img_shape = infer_leading_dims(img, 3)
  
          fc_out = self.conv(img.view(T * B, *img_shape))
          lstm_input = torch.cat(
              [
                  fc_out.view(T, B, -1),
                  prev_action.view(T, B, -1),  # Assumed onehot.
                  prev_reward.view(T, B, 1),
              ],
              dim=2,
          )
          init_rnn_state = None if init_rnn_state is None else tuple(init_rnn_state)
          lstm_out, (hn, cn) = self.lstm(lstm_input, init_rnn_state)
          pi = F.softmax(self.pi(lstm_out.view(T * B, -1)), dim=-1)
          v = self.value(lstm_out.view(T * B, -1)).squeeze(-1)
  
          # Restore leading dimensions: [T,B], [B], or [], as input.
          pi, v = restore_leading_dims((pi, v), lead_dim, T, B)
          # Model should always leave B-dimension in rnn state: [N,B,H].
          next_rnn_state = RnnState(h=hn, c=cn)
  
          return pi, v, next_rnn_state
  
apis: torch.nn.LSTM torch.nn.Linear