class_docstr: Describes the sub-module-architectures to be used in RLModules.
  
  RLlib's native RLModules get their Models from a Catalog object.
  By default, that Catalog builds the configs it has as attributes.
  This component was build to be hackable and extensible. You can inject custom
  components into RL Modules by overriding the `build_xxx` methods of this class.
  Note that it is recommended to write a custom RL Module for a single use-case.
  Modifications to Catalogs mostly make sense if you want to reuse the same
  Catalog for different RL Modules. For example if you have written a custom
  encoder and want to inject it into different RL Modules (e.g. for PPO, DQN, etc.).
  You can influence the decision tree that determines the sub-components by modifying
  `Catalog._determine_components_hook`.
  
  Usage example:
  
  # Define a custom catalog
  
  .. testcode::
  
      import torch
      import gymnasium as gym
      from ray.rllib.core.models.configs import MLPHeadConfig
      from ray.rllib.core.models.catalog import Catalog
  
      class MyCatalog(Catalog):
          def __init__(
              self,
              observation_space: gym.Space,
              action_space: gym.Space,
              model_config_dict: dict,
          ):
              super().__init__(observation_space, action_space, model_config_dict)
              self.my_model_config = MLPHeadConfig(
                  hidden_layer_dims=[64, 32],
                  input_dims=[self.observation_space.shape[0]],
              )
  
          def build_my_head(self, framework: str):
              return self.my_model_config.build(framework=framework)
  
      # With that, RLlib can build and use models from this catalog like this:
      catalog = MyCatalog(gym.spaces.Box(0, 1), gym.spaces.Box(0, 1), {})
      my_head = catalog.build_my_head(framework="torch")
  
      # Make a call to the built model.
      out = my_head(torch.Tensor([[1]]))
class_docstr_tok: ["Describes","the","sub-module-architectures","to","be","used","in","RLModules",".","RLlib","'s","native","RLModules","get","their","Models","from","a","Catalog","object",".","By","default",",","that","Catalog","builds","the","configs","it","has","as","attributes",".","This","component","was","build","to","be","hackable","and","extensible",".","You","can","inject","custom","components","into","RL","Modules","by","overriding","the","`","build_xxx","`","methods","of","this","class",".","Note","that","it","is","recommended","to","write","a","custom","RL","Module","for","a","single","use-case",".","Modifications","to","Catalogs","mostly","make","sense","if","you","want","to","reuse","the","same","Catalog","for","different","RL","Modules",".","For","example","if","you","have","written","a","custom","encoder","and","want","to","inject","it","into","different","RL","Modules","(","e.g",".","for","PPO",",","DQN",",","etc.",")",".","You","can","influence","the","decision","tree","that","determines","the","sub-components","by","modifying","`","Catalog._determine_components_hook","`",".","Usage","example",":","#","Define","a","custom","catalog","..","testcode",":",":","import","torch","import","gymnasium","as","gym","from","ray.rllib.core.models.configs","import","MLPHeadConfig","from","ray.rllib.core.models.catalog","import","Catalog","class","MyCatalog","(","Catalog",")",":","def","__init__","(","self",",","observation_space",":","gym.Space",",","action_space",":","gym.Space",",","model_config_dict",":","dict",",",")",":","super","(",")",".__init__","(","observation_space",",","action_space",",","model_config_dict",")","self.my_model_config","=","MLPHeadConfig","(","hidden_layer_dims=","[","64",",","32","]",",","input_dims=","[","self.observation_space.shape","[","0","]","]",",",")","def","build_my_head","(","self",",","framework",":","str",")",":","return","self.my_model_config.build","(","framework=framework",")","#","With","that",",","RLlib","can","build","and","use","models","from","this","catalog","like","this",":","catalog","=","MyCatalog","(","gym.spaces.Box","(","0",",","1",")",",","gym.spaces.Box","(","0",",","1",")",",","{","}",")","my_head","=","catalog.build_my_head","(","framework=","''","torch","''",")","#","Make","a","call","to","the","built","model",".","out","=","my_head","(","torch.Tensor","(","[","[","1","]","]",")",")"]
docstr: Returns a distribution class for the given action space.
  
  You can get the required input dimension for the distribution by calling
  `action_dict_cls.required_input_dim(action_space)`
  on the retrieved class. This is useful, because the Catalog needs to find out
  about the required input dimension for the distribution before the model that
  outputs these inputs is configured.
  
  Args:
      action_space: Action space of the target gym env.
      framework: The framework to use.
  
  Returns:
      The distribution class for the given action space.
docstr_tok: ["Returns","a","distribution","class","for","the","given","action","space",".","You","can","get","the","required","input","dimension","for","the","distribution","by","calling","`","action_dict_cls.required_input_dim","(","action_space",")","`","on","the","retrieved","class",".","This","is","useful",",","because","the","Catalog","needs","to","find","out","about","the","required","input","dimension","for","the","distribution","before","the","model","that","outputs","these","inputs","is","configured",".","Args",":","action_space",":","Action","space","of","the","target","gym","env",".","framework",":","The","framework","to","use",".","Returns",":","The","distribution","class","for","the","given","action","space","."]
code: def _get_dist_cls_from_action_space(
      cls,
      action_space: gym.Space,
      *,
      framework: Optional[str] = None,
  ) -> Distribution:
      """Returns a distribution class for the given action space.
  
      You can get the required input dimension for the distribution by calling
      `action_dict_cls.required_input_dim(action_space)`
      on the retrieved class. This is useful, because the Catalog needs to find out
      about the required input dimension for the distribution before the model that
      outputs these inputs is configured.
  
      Args:
          action_space: Action space of the target gym env.
          framework: The framework to use.
  
      Returns:
          The distribution class for the given action space.
      """
      # This method is structured in two steps:
      # Firstly, construct a dictionary containing the available distribution classes.
      # Secondly, return the correct distribution class for the given action space.
  
      # Step 1: Construct the dictionary.
  
      class DistEnum(enum.Enum):
          Categorical = "Categorical"
          DiagGaussian = "Gaussian"
          Deterministic = "Deterministic"
          MultiDistribution = "MultiDistribution"
          MultiCategorical = "MultiCategorical"
  
      if framework == "torch":
          from ray.rllib.models.torch.torch_distributions import (
              TorchCategorical,
              TorchDeterministic,
              TorchDiagGaussian,
          )
  
          distribution_dicts = {
              DistEnum.Deterministic: TorchDeterministic,
              DistEnum.DiagGaussian: TorchDiagGaussian,
              DistEnum.Categorical: TorchCategorical,
          }
      elif framework == "tf2":
          from ray.rllib.models.tf.tf_distributions import (
              TfCategorical,
              TfDeterministic,
              TfDiagGaussian,
          )
  
          distribution_dicts = {
              DistEnum.Deterministic: TfDeterministic,
              DistEnum.DiagGaussian: TfDiagGaussian,
              DistEnum.Categorical: TfCategorical,
          }
      else:
          raise ValueError(
              f"Unknown framework: {framework}. Only 'torch' and 'tf2' are "
              "supported for RLModule Catalogs."
          )
  
      # Only add a MultiAction distribution class to the dict if we can compute its
      # components (we need a Tuple/Dict space for this).
      if isinstance(action_space, (Tuple, Dict)):
          partial_multi_action_distribution_cls = (
              _multi_action_dist_partial_helper(
                  catalog_cls=cls,
                  action_space=action_space,
                  framework=framework,
              )
          )
  
          distribution_dicts[
              DistEnum.MultiDistribution
          ] = partial_multi_action_distribution_cls
  
      # Only add a MultiCategorical distribution class to the dict if we can compute
      # its components (we need a MultiDiscrete space for this).
      if isinstance(action_space, MultiDiscrete):
          partial_multi_categorical_distribution_cls = (
              _multi_categorical_dist_partial_helper(
                  action_space=action_space,
                  framework=framework,
              )
          )
  
          distribution_dicts[
              DistEnum.MultiCategorical
          ] = partial_multi_categorical_distribution_cls
  
      # Step 2: Return the correct distribution class for the given action space.
  
      # Box space -> DiagGaussian OR Deterministic.
      if isinstance(action_space, Box):
          if action_space.dtype.char in np.typecodes["AllInteger"]:
              raise ValueError(
                  "Box(..., `int`) action spaces are not supported. "
                  "Use MultiDiscrete  or Box(..., `float`)."
              )
          else:
              if len(action_space.shape) > 1:
                  raise UnsupportedSpaceException(
                      f"Action space has multiple dimensions {action_space.shape}. "
                      f"Consider reshaping this into a single dimension, using a "
                      f"custom action distribution, using a Tuple action space, "
                      f"or the multi-agent API."
                  )
              return distribution_dicts[DistEnum.DiagGaussian]
  
      # Discrete Space -> Categorical.
      elif isinstance(action_space, Discrete):
          return distribution_dicts[DistEnum.Categorical]
  
      # Tuple/Dict Spaces -> MultiAction.
      elif isinstance(action_space, (Tuple, Dict)):
          return distribution_dicts[DistEnum.MultiDistribution]
  
      # Simplex -> Dirichlet.
      elif isinstance(action_space, Simplex):
          # TODO(Artur): Supported Simplex (in torch).
          raise NotImplementedError(
              "Simplex action space not yet supported."
          )
  
      # MultiDiscrete -> MultiCategorical.
      elif isinstance(action_space, MultiDiscrete):
          return distribution_dicts[DistEnum.MultiCategorical]
  
      # Unknown type -> Error.
      else:
          raise NotImplementedError(
              f"Unsupported action space: `{action_space}`"
          )
  
imports: import enum
  import functools
  from typing import Optional
  import gymnasium
  import numpy
  import tree
  from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
  from ray.rllib.core.models.base import Encoder
  from ray.rllib.core.models.configs import CNNEncoderConfig, MLPEncoderConfig, RecurrentEncoderConfig
  from ray.rllib.core.models.configs import ModelConfig
  from ray.rllib.models import MODEL_DEFAULTS
  from ray.rllib.models.distributions import Distribution
  from ray.rllib.models.preprocessors import get_preprocessor, Preprocessor
  from ray.rllib.models.utils import get_filter_config
  from ray.rllib.utils.deprecation import deprecation_warning
  from ray.rllib.utils.error import UnsupportedSpaceException
  from ray.rllib.utils.spaces.simplex import Simplex
  from ray.rllib.utils.spaces.space_utils import flatten_space
  from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space
  from ray.rllib.utils.typing import ViewRequirementsDict
  from ray.rllib.utils.annotations import OverrideToImplementCustomLogic, OverrideToImplementCustomLogic_CallToSuperRecommended
label_desc: Use image node_id 9 for calling the Catalog obj's underlying member method code with example usage: obj._get_dist_cls_from_action_space(cls, action_space) and returns: distribution_dicts, distribution_dicts, distribution_dicts, distribution_dicts