From c773824f4f22a14375a2451264d1b5708addd0f9 Mon Sep 17 00:00:00 2001
From: Tanay Wakhare <twakhare@gmail.com>
Date: Wed, 3 Jun 2020 13:06:06 -0400
Subject: [PATCH] [RLlib] Bug fixes and tests in DiagGaussian (#8676)

---
 rllib/BUILD                              |  2 +-
 rllib/models/tests/test_distributions.py | 75 +++++++++++++++++++++++-
 rllib/models/torch/torch_action_dist.py  |  2 +-
 3 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/rllib/BUILD b/rllib/BUILD
index 29aeead47..834093114 100644
--- a/rllib/BUILD
+++ b/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"]
 )
 
diff --git a/rllib/models/tests/test_distributions.py b/rllib/models/tests/test_distributions.py
index e4ce76a50..4531dd00f 100644
--- a/rllib/models/tests/test_distributions.py
+++ b/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
diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py
index 64cc07a9f..748012bf8 100644
--- a/rllib/models/torch/torch_action_dist.py
+++ b/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)