class_docstr: An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.
code: def __init__(
      self,
      obs_space: gym.spaces.Space,
      action_space: gym.spaces.Space,
      num_outputs: int,
      model_config: ModelConfigDict,
      name: str,
  ):
      nn.Module.__init__(self)
      super(LSTMWrapper, self).__init__(
          obs_space, action_space, None, model_config, name
      )
  
      # At this point, self.num_outputs is the number of nodes coming
      # from the wrapped (underlying) model. In other words, self.num_outputs
      # is the input size for the LSTM layer.
      # If None, set it to the observation space.
      if self.num_outputs is None:
          self.num_outputs = int(np.product(self.obs_space.shape))
  
      self.cell_size = model_config["lstm_cell_size"]
      self.time_major = model_config.get("_time_major", False)
      self.use_prev_action = model_config["lstm_use_prev_action"]
      self.use_prev_reward = model_config["lstm_use_prev_reward"]
  
      self.action_space_struct = get_base_struct_from_space(
          self.action_space
      )
      self.action_dim = 0
  
      for space in tree.flatten(self.action_space_struct):
          if isinstance(space, Discrete):
              self.action_dim += space.n
          elif isinstance(space, MultiDiscrete):
              self.action_dim += np.sum(space.nvec)
          elif space.shape is not None:
              self.action_dim += int(np.product(space.shape))
          else:
              self.action_dim += int(len(space))
  
      # Add prev-action/reward nodes to input to LSTM.
      if self.use_prev_action:
          self.num_outputs += self.action_dim
      if self.use_prev_reward:
          self.num_outputs += 1
  
      # Define actual LSTM layer (with num_outputs being the nodes coming
      # from the wrapped (underlying) layer).
      self.lstm = nn.LSTM(
          self.num_outputs,
          self.cell_size,
          batch_first=not self.time_major,
      )
  
      # Set self.num_outputs to the number of output nodes desired by the
      # caller of this constructor.
      self.num_outputs = num_outputs
  
      # Postprocess LSTM output with another hidden layer and compute values.
      self._logits_branch = SlimFC(
          in_size=self.cell_size,
          out_size=self.num_outputs,
          activation_fn=None,
          initializer=torch.nn.init.xavier_uniform_,
      )
      self._value_branch = SlimFC(
          in_size=self.cell_size,
          out_size=1,
          activation_fn=None,
          initializer=torch.nn.init.xavier_uniform_,
      )
  
      # __sphinx_doc_begin__
      # Add prev-a/r to this model's view, if required.
      if model_config["lstm_use_prev_action"]:
          self.view_requirements[
              SampleBatch.PREV_ACTIONS
          ] = ViewRequirement(
              SampleBatch.ACTIONS, space=self.action_space, shift=-1
          )
      if model_config["lstm_use_prev_reward"]:
          self.view_requirements[
              SampleBatch.PREV_REWARDS
          ] = ViewRequirement(SampleBatch.REWARDS, shift=-1)
  
imports: import numpy
  import gymnasium
  from gymnasium.spaces import Discrete, MultiDiscrete
  import tree
  from typing import Dict, List, Union, Tuple
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.misc import SlimFC
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.policy.rnn_sequencing import add_time_dimension
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.utils.annotations import override, DeveloperAPI
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space
  from ray.rllib.utils.torch_utils import flatten_inputs_to_1d_tensor, one_hot
  from ray.rllib.utils.typing import ModelConfigDict, TensorType
  from ray.rllib.utils.deprecation import deprecation_warning
  from ray.util.debug import log_once
label_desc: Use image node_id 1 to create a new LSTMWrapper object from inherited base classes: RecurrentNetwork, nn with example: obj = LSTMWrapper(obs_space, action_space, num_outputs, model_config, name)