class_docstr: Tests ActionDistribution classes.
docstr: Tests the SquashedGaussian ActionDistribution for all frameworks.
docstr_tok: ["Tests","the","SquashedGaussian","ActionDistribution","for","all","frameworks","."]
code: def test_squashed_gaussian(self):
      """Tests the SquashedGaussian ActionDistribution for all frameworks."""
      input_space = Box(-2.0, 2.0, shape=(2000, 10))
      input_space.seed(42)
  
      low, high = -2.0, 1.0
  
      for fw, sess in framework_iterator(session=True):
          cls = (
              SquashedGaussian
              if fw != "torch"
              else TorchSquashedGaussian
          )
  
          # 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,
              bounds=(low, high),
          )
  
          # Batch of size=n and deterministic.
          inputs = input_space.sample()
          means, _ = np.split(inputs, 2, axis=-1)
          squashed_distribution = cls(inputs, {}, low=low, high=high)
          expected = ((np.tanh(means) + 1.0) / 2.0) * (high - low) + low
          # Sample n times, expect always mean value (deterministic draw).
          out = squashed_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)
          squashed_distribution = cls(inputs, {}, low=low, high=high)
          expected = ((np.tanh(means) + 1.0) / 2.0) * (high - low) + low
          values = squashed_distribution.sample()
          if sess:
              values = sess.run(values)
          else:
              values = values.numpy()
          self.assertTrue(np.max(values) <= high)
          self.assertTrue(np.min(values) >= low)
  
          check(np.mean(values), expected.mean(), decimals=1)
  
          # Test log-likelihood outputs.
          sampled_action_logp = squashed_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()
          # Convert to parameters for distr.
          stds = np.exp(
              np.clip(log_stds, MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT)
          )
          # Unsquash values, then get log-llh from regular gaussian.
          # atanh_in = np.clip((values - low) / (high - low) * 2.0 - 1.0,
          #   -1.0 + SMALL_NUMBER, 1.0 - SMALL_NUMBER)
          normed_values = (values - low) / (high - low) * 2.0 - 1.0
          save_normed_values = np.clip(
              normed_values, -1.0 + SMALL_NUMBER, 1.0 - SMALL_NUMBER
          )
          unsquashed_values = np.arctanh(save_normed_values)
          log_prob_unsquashed = np.sum(
              np.log(norm.pdf(unsquashed_values, means, stds)), -1
          )
          log_prob = log_prob_unsquashed - np.sum(
              np.log(1 - np.tanh(unsquashed_values) ** 2), axis=-1
          )
          check(
              np.sum(sampled_action_logp), np.sum(log_prob), rtol=0.05
          )
  
          # 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],
              ]
          )
          log_stds = np.array(
              [[0.8, -0.2, 0.3, -1.0, 2.0], [0.7, -0.3, 0.4, -0.9, 2.0]]
          )
          squashed_distribution = cls(
              inputs=np.concatenate([means, log_stds], axis=-1),
              model={},
              low=low,
              high=high,
          )
          # 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],
              ]
          )
  
          # Unsquash values, then get log-llh from regular gaussian.
          unsquashed_values = np.arctanh(
              (values - low) / (high - low) * 2.0 - 1.0
          )
          log_prob_unsquashed = np.sum(
              np.log(norm.pdf(unsquashed_values, means, stds)), -1
          )
          log_prob = log_prob_unsquashed - np.sum(
              np.log(1 - np.tanh(unsquashed_values) ** 2), axis=-1
          )
  
          outs = squashed_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 5 for calling the TestActionDistributions obj's underlying member method code with example usage: obj.test_squashed_gaussian() without return types