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: Get the initial recurrent state values for the model.
  
  Returns:
      list of np.array objects, if any
  
  Sample implementation for the ``MyRNNClass`` example::
  
      def get_initial_state(self):
          return [
              np.zeros(self.cell_size, np.float32),
              np.zeros(self.cell_size, np.float32),
          ]
docstr_tok: ["Get","the","initial","recurrent","state","values","for","the","model",".","Returns",":","list","of","np.array","objects",",","if","any","Sample","implementation","for","the","``","MyRNNClass","``","example",":",":","def","get_initial_state","(","self",")",":","return","[","np.zeros","(","self.cell_size",",","np.float32",")",",","np.zeros","(","self.cell_size",",","np.float32",")",",","]"]
code: def get_initial_state(self) -> List[TensorType]:
      """Get the initial recurrent state values for the model.
  
      Returns:
          list of np.array objects, if any
  
      Sample implementation for the ``MyRNNClass`` example::
  
          def get_initial_state(self):
              return [
                  np.zeros(self.cell_size, np.float32),
                  np.zeros(self.cell_size, np.float32),
              ]
      """
      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 3 for calling the RecurrentNetwork obj's underlying member method code with example usage: obj.get_initial_state() without return types