class_docstr: Extension of the standard TorchModelV2 for SAC.
  
  To customize, do one of the following:
  - sub-class SACTorchModel and override one or more of its methods.
  - Use SAC's `q_model_config` and `policy_model` keys to tweak the default model
    behaviors (e.g. fcnet_hiddens, conv_filters, etc..).
  - Use SAC's `q_model_config->custom_model` and `policy_model->custom_model` keys
    to specify your own custom Q-model(s) and policy-models, which will be
    created within this SACTFModel (see `build_policy_model` and
    `build_q_model`.
  
  Note: It is not recommended to override the `forward` method for SAC. This
  would lead to shared weights (between policy and Q-nets), which will then
  not be optimized by either of the critic- or actor-optimizers!
  
  Data flow:
      `obs` -> forward() (should stay a noop method!) -> `model_out`
      `model_out` -> get_policy_output() -> pi(actions|obs)
      `model_out`, `actions` -> get_q_values() -> Q(s, a)
      `model_out`, `actions` -> get_twin_q_values() -> Q_twin(s, a)
class_docstr_tok: ["Extension","of","the","standard","TorchModelV2","for","SAC",".","To","customize",",","do","one","of","the","following",":","-","sub-class","SACTorchModel","and","override","one","or","more","of","its","methods",".","-","Use","SAC","'s","`","q_model_config","`","and","`","policy_model","`","keys","to","tweak","the","default","model","behaviors","(","e.g",".","fcnet_hiddens",",","conv_filters",",","etc","..",")",".","-","Use","SAC","'s","`","q_model_config-",">","custom_model","`","and","`","policy_model-",">","custom_model","`","keys","to","specify","your","own","custom","Q-model","(","s",")","and","policy-models",",","which","will","be","created","within","this","SACTFModel","(","see","`","build_policy_model","`","and","`","build_q_model","`",".","Note",":","It","is","not","recommended","to","override","the","`","forward","`","method","for","SAC",".","This","would","lead","to","shared","weights","(","between","policy","and","Q-nets",")",",","which","will","then","not","be","optimized","by","either","of","the","critic-","or","actor-optimizers","!","Data","flow",":","`","obs","`","-",">","forward","(",")","(","should","stay","a","noop","method","!",")","-",">","`","model_out","`","`","model_out","`","-",">","get_policy_output","(",")","-",">","pi","(","actions|obs",")","`","model_out","`",",","`","actions","`","-",">","get_q_values","(",")","-",">","Q","(","s",",","a",")","`","model_out","`",",","`","actions","`","-",">","get_twin_q_values","(",")","-",">","Q_twin","(","s",",","a",")"]
docstr: Returns distribution inputs and states given the output of
  policy.model().
  
  For continuous action spaces, these will be the mean/stddev
  distribution inputs for the (SquashedGaussian) action distribution.
  For discrete action spaces, these will be the logits for a categorical
  distribution.
  
  Args:
      model_out: Feature outputs from the model layers
          (result of doing `model(obs)`).
      state_in List(TensorType): State input for recurrent cells
      seq_lens: Sequence lengths of input- and state
          sequences
  
  Returns:
      TensorType: Distribution inputs for sampling actions.
docstr_tok: ["Returns","distribution","inputs","and","states","given","the","output","of","policy.model","(",")",".","For","continuous","action","spaces",",","these","will","be","the","mean\/stddev","distribution","inputs","for","the","(","SquashedGaussian",")","action","distribution",".","For","discrete","action","spaces",",","these","will","be","the","logits","for","a","categorical","distribution",".","Args",":","model_out",":","Feature","outputs","from","the","model","layers","(","result","of","doing","`","model","(","obs",")","`",")",".","state_in","List","(","TensorType",")",":","State","input","for","recurrent","cells","seq_lens",":","Sequence","lengths","of","input-","and","state","sequences","Returns",":","TensorType",":","Distribution","inputs","for","sampling","actions","."]
code: def get_action_model_outputs(
      self,
      model_out: TensorType,
      state_in: List[TensorType] = None,
      seq_lens: TensorType = None,
  ) -> (TensorType, List[TensorType]):
      """Returns distribution inputs and states given the output of
      policy.model().
  
      For continuous action spaces, these will be the mean/stddev
      distribution inputs for the (SquashedGaussian) action distribution.
      For discrete action spaces, these will be the logits for a categorical
      distribution.
  
      Args:
          model_out: Feature outputs from the model layers
              (result of doing `model(obs)`).
          state_in List(TensorType): State input for recurrent cells
          seq_lens: Sequence lengths of input- and state
              sequences
  
      Returns:
          TensorType: Distribution inputs for sampling actions.
      """
  
      def concat_obs_if_necessary(obs: TensorStructType):
          """Concat model outs if they come as original tuple observations."""
          if isinstance(obs, (list, tuple)):
              obs = torch.cat(obs, dim=-1)
          elif isinstance(obs, dict):
              obs = torch.cat(
                  [
                      torch.unsqueeze(val, 1)
                      if len(val.shape) == 1
                      else val
                      for val in tree.flatten(obs.values())
                  ],
                  dim=-1,
              )
          return obs
  
      if state_in is None:
          state_in = []
  
      if isinstance(model_out, dict) and "obs" in model_out:
          # Model outs may come as original Tuple observations
          if isinstance(self.action_model.obs_space, Box):
              model_out["obs"] = concat_obs_if_necessary(
                  model_out["obs"]
              )
          return self.action_model(model_out, state_in, seq_lens)
      else:
          if isinstance(self.action_model.obs_space, Box):
              model_out = concat_obs_if_necessary(model_out)
          return self.action_model(
              {"obs": model_out}, state_in, seq_lens
          )
  
imports: import gymnasium
  from gymnasium.spaces import Box, Discrete
  import numpy
  import tree
  from typing import Dict, List, Optional
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.spaces.simplex import Simplex
  from ray.rllib.utils.typing import ModelConfigDict, TensorType, TensorStructType
label_desc: Use image node_id 8 for calling the SACTorchModel obj's underlying member method code with example usage: obj.get_action_model_outputs(model_out, state_in, seq_lens) and returns: obs, self, self