[wingman -> rllib] IMPALA MultiDiscrete changes (#3967)

python/ray/rllib/models/action_dist.py

@@ -114,6 +114,31 @@ class Categorical(ActionDistribution):
         return tf.squeeze(tf.multinomial(self.inputs, 1), axis=1)
 
 
+class MultiCategorical(ActionDistribution):
+    """Categorical distribution for discrete action spaces."""
+
+    def __init__(self, inputs):
+        self.cats = [Categorical(input_) for input_ in inputs]
+        self.sample_op = self._build_sample_op()
+
+    def logp(self, actions):
+        # If tensor is provided, unstack it into list
+        if isinstance(actions, tf.Tensor):
+            actions = tf.unstack(actions, axis=1)
+        logps = tf.stack(
+            [cat.logp(act) for cat, act in zip(self.cats, actions)])
+        return tf.reduce_sum(logps, axis=0)
+
+    def entropy(self):
+        return tf.stack([cat.entropy() for cat in self.cats], axis=1)
+
+    def kl(self, other):
+        return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)]
+
+    def _build_sample_op(self):
+        return tf.stack([cat.sample() for cat in self.cats], axis=1)
+
+
 class DiagGaussian(ActionDistribution):
     """Action distribution where each vector element is a gaussian.