class_docstr: An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.
class_docstr_tok: ["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,
  ):
      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.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 input layers.
      input_layer = tf.keras.layers.Input(
          shape=(None, self.num_outputs), name="inputs"
      )
  
      # Set self.num_outputs to the number of output nodes desired by the
      # caller of this constructor.
      self.num_outputs = num_outputs
  
      state_in_h = tf.keras.layers.Input(
          shape=(self.cell_size,), name="h"
      )
      state_in_c = tf.keras.layers.Input(
          shape=(self.cell_size,), name="c"
      )
      seq_in = tf.keras.layers.Input(
          shape=(), name="seq_in", dtype=tf.int32
      )
  
      # Preprocess observation with a hidden layer and send to LSTM cell
      lstm_out, state_h, state_c = tf.keras.layers.LSTM(
          self.cell_size,
          return_sequences=True,
          return_state=True,
          name="lstm",
      )(
          inputs=input_layer,
          mask=tf.sequence_mask(seq_in),
          initial_state=[state_in_h, state_in_c],
      )
  
      # Postprocess LSTM output with another hidden layer and compute values
      logits = tf.keras.layers.Dense(
          self.num_outputs,
          activation=tf.keras.activations.linear,
          name="logits",
      )(lstm_out)
      values = tf.keras.layers.Dense(1, activation=None, name="values")(
          lstm_out
      )
  
      # Create the RNN model
      self._rnn_model = tf.keras.Model(
          inputs=[input_layer, seq_in, state_in_h, state_in_c],
          outputs=[logits, values, state_h, state_c],
      )
      # Print out model summary in INFO logging mode.
      if logger.isEnabledFor(logging.INFO):
          self._rnn_model.summary()
  
      # 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 logging
  import tree
  from typing import Dict, List, Tuple
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  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_tf
  from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space
  from ray.rllib.utils.tf_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 with example: obj = LSTMWrapper(obs_space, action_space, num_outputs, model_config, name)