[rllib] MultiCategorical shouldn't return array for kl or entropy (ray-project#5215)

* wip

* fix

Author: ericl

ci/travis/format.sh

@@ -29,10 +29,6 @@ YAPF_VERSION=$(yapf --version | awk '{print $2}')
 tool_version_check() {
     if [[ $2 != $3 ]]; then
         echo "WARNING: Ray uses $1 $3, You currently are using $2. This might generate different results."
-        read -p "Do you want to continue? [y/n] " answer
-        if ! [ $answer = 'y' ] && ! [ $answer = 'Y' ]; then
-            exit 1
-        fi
     fi
 }
 

python/ray/rllib/agents/impala/vtrace_policy.py

@@ -184,7 +184,8 @@ def make_time_major(*args, **kw):
         actions=make_time_major(loss_actions, drop_last=True),
         actions_logp=make_time_major(
             action_dist.logp(actions), drop_last=True),
-        actions_entropy=make_time_major(action_dist.entropy(), drop_last=True),
+        actions_entropy=make_time_major(
+            action_dist.multi_entropy(), drop_last=True),
         dones=make_time_major(dones, drop_last=True),
         behaviour_logits=make_time_major(
             unpacked_behaviour_logits, drop_last=True),

python/ray/rllib/agents/ppo/appo_policy.py

@@ -205,9 +205,9 @@ def make_time_major(*args, **kw):
                 prev_action_dist.logp(actions), drop_last=True),
             actions_logp=make_time_major(
                 action_dist.logp(actions), drop_last=True),
-            action_kl=prev_action_dist.kl(action_dist),
+            action_kl=prev_action_dist.multi_kl(action_dist),
             actions_entropy=make_time_major(
-                action_dist.entropy(), drop_last=True),
+                action_dist.multi_entropy(), drop_last=True),
             dones=make_time_major(dones, drop_last=True),
             behaviour_logits=make_time_major(
                 unpacked_behaviour_logits, drop_last=True),
@@ -229,8 +229,8 @@ def make_time_major(*args, **kw):
         policy.loss = PPOSurrogateLoss(
             prev_actions_logp=make_time_major(prev_action_dist.logp(actions)),
             actions_logp=make_time_major(action_dist.logp(actions)),
-            action_kl=prev_action_dist.kl(action_dist),
-            actions_entropy=make_time_major(action_dist.entropy()),
+            action_kl=prev_action_dist.multi_kl(action_dist),
+            actions_entropy=make_time_major(action_dist.multi_entropy()),
             values=make_time_major(values),
             valid_mask=make_time_major(mask),
             advantages=make_time_major(

python/ray/rllib/models/action_dist.py

@@ -69,6 +69,22 @@ def sampled_action_prob(self):
         """Returns the log probability of the sampled action."""
         return tf.exp(self.logp(self.sample_op))
 
+    def multi_kl(self, other):
+        """The KL-divergence between two action distributions.
+
+        This differs from kl() in that it can return an array for
+        MultiDiscrete. TODO(ekl) consider removing this.
+        """
+        return self.kl(other)
+
+    def multi_entropy(self):
+        """The entropy of the action distribution.
+
+        This differs from entropy() in that it can return an array for
+        MultiDiscrete. TODO(ekl) consider removing this.
+        """
+        return self.entropy()
+
 
 class Categorical(ActionDistribution):
     """Categorical distribution for discrete action spaces."""
@@ -133,6 +149,7 @@ def __init__(self, inputs, input_lens):
         ]
         self.sample_op = self._build_sample_op()
 
+    @override(ActionDistribution)
     def logp(self, actions):
         # If tensor is provided, unstack it into list
         if isinstance(actions, tf.Tensor):
@@ -141,12 +158,23 @@ def logp(self, actions):
             [cat.logp(act) for cat, act in zip(self.cats, actions)])
         return tf.reduce_sum(logps, axis=0)
 
-    def entropy(self):
+    @override(ActionDistribution)
+    def multi_entropy(self):
         return tf.stack([cat.entropy() for cat in self.cats], axis=1)
 
-    def kl(self, other):
+    @override(ActionDistribution)
+    def entropy(self):
+        return tf.reduce_sum(self.multi_entropy(), axis=1)
+
+    @override(ActionDistribution)
+    def multi_kl(self, other):
         return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)]
 
+    @override(ActionDistribution)
+    def kl(self, other):
+        return tf.reduce_sum(self.multi_kl(other), axis=1)
+
+    @override(ActionDistribution)
     def _build_sample_op(self):
         return tf.stack([cat.sample() for cat in self.cats], axis=1)