class_docstr: Tests ActionDistribution classes.
code: def test_multi_categorical(self):
      batch_size = 100
      num_categories = 3
      num_sub_distributions = 5
      # Create 5 categorical distributions of 3 categories each.
      inputs_space = Box(
          -1.0,
          2.0,
          shape=(batch_size, num_sub_distributions * num_categories),
      )
      inputs_space.seed(42)
      values_space = Box(
          0,
          num_categories - 1,
          shape=(num_sub_distributions, batch_size),
          dtype=np.int32,
      )
      values_space.seed(42)
  
      inputs = inputs_space.sample()
      input_lengths = [num_categories] * num_sub_distributions
      inputs_split = np.split(inputs, num_sub_distributions, axis=1)
  
      for fw, sess in framework_iterator(session=True):
          # Create the correct distribution object.
          cls = (
              MultiCategorical
              if fw != "torch"
              else TorchMultiCategorical
          )
          multi_categorical = cls(inputs, None, input_lengths)
  
          # Do a stability test using extreme NN outputs to see whether
          # sampling and logp'ing result in NaN or +/-inf values.
          self._stability_test(
              cls,
              inputs_space.shape,
              fw=fw,
              sess=sess,
              bounds=(0, num_categories - 1),
              extra_kwargs={"input_lens": input_lengths},
          )
  
          # Batch of size=3 and deterministic (True).
          expected = np.transpose(np.argmax(inputs_split, axis=-1))
          # Sample, expect always max value
          # (max likelihood for deterministic draw).
          out = multi_categorical.deterministic_sample()
          check(out, expected)
  
          # Batch of size=3 and non-deterministic -> expect roughly the mean.
          out = multi_categorical.sample()
          check(
              tf.reduce_mean(out)
              if fw != "torch"
              else torch.mean(out.float()),
              1.0,
              decimals=0,
          )
  
          # Test log-likelihood outputs.
          probs = softmax(inputs_split)
          values = values_space.sample()
  
          out = multi_categorical.logp(
              values
              if fw != "torch"
              else [
                  torch.Tensor(values[i])
                  for i in range(num_sub_distributions)
              ]
          )  # v in np.stack(values, 1)])
          expected = []
          for i in range(batch_size):
              expected.append(
                  np.sum(
                      np.log(
                          np.array(
                              [
                                  probs[j][i][values[j][i]]
                                  for j in range(num_sub_distributions)
                              ]
                          )
                      )
                  )
              )
          check(out, expected, decimals=4)
  
          # Test entropy outputs.
          out = multi_categorical.entropy()
          expected_entropy = -np.sum(
              np.sum(probs * np.log(probs), 0), -1
          )
          check(out, expected_entropy)
  
imports: from functools import partial
  from gymnasium.spaces import Box, Dict, Tuple
  import numpy
  from scipy.stats import beta, norm
  import tree
  import unittest
  from ray.rllib.models.tf.tf_action_dist import Beta, Categorical, DiagGaussian, GumbelSoftmax, MultiActionDistribution, MultiCategorical, SquashedGaussian
  from ray.rllib.models.torch.torch_action_dist import TorchBeta, TorchCategorical, TorchDiagGaussian, TorchMultiActionDistribution, TorchMultiCategorical, TorchSquashedGaussian
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, softmax, SMALL_NUMBER, LARGE_INTEGER
  from ray.rllib.utils.test_utils import check, framework_iterator
label_desc: Use image node_id 4 for calling the TestActionDistributions obj's underlying member method code with example usage: obj.test_multi_categorical() without return types