Test for differential entropy

tests/test_distributions.py

@@ -1,3 +1,4 @@
+import pytest
 import numpy as np
 import torch as th
 
@@ -5,6 +6,7 @@ from torchy_baselines.common.utils import set_random_seed
 from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution,\
     CategoricalDistribution, TanhBijector, StateDependentNoiseDistribution
 
+
 # TODO: more tests for the other distributions
 def test_bijector():
     """
@@ -37,3 +39,32 @@ def test_sde_distribution():
 
     assert th.allclose(actions.mean(), dist.distribution.mean.mean(), rtol=1e-3)
     assert th.allclose(actions.std(), dist.distribution.scale.mean(), rtol=1e-3)
+
+
+N_ACTIONS = 1
+
+# TODO: fix for num action > 1
+# TODO: analytical form for squashed Gaussian?
+@pytest.mark.parametrize("dist", [
+    DiagGaussianDistribution(N_ACTIONS),
+    StateDependentNoiseDistribution(N_ACTIONS, squash_output=False),
+])
+def test_entropy(dist):
+    # The entropy can be approximated by averaging the negative log likelihood
+    # mean negative log likelihood == differential entropy
+    n_samples = int(5e6)
+    n_features = 3
+    set_random_seed(1)
+    state = th.rand(n_samples, n_features)
+    deterministic_actions = th.rand(n_samples, N_ACTIONS)
+    _, log_std = dist.proba_distribution_net(n_features, log_std_init=th.log(th.tensor(0.2)))
+
+    if isinstance(dist, DiagGaussianDistribution):
+        actions, dist = dist.proba_distribution(deterministic_actions, log_std)
+    else:
+        dist.sample_weights(log_std, batch_size=n_samples)
+        actions, dist = dist.proba_distribution(deterministic_actions, log_std, state)
+
+    entropy = dist.entropy()
+    log_prob = dist.log_prob(actions)
+    assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3)