class_docstr: Registry of models, preprocessors, and action distributions for envs.
  
  .. testcode::
      :skipif: True
  
      prep = ModelCatalog.get_preprocessor(env)
      observation = prep.transform(raw_observation)
  
      dist_class, dist_dim = ModelCatalog.get_action_dist(
          env.action_space, {})
      model = ModelCatalog.get_model_v2(
          obs_space, action_space, num_outputs, options)
      dist = dist_class(model.outputs, model)
      action = dist.sample()
class_docstr_tok: ["Registry","of","models",",","preprocessors",",","and","action","distributions","for","envs",".","..","testcode",":",":",":","skipif",":","True","prep","=","ModelCatalog.get_preprocessor","(","env",")","observation","=","prep.transform","(","raw_observation",")","dist_class",",","dist_dim","=","ModelCatalog.get_action_dist","(","env.action_space",",","{","}",")","model","=","ModelCatalog.get_model_v2","(","obs_space",",","action_space",",","num_outputs",",","options",")","dist","=","dist_class","(","model.outputs",",","model",")","action","=","dist.sample","(",")"]
docstr: Returns a distribution class and size for the given action space.
  
  Args:
      action_space: Action space of the target gym env.
      config (Optional[dict]): Optional model config.
      dist_type (Optional[Union[str, Type[ActionDistribution]]]):
          Identifier of the action distribution (str) interpreted as a
          hint or the actual ActionDistribution class to use.
      framework: One of "tf2", "tf", "torch", or "jax".
      kwargs: Optional kwargs to pass on to the Distribution's
          constructor.
  
  Returns:
      Tuple:
          - dist_class (ActionDistribution): Python class of the
              distribution.
          - dist_dim (int): The size of the input vector to the
              distribution.
docstr_tok: ["Returns","a","distribution","class","and","size","for","the","given","action","space",".","Args",":","action_space",":","Action","space","of","the","target","gym","env",".","config","(","Optional","[","dict","]",")",":","Optional","model","config",".","dist_type","(","Optional","[","Union","[","str",",","Type","[","ActionDistribution","]","]","]",")",":","Identifier","of","the","action","distribution","(","str",")","interpreted","as","a","hint","or","the","actual","ActionDistribution","class","to","use",".","framework",":","One","of","``","tf2","''",",","``","tf","''",",","``","torch","''",",","or","``","jax","''",".","kwargs",":","Optional","kwargs","to","pass","on","to","the","Distribution's","constructor",".","Returns",":","Tuple",":","-","dist_class","(","ActionDistribution",")",":","Python","class","of","the","distribution",".","-","dist_dim","(","int",")",":","The","size","of","the","input","vector","to","the","distribution","."]
code: def get_action_dist(
      action_space: gym.Space,
      config: ModelConfigDict,
      dist_type: Optional[Union[str, Type[ActionDistribution]]] = None,
      framework: str = "tf",
      **kwargs,
  ) -> (type, int):
      """Returns a distribution class and size for the given action space.
  
      Args:
          action_space: Action space of the target gym env.
          config (Optional[dict]): Optional model config.
          dist_type (Optional[Union[str, Type[ActionDistribution]]]):
              Identifier of the action distribution (str) interpreted as a
              hint or the actual ActionDistribution class to use.
          framework: One of "tf2", "tf", "torch", or "jax".
          kwargs: Optional kwargs to pass on to the Distribution's
              constructor.
  
      Returns:
          Tuple:
              - dist_class (ActionDistribution): Python class of the
                  distribution.
              - dist_dim (int): The size of the input vector to the
                  distribution.
      """
  
      dist_cls = None
      config = config or MODEL_DEFAULTS
      # Custom distribution given.
      if config.get("custom_action_dist"):
          custom_action_config = config.copy()
          action_dist_name = custom_action_config.pop(
              "custom_action_dist"
          )
          logger.debug(
              "Using custom action distribution {}".format(
                  action_dist_name
              )
          )
          dist_cls = _global_registry.get(
              RLLIB_ACTION_DIST, action_dist_name
          )
          return ModelCatalog._get_multi_action_distribution(
              dist_cls, action_space, custom_action_config, framework
          )
  
      # Dist_type is given directly as a class.
      elif (
          type(dist_type) is type
          and issubclass(dist_type, ActionDistribution)
          and dist_type
          not in (MultiActionDistribution, TorchMultiActionDistribution)
      ):
          dist_cls = dist_type
      # Box space -> DiagGaussian OR Deterministic.
      elif isinstance(action_space, Box):
          if action_space.dtype.name.startswith("int"):
              low_ = np.min(action_space.low)
              high_ = np.max(action_space.high)
              dist_cls = (
                  TorchMultiCategorical
                  if framework == "torch"
                  else MultiCategorical
              )
              num_cats = int(np.product(action_space.shape))
              return (
                  partial(
                      dist_cls,
                      input_lens=[
                          high_ - low_ + 1 for _ in range(num_cats)
                      ],
                      action_space=action_space,
                  ),
                  num_cats * (high_ - low_ + 1),
              )
          else:
              if len(action_space.shape) > 1:
                  raise UnsupportedSpaceException(
                      "Action space has multiple dimensions "
                      "{}. ".format(action_space.shape)
                      + "Consider reshaping this into a single dimension, "
                      "using a custom action distribution, "
                      "using a Tuple action space, or the multi-agent API."
                  )
              # TODO(sven): Check for bounds and return SquashedNormal, etc..
              if dist_type is None:
                  return (
                      partial(
                          TorchDiagGaussian
                          if framework == "torch"
                          else DiagGaussian,
                          action_space=action_space,
                      ),
                      DiagGaussian.required_model_output_shape(
                          action_space, config
                      ),
                  )
              elif dist_type == "deterministic":
                  dist_cls = (
                      TorchDeterministic
                      if framework == "torch"
                      else Deterministic
                  )
      # Discrete Space -> Categorical.
      elif isinstance(action_space, Discrete):
          if framework == "torch":
              dist_cls = TorchCategorical
          elif framework == "jax":
              from ray.rllib.models.jax.jax_action_dist import (
                  JAXCategorical,
              )
  
              dist_cls = JAXCategorical
          else:
              dist_cls = Categorical
      # Tuple/Dict Spaces -> MultiAction.
      elif dist_type in (
          MultiActionDistribution,
          TorchMultiActionDistribution,
      ) or isinstance(action_space, (Tuple, Dict)):
          return ModelCatalog._get_multi_action_distribution(
              (
                  MultiActionDistribution
                  if framework == "tf"
                  else TorchMultiActionDistribution
              ),
              action_space,
              config,
              framework,
          )
      # Simplex -> Dirichlet.
      elif isinstance(action_space, Simplex):
          dist_cls = (
              TorchDirichlet if framework == "torch" else Dirichlet
          )
      # MultiDiscrete -> MultiCategorical.
      elif isinstance(action_space, MultiDiscrete):
          dist_cls = (
              TorchMultiCategorical
              if framework == "torch"
              else MultiCategorical
          )
          return partial(dist_cls, input_lens=action_space.nvec), int(
              sum(action_space.nvec)
          )
      # Unknown type -> Error.
      else:
          raise NotImplementedError(
              "Unsupported args: {} {}".format(action_space, dist_type)
          )
  
      return dist_cls, int(
          dist_cls.required_model_output_shape(action_space, config)
      )
  
imports: from functools import partial
  import gymnasium
  from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
  import logging
  import numpy
  import tree
  from typing import List, Optional, Type, Union
  from ray.tune.registry import RLLIB_MODEL, RLLIB_ACTION_DIST, _global_registry
  from ray.rllib.models.action_dist import ActionDistribution
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.preprocessors import get_preprocessor, Preprocessor
  from ray.rllib.models.tf.tf_action_dist import Categorical, Deterministic, DiagGaussian, Dirichlet, MultiActionDistribution, MultiCategorical
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical, TorchDeterministic, TorchDirichlet, TorchDiagGaussian, TorchMultiActionDistribution, TorchMultiCategorical
  from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI
  from ray.rllib.utils.deprecation import DEPRECATED_VALUE, deprecation_warning
  from ray.rllib.utils.error import UnsupportedSpaceException
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.from_config import from_config
  from ray.rllib.utils.spaces.simplex import Simplex
  from ray.rllib.utils.spaces.space_utils import flatten_space
  from ray.rllib.utils.typing import ModelConfigDict, TensorType
label_desc: Use image node_id 1 for calling the ModelCatalog obj's underlying member method code with example usage: obj.get_action_dist(action_space, config, dist_type, framework) and returns: dist_cls,  int, ModelCatalog, partial,  unknown, partial,  DiagGaussian, ModelCatalog, partial,  int