class_docstr: Helper class to simplify implementing RNN models with TorchModelV2.
  
  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, nn.Module)``::
  
      def __init__(self, obs_space, num_outputs):
          nn.Module.__init__(self)
          super().__init__(obs_space, action_space, num_outputs,
                           model_config, name)
          self.obs_size = _get_size(obs_space)
          self.rnn_hidden_dim = model_config["lstm_cell_size"]
          self.fc1 = nn.Linear(self.obs_size, self.rnn_hidden_dim)
          self.rnn = nn.GRUCell(self.rnn_hidden_dim, self.rnn_hidden_dim)
          self.fc2 = nn.Linear(self.rnn_hidden_dim, num_outputs)
  
          self.value_branch = nn.Linear(self.rnn_hidden_dim, 1)
          self._cur_value = None
  
      @override(ModelV2)
      def get_initial_state(self):
          # Place hidden states on same device as model.
          h = [self.fc1.weight.new(
              1, self.rnn_hidden_dim).zero_().squeeze(0)]
          return h
  
      @override(ModelV2)
      def value_function(self):
          assert self._cur_value is not None, "must call forward() first"
          return self._cur_value
  
      @override(RecurrentNetwork)
      def forward_rnn(self, input_dict, state, seq_lens):
          x = nn.functional.relu(self.fc1(input_dict["obs_flat"].float()))
          h_in = state[0].reshape(-1, self.rnn_hidden_dim)
          h = self.rnn(x, h_in)
          q = self.fc2(h)
          self._cur_value = self.value_branch(h).squeeze(1)
          return q, [h]
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, size].
      seq_lens: 1D tensor holding input sequence lengths.
          Note: len(seq_lens) == B.
  
  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].
  
  Examples:
      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",",","size","]",".","seq_lens",":","1D","tensor","holding","input","sequence","lengths",".","Note",":","len","(","seq_lens",")","==","B",".","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","]",".","Examples",":","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, size].
          seq_lens: 1D tensor holding input sequence lengths.
              Note: len(seq_lens) == B.
  
      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].
  
      Examples:
          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 an RNN model"
      )
  
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 2 for calling the RecurrentNetwork obj's underlying member method code with example usage: obj.forward_rnn(inputs, state, seq_lens) without return types