docstr: Tests if we can create a bunch of action distributions.
  
  Action distributions are created from the base catalog class. Things this
  test checks:
      - Whether we output the correct action distributions classes.
      - Whether we can instantiate the action distributions, query their
          required input dimensions and sample from them.
docstr_tok: ["Tests","if","we","can","create","a","bunch","of","action","distributions",".","Action","distributions","are","created","from","the","base","catalog","class",".","Things","this","test","checks",":","-","Whether","we","output","the","correct","action","distributions","classes",".","-","Whether","we","can","instantiate","the","action","distributions",",","query","their","required","input","dimensions","and","sample","from","them","."]
code: def test_get_dist_cls_from_action_space(self):
      """Tests if we can create a bunch of action distributions.
  
      Action distributions are created from the base catalog class. Things this
      test checks:
          - Whether we output the correct action distributions classes.
          - Whether we can instantiate the action distributions, query their
              required input dimensions and sample from them.
  
      """
      TestConfig = namedtuple(
          "TestConfig", ("action_space", "expected_dist_cls_dict")
      )
      test_configs = [
          # Box
          TestConfig(
              Box(-np.inf, np.inf, (7,), dtype=np.float32),
              {"torch": TorchDiagGaussian, "tf2": TfDiagGaussian},
          ),
          # Discrete
          TestConfig(
              Discrete(5),
              {"torch": TorchCategorical, "tf2": TfCategorical},
          ),
          # Nested Dict
          TestConfig(
              Dict(
                  {
                      "a": Box(-np.inf, np.inf, (7,), dtype=np.float32),
                      "b": Dict({"c": Discrete(5)}),
                  }
              ),
              {
                  "torch": TorchMultiDistribution,
                  "tf2": TfMultiDistribution,
              },
          ),
          # Nested Tuple
          TestConfig(
              Tuple(
                  (
                      Box(-np.inf, np.inf, (7,), dtype=np.float32),
                      Tuple((Discrete(5), Discrete(5))),
                  )
              ),
              {
                  "torch": TorchMultiDistribution,
                  "tf2": TfMultiDistribution,
              },
          ),
          # Tuple nested inside Dict
          TestConfig(
              Dict(
                  {
                      "a": Box(-np.inf, np.inf, (7,), dtype=np.float32),
                      "b": Dict(
                          {
                              "c": Tuple(
                                  (
                                      Box(
                                          -np.inf,
                                          np.inf,
                                          (7,),
                                          dtype=np.float32,
                                      ),
                                      Tuple((Discrete(5), Discrete(5))),
                                  )
                              )
                          }
                      ),
                  }
              ),
              {
                  "torch": TorchMultiDistribution,
                  "tf2": TfMultiDistribution,
              },
          ),
          # Dict nested inside Tuple
          TestConfig(
              Tuple(
                  (
                      Box(-np.inf, np.inf, (7,), dtype=np.float32),
                      Tuple(
                          (
                              Discrete(5),
                              Dict(
                                  {
                                      "a": Box(
                                          -np.inf,
                                          np.inf,
                                          (7,),
                                          dtype=np.float32,
                                      ),
                                      "b": Dict({"c": Discrete(5)}),
                                  }
                              ),
                          )
                      ),
                  )
              ),
              {
                  "torch": TorchMultiDistribution,
                  "tf2": TfMultiDistribution,
              },
          ),
          # MultiDiscrete
          TestConfig(
              MultiDiscrete([5, 5, 5]),
              {
                  "torch": TorchMultiCategorical,
                  "tf2": TfMultiCategorical,
              },
          ),
      ]
  
      for (
          action_space,
          expected_cls_dict,
      ) in test_configs:
          print(f"Testing action space: {action_space}:")
          catalog = Catalog(
              observation_space=Box(
                  -1.0, 1.0, (84, 84, 1), dtype=np.float32
              ),
              action_space=action_space,
              model_config_dict=MODEL_DEFAULTS.copy(),
          )
  
          for framework in framework_iterator(
              frameworks=["tf2", "torch"]
          ):
              if framework == "tf2":
                  framework = "tf2"
  
              dist_cls = catalog._get_dist_cls_from_action_space(
                  action_space=action_space,
                  framework=framework,
              )
  
              # Check if we can query the required input dimensions
              expected_cls = expected_cls_dict[framework]
              if (
                  expected_cls is TorchMultiDistribution
                  or expected_cls is TfMultiDistribution
              ):
                  # For these special cases, we need to create partials of the
                  # expected classes so that we can calculate the required inputs
                  expected_cls = _multi_action_dist_partial_helper(
                      catalog_cls=catalog,
                      action_space=action_space,
                      framework=framework,
                  )
              elif (
                  expected_cls is TorchMultiCategorical
                  or expected_cls is TfMultiCategorical
              ):
                  # For these special cases, we need to create partials of the
                  # expected classes so that we can calculate the required inputs
                  expected_cls = _multi_categorical_dist_partial_helper(
                      action_space=action_space, framework=framework
                  )
  
              # Now that we have sorted out special cases, we can finally get the
              # input_dim
              input_dim = expected_cls.required_input_dim(action_space)
              logits = np.ones((32, input_dim), dtype=np.float32)
              if framework == "torch":
                  logits = torch.from_numpy(logits)
              else:
                  logits = tf.convert_to_tensor(logits)
              # We don't need a model if we input tensors
              dist = dist_cls.from_logits(logits=logits)
              self.assertTrue(
                  isinstance(dist, expected_cls_dict[framework]),
                  msg=f"Expected {expected_cls_dict[framework]}, "
                  f"got {type(dist)}",
              )
              # Test if sampling works
              actions = dist.sample()
              # Test is logp works
              dist.logp(actions)
  
              # For any array of actions in a possibly nested space, convert to
              # numpy and pick the first one to check if it is in the action space.
              action = tree.map_structure(
                  lambda a: convert_to_numpy(a)[0], actions
              )
              self.assertTrue(action_space.contains(action))
  
imports: from collections import namedtuple
  import functools
  import itertools
  import unittest
  import gymnasium
  from gymnasium.spaces import Box, Discrete, Dict, Tuple, MultiDiscrete
  import numpy
  import tree
  from ray.rllib.algorithms.ppo.ppo import PPOConfig
  from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
  from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
  from ray.rllib.core.models.base import Encoder, STATE_IN, ENCODER_OUT, STATE_OUT
  from ray.rllib.core.models.catalog import Catalog, _multi_action_dist_partial_helper, _multi_categorical_dist_partial_helper
  from ray.rllib.core.models.configs import MLPEncoderConfig, ModelConfig, CNNEncoderConfig
  from ray.rllib.core.models.torch.base import TorchModel
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
  from ray.rllib.models import MODEL_DEFAULTS
  from ray.rllib.models.tf.tf_distributions import TfCategorical, TfDiagGaussian, TfMultiCategorical, TfMultiDistribution
  from ray.rllib.models.torch.torch_distributions import TorchCategorical, TorchDiagGaussian, TorchMultiCategorical, TorchMultiDistribution
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.spaces.space_utils import get_dummy_batch_for_space
  from ray.rllib.utils.test_utils import framework_iterator
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
label_desc: Use image node_id 3 for calling the TestCatalog obj's underlying member method code with example usage: obj.test_get_dist_cls_from_action_space() without return types