Torch multicat support (7419)

rllib/models/catalog.py

@@ -8,8 +8,8 @@ from ray.tune.registry import RLLIB_MODEL, RLLIB_PREPROCESSOR, \
 
 from ray.rllib.models.extra_spaces import Simplex
 from ray.rllib.models.action_dist import ActionDistribution
-from ray.rllib.models.torch.torch_action_dist import (TorchCategorical,
-                                                      TorchDiagGaussian)
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \
+    TorchMultiCategorical, TorchDiagGaussian
 from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork as FCNetV2
 from ray.rllib.models.tf.visionnet_v2 import VisionNetwork as VisionNetV2
 from ray.rllib.models.tf.tf_action_dist import Categorical, MultiCategorical, \
@@ -183,11 +183,9 @@ class ModelCatalog:
             dist = Dirichlet
         # MultiDiscrete -> MultiCategorical.
         elif isinstance(action_space, gym.spaces.MultiDiscrete):
-            if framework == "torch":
-                # TODO(sven): implement
-                raise NotImplementedError(
-                    "MultiDiscrete action spaces not supported for Pytorch.")
-            return partial(MultiCategorical, input_lens=action_space.nvec), \
+            dist = MultiCategorical if framework == "tf" else \
+                TorchMultiCategorical
+            return partial(dist, input_lens=action_space.nvec), \
                 int(sum(action_space.nvec))
         # Dict -> TODO(sven)
         elif isinstance(action_space, gym.spaces.Dict):

rllib/models/tests/test_distributions.py

@@ -4,12 +4,15 @@ from scipy.stats import norm
 from tensorflow.python.eager.context import eager_mode
 import unittest
 
-from ray.rllib.models.tf.tf_action_dist import Categorical, SquashedGaussian
-from ray.rllib.utils import try_import_tf
-from ray.rllib.utils.numpy import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT
+from ray.rllib.models.tf.tf_action_dist import Categorical, MultiCategorical, \
+    SquashedGaussian
+from ray.rllib.models.torch.torch_action_dist import TorchMultiCategorical
+from ray.rllib.utils import try_import_tf, try_import_torch
+from ray.rllib.utils.numpy import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, softmax
 from ray.rllib.utils.test_utils import check
 
 tf = try_import_tf()
+torch, _ = try_import_torch()
 
 
 class TestDistributions(unittest.TestCase):
@@ -32,6 +35,70 @@ class TestDistributions(unittest.TestCase):
         probs = np.exp(z) / np.sum(np.exp(z))
         self.assertTrue(np.sum(np.abs(probs - counts / num_samples)) <= 0.01)
 
+    def test_multi_categorical(self):
+        batch_size = 100
+        num_categories = 3
+        num_sub_distributions = 5
+        # Create 5 categorical distributions of 3 categories each.
+        inputs_space = Box(
+            -1.0,
+            2.0,
+            shape=(batch_size, num_sub_distributions * num_categories))
+        values_space = Box(
+            0,
+            num_categories - 1,
+            shape=(num_sub_distributions, batch_size),
+            dtype=np.int32)
+
+        # The Component to test.
+        inputs = inputs_space.sample()
+        input_lengths = [num_categories] * num_sub_distributions
+        inputs_split = np.split(inputs, num_sub_distributions, axis=1)
+
+        for fw in ["tf", "eager", "torch"]:
+            print("framework={}".format(fw))
+
+            cls = MultiCategorical if fw != "torch" else TorchMultiCategorical
+            multi_categorical = cls(inputs, None, input_lengths)
+
+            # Batch of size=3 and deterministic (True).
+            expected = np.transpose(np.argmax(inputs_split, axis=-1))
+            # Sample, expect always max value
+            # (max likelihood for deterministic draw).
+            out = multi_categorical.deterministic_sample()
+            check(out, expected)
+
+            # Batch of size=3 and non-deterministic -> expect roughly the mean.
+            out = multi_categorical.sample()
+            check(
+                tf.reduce_mean(out)
+                if fw != "torch" else torch.mean(out.float()),
+                1.0,
+                decimals=0)
+
+            # Test log-likelihood outputs.
+            probs = softmax(inputs_split)
+            values = values_space.sample()
+
+            out = multi_categorical.logp(values if fw != "torch" else [
+                torch.Tensor(values[i]) for i in range(num_sub_distributions)
+            ])  # v in np.stack(values, 1)])
+            expected = []
+            for i in range(batch_size):
+                expected.append(
+                    np.sum(
+                        np.log(
+                            np.array([
+                                probs[j][i][values[j][i]]
+                                for j in range(num_sub_distributions)
+                            ]))))
+            check(out, expected, decimals=4)
+
+            # Test entropy outputs.
+            out = multi_categorical.entropy()
+            expected_entropy = -np.sum(np.sum(probs * np.log(probs), 0), -1)
+            check(out, expected_entropy)
+
     def test_squashed_gaussian(self):
         """Tests the SquashedGaussia ActionDistribution (tf-eager only)."""
         with eager_mode():

rllib/models/tf/tf_action_dist.py

@@ -102,11 +102,13 @@ class MultiCategorical(TFActionDistribution):
 
     @override(ActionDistribution)
     def deterministic_sample(self):
-        return tf.math.argmax(self.inputs, axis=-1)
+        return tf.stack(
+            [cat.deterministic_sample() for cat in self.cats],
+            axis=1)
 
     @override(ActionDistribution)
     def logp(self, actions):
-        # If tensor is provided, unstack it into list
+        # If tensor is provided, unstack it into list.
         if isinstance(actions, tf.Tensor):
             actions = tf.unstack(tf.cast(actions, tf.int32), axis=1)
         logps = tf.stack(

rllib/models/torch/torch_action_dist.py

@@ -10,8 +10,10 @@ torch, nn = try_import_torch()
 class TorchDistributionWrapper(ActionDistribution):
     """Wrapper class for torch.distributions."""
 
-    def __init_(self, inputs):
-        super().__init__(inputs)
+    @override(ActionDistribution)
+    def __init__(self, inputs, model):
+        inputs = torch.Tensor(inputs)
+        super().__init__(inputs, model)
         # Store the last sample here.
         self.last_sample = None
 
@@ -56,6 +58,67 @@ class TorchCategorical(TorchDistributionWrapper):
         return action_space.n
 
 
+class TorchMultiCategorical(TorchDistributionWrapper):
+    """MultiCategorical distribution for MultiDiscrete action spaces."""
+
+    @override(TorchDistributionWrapper)
+    def __init__(self, inputs, model, input_lens):
+        super().__init__(inputs, model)
+        inputs_split = self.inputs.split(input_lens, dim=1)
+        self.cats = [
+            torch.distributions.categorical.Categorical(logits=input_)
+            for input_ in inputs_split
+        ]
+
+    @override(TorchDistributionWrapper)
+    def sample(self):
+        arr = [cat.sample() for cat in self.cats]
+        ret = torch.stack(arr, dim=1)
+        return ret
+
+    @override(ActionDistribution)
+    def deterministic_sample(self):
+        arr = [torch.argmax(cat.probs, -1) for cat in self.cats]
+        ret = torch.stack(arr, dim=1)
+        return ret
+
+    @override(TorchDistributionWrapper)
+    def logp(self, actions):
+        # # If tensor is provided, unstack it into list.
+        if isinstance(actions, torch.Tensor):
+            actions = torch.unbind(actions, dim=1)
+        logps = torch.stack(
+            [cat.log_prob(act) for cat, act in zip(self.cats, actions)])
+        return torch.sum(logps, dim=0)
+
+    @override(ActionDistribution)
+    def multi_entropy(self):
+        return torch.stack([cat.entropy() for cat in self.cats], dim=1)
+
+    @override(TorchDistributionWrapper)
+    def entropy(self):
+        return torch.sum(self.multi_entropy(), dim=1)
+
+    @override(ActionDistribution)
+    def multi_kl(self, other):
+        return torch.stack(
+            [
+                torch.distributions.kl.kl_divergence(cat, oth_cat)
+                for cat, oth_cat in zip(self.cats, other.cats)
+            ],
+            dim=1,
+        )
+
+    @override(TorchDistributionWrapper)
+    def kl(self, other):
+        return torch.sum(self.multi_kl(other), dim=1)
+
+    @staticmethod
+    @override(ActionDistribution)
+    def required_model_output_shape(action_space, model_config):
+        return np.sum(action_space.nvec)
+
+
 class TorchDiagGaussian(TorchDistributionWrapper):
     """Wrapper class for PyTorch Normal distribution."""