From 658e002cdfc32527cb112b5fe00fea0664c27a10 Mon Sep 17 00:00:00 2001 From: jon-chuang <9093549+jon-chuang@users.noreply.github.com> Date: Mon, 19 Aug 2019 09:34:50 +0300 Subject: [PATCH] [rllib] Add autoregressive KL (#5469) --- rllib/examples/autoregressive_action_dist.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/rllib/examples/autoregressive_action_dist.py b/rllib/examples/autoregressive_action_dist.py index 819341594..5a6c0fa86 100644 --- a/rllib/examples/autoregressive_action_dist.py +++ b/rllib/examples/autoregressive_action_dist.py @@ -99,8 +99,12 @@ class BinaryAutoregressiveOutput(ActionDistribution): return a1_dist.entropy() + a2_dist.entropy() def kl(self, other): - # TODO: implement this properly - return tf.zeros_like(self.entropy()) + a1_dist = self._a1_distribution() + a1_terms = a1_dist.kl(other._a1_distribution()) + + a1 = a1_dist.sample() + a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1)) + return a1_terms + a2_terms def _a1_distribution(self): BATCH = tf.shape(self.inputs)[0]