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 forward(
      self,
      input_dict: Dict[str, TensorType],
      state: List[TensorType],
      seq_lens: TensorType,
  ) -> Tuple[TensorType, List[TensorType]]:
      assert seq_lens is not None
      # Push obs through "unwrapped" net's `forward()` first.
      wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
  
      # Concat. prev-action/reward if required.
      prev_a_r = []
  
      # Prev actions.
      if self.model_config["lstm_use_prev_action"]:
          prev_a = input_dict[SampleBatch.PREV_ACTIONS]
          # If actions are not processed yet (in their original form as
          # have been sent to environment):
          # Flatten/one-hot into 1D array.
          if self.model_config["_disable_action_flattening"]:
              prev_a_r.append(
                  flatten_inputs_to_1d_tensor(
                      prev_a,
                      spaces_struct=self.action_space_struct,
                      time_axis=False,
                  )
              )
          # If actions are already flattened (but not one-hot'd yet!),
          # one-hot discrete/multi-discrete actions here.
          else:
              if isinstance(
                  self.action_space, (Discrete, MultiDiscrete)
              ):
                  prev_a = one_hot(prev_a, self.action_space)
              prev_a_r.append(
                  tf.reshape(
                      tf.cast(prev_a, tf.float32), [-1, self.action_dim]
                  )
              )
      # Prev rewards.
      if self.model_config["lstm_use_prev_reward"]:
          prev_a_r.append(
              tf.reshape(
                  tf.cast(
                      input_dict[SampleBatch.PREV_REWARDS], tf.float32
                  ),
                  [-1, 1],
              )
          )
  
      # Concat prev. actions + rewards to the "main" input.
      if prev_a_r:
          wrapped_out = tf.concat([wrapped_out] + prev_a_r, axis=1)
  
      # Push everything through our LSTM.
      input_dict["obs_flat"] = wrapped_out
      return super().forward(input_dict, state, seq_lens)
  
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 2 for calling the LSTMWrapper obj's underlying member method code with example usage: obj.forward(input_dict, state, seq_lens) and returns: super