class_docstr: Helper class to simplify implementing RNN models with TFModelV2.
  
  Instead of implementing forward(), you can implement forward_rnn() which
  takes batches with the time dimension added already.
  
  Here is an example implementation for a subclass
  ``MyRNNClass(RecurrentNetwork)``::
  
      def __init__(self, *args, **kwargs):
          super(MyModelClass, self).__init__(*args, **kwargs)
          cell_size = 256
  
          # Define input layers
          input_layer = tf.keras.layers.Input(
              shape=(None, obs_space.shape[0]))
          state_in_h = tf.keras.layers.Input(shape=(256, ))
          state_in_c = tf.keras.layers.Input(shape=(256, ))
          seq_in = tf.keras.layers.Input(shape=(), dtype=tf.int32)
  
          # Send to LSTM cell
          lstm_out, state_h, state_c = tf.keras.layers.LSTM(
              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])
          output_layer = tf.keras.layers.Dense(...)(lstm_out)
  
          # Create the RNN model
          self.rnn_model = tf.keras.Model(
              inputs=[input_layer, seq_in, state_in_h, state_in_c],
              outputs=[output_layer, state_h, state_c])
          self.rnn_model.summary()
class_docstr_tok: ["Helper","class","to","simplify","implementing","RNN","models","with","TFModelV2",".","Instead","of","implementing","forward","(",")",",","you","can","implement","forward_rnn","(",")","which","takes","batches","with","the","time","dimension","added","already",".","Here","is","an","example","implementation","for","a","subclass","``","MyRNNClass","(","RecurrentNetwork",")","``",":",":","def","__init__","(","self",",","*","args",",","*","*","kwargs",")",":","super","(","MyModelClass",",","self",")",".__init__","(","*","args",",","*","*","kwargs",")","cell_size","=","256","#","Define","input","layers","input_layer","=","tf.keras.layers.Input","(","shape=","(","None",",","obs_space.shape","[","0","]",")",")","state_in_h","=","tf.keras.layers.Input","(","shape=","(","256",",",")",")","state_in_c","=","tf.keras.layers.Input","(","shape=","(","256",",",")",")","seq_in","=","tf.keras.layers.Input","(","shape=","(",")",",","dtype=tf.int32",")","#","Send","to","LSTM","cell","lstm_out",",","state_h",",","state_c","=","tf.keras.layers.LSTM","(","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","]",")","output_layer","=","tf.keras.layers.Dense","(","...",")","(","lstm_out",")","#","Create","the","RNN","model","self.rnn_model","=","tf.keras.Model","(","inputs=","[","input_layer",",","seq_in",",","state_in_h",",","state_in_c","]",",","outputs=","[","output_layer",",","state_h",",","state_c","]",")","self.rnn_model.summary","(",")"]
docstr: Call the model with the given input tensors and state.
  
  Args:
      inputs: observation tensor with shape [B, T, obs_size].
      state: list of state tensors, each with shape [B, T, size].
      seq_lens: 1d tensor holding input sequence lengths.
  
  Returns:
      (outputs, new_state): The model output tensor of shape
          [B, T, num_outputs] and the list of new state tensors each with
          shape [B, size].
  
  Sample implementation for the ``MyRNNClass`` example::
  
      def forward_rnn(self, inputs, state, seq_lens):
          model_out, h, c = self.rnn_model([inputs, seq_lens] + state)
          return model_out, [h, c]
docstr_tok: ["Call","the","model","with","the","given","input","tensors","and","state",".","Args",":","inputs",":","observation","tensor","with","shape","[","B",",","T",",","obs_size","]",".","state",":","list","of","state","tensors",",","each","with","shape","[","B",",","T",",","size","]",".","seq_lens",":","1d","tensor","holding","input","sequence","lengths",".","Returns",":","(","outputs",",","new_state",")",":","The","model","output","tensor","of","shape","[","B",",","T",",","num_outputs","]","and","the","list","of","new","state","tensors","each","with","shape","[","B",",","size","]",".","Sample","implementation","for","the","``","MyRNNClass","``","example",":",":","def","forward_rnn","(","self",",","inputs",",","state",",","seq_lens",")",":","model_out",",","h",",","c","=","self.rnn_model","(","[","inputs",",","seq_lens","]","+","state",")","return","model_out",",","[","h",",","c","]"]
code: def forward_rnn(
      self,
      inputs: TensorType,
      state: List[TensorType],
      seq_lens: TensorType,
  ) -> Tuple[TensorType, List[TensorType]]:
      """Call the model with the given input tensors and state.
  
      Args:
          inputs: observation tensor with shape [B, T, obs_size].
          state: list of state tensors, each with shape [B, T, size].
          seq_lens: 1d tensor holding input sequence lengths.
  
      Returns:
          (outputs, new_state): The model output tensor of shape
              [B, T, num_outputs] and the list of new state tensors each with
              shape [B, size].
  
      Sample implementation for the ``MyRNNClass`` example::
  
          def forward_rnn(self, inputs, state, seq_lens):
              model_out, h, c = self.rnn_model([inputs, seq_lens] + state)
              return model_out, [h, c]
      """
      raise NotImplementedError(
          "You must implement this for a RNN model"
      )
  
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 RecurrentNetwork obj's underlying member method code with example usage: obj.forward_rnn(inputs, state, seq_lens) without return types