class_docstr: Tests ActionDistribution classes.
docstr: Tests the DiagGaussian ActionDistribution for all frameworks.
docstr_tok: ["Tests","the","DiagGaussian","ActionDistribution","for","all","frameworks","."]
code: def test_diag_gaussian(self):
      """Tests the DiagGaussian ActionDistribution for all frameworks."""
      input_space = Box(-2.0, 1.0, shape=(2000, 10))
      input_space.seed(42)
  
      for fw, sess in framework_iterator(session=True):
          cls = DiagGaussian if fw != "torch" else TorchDiagGaussian
  
          # 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, input_space.shape, fw=fw, sess=sess)
  
          # Batch of size=n and deterministic.
          inputs = input_space.sample()
          means, _ = np.split(inputs, 2, axis=-1)
          diag_distribution = cls(inputs, {})
          expected = means
          # Sample n times, expect always mean value (deterministic draw).
          out = diag_distribution.deterministic_sample()
          check(out, expected)
  
          # Batch of size=n and non-deterministic -> expect roughly the mean.
          inputs = input_space.sample()
          means, log_stds = np.split(inputs, 2, axis=-1)
          diag_distribution = cls(inputs, {})
          expected = means
          values = diag_distribution.sample()
          if sess:
              values = sess.run(values)
          else:
              values = values.numpy()
          check(np.mean(values), expected.mean(), decimals=1)
  
          # Test log-likelihood outputs.
          sampled_action_logp = diag_distribution.logp(
              values if fw != "torch" else torch.Tensor(values)
          )
          if sess:
              sampled_action_logp = sess.run(sampled_action_logp)
          else:
              sampled_action_logp = sampled_action_logp.numpy()
  
          # NN output.
          means = np.array(
              [
                  [0.1, 0.2, 0.3, 0.4, 50.0],
                  [-0.1, -0.2, -0.3, -0.4, -1.0],
              ],
              dtype=np.float32,
          )
          log_stds = np.array(
              [
                  [0.8, -0.2, 0.3, -1.0, 2.0],
                  [0.7, -0.3, 0.4, -0.9, 2.0],
              ],
              dtype=np.float32,
          )
  
          diag_distribution = cls(
              inputs=np.concatenate([means, log_stds], axis=-1),
              model={},
          )
          # Convert to parameters for distr.
          stds = np.exp(log_stds)
          # Values to get log-likelihoods for.
          values = np.array(
              [
                  [0.9, 0.2, 0.4, -0.1, -1.05],
                  [-0.9, -0.2, 0.4, -0.1, -1.05],
              ]
          )
  
          # get log-llh from regular gaussian.
          log_prob = np.sum(np.log(norm.pdf(values, means, stds)), -1)
  
          outs = diag_distribution.logp(
              values if fw != "torch" else torch.Tensor(values)
          )
          if sess:
              outs = sess.run(outs)
          check(outs, log_prob, decimals=4)
  
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 6 for calling the TestActionDistributions obj's underlying member method code with example usage: obj.test_diag_gaussian() without return types