From 07345e5e2752de58ff694257641e955246935180 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 18 Dec 2019 13:45:56 +0100 Subject: [PATCH] Test for differential entropy --- tests/test_distributions.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 130a1bc..3d896ab 100644 --- a/tests/test_distributions.py +++ b/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)