[RLlib] Bug fixes and tests in DiagGaussian (#8676)

rllib/BUILD

@@ -1016,7 +1016,7 @@ sh_test(
 py_test(
     name = "test_distributions",
     tags = ["models"],
-    size = "small",
+    size = "medium",
     srcs = ["models/tests/test_distributions.py"]
 )
 

rllib/models/tests/test_distributions.py

@@ -43,8 +43,15 @@ class TestDistributions(unittest.TestCase):
         ]
         inputs = np.zeros(shape=network_output_shape, dtype=np.float32)
         for batch_item in range(network_output_shape[0]):
-            for num in range(len(inputs[batch_item])):
+            for num in range(len(inputs[batch_item]) // 2):
                 inputs[batch_item][num] = np.random.choice(extreme_values)
+            else:
+                # For Gaussians, the second half of the vector is
+                # log standard deviations, and should therefore be
+                # the log of a positive number >= 1.
+                inputs[batch_item][num] = np.log(
+                    max(1, np.random.choice((extreme_values))))
+
         dist = distribution_cls(inputs, {})
         for _ in range(100):
             sample = dist.sample()
@@ -237,6 +244,72 @@ class TestDistributions(unittest.TestCase):
                 outs = sess.run(outs)
             check(outs, log_prob, decimals=4)
 
+    def test_diag_gaussian(self):
+        """Tests the DiagGaussian ActionDistribution for all frameworks."""
+        input_space = Box(-2.0, 2.0, shape=(200, 10))
+
+        for fw, sess in framework_iterator(
+                frameworks=("torch", "tf", "tfe"), 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)
+
     def test_beta(self):
         input_space = Box(-2.0, 1.0, shape=(200, 10))
         low, high = -1.0, 2.0

rllib/models/torch/torch_action_dist.py

@@ -141,7 +141,7 @@ class TorchDiagGaussian(TorchDistributionWrapper):
     @override(ActionDistribution)
     def __init__(self, inputs, model):
         super().__init__(inputs, model)
-        mean, log_std = torch.chunk(inputs, 2, dim=1)
+        mean, log_std = torch.chunk(self.inputs, 2, dim=1)
         self.dist = torch.distributions.normal.Normal(mean, torch.exp(log_std))
 
     @override(ActionDistribution)