From 4198db5038dd465a5d5db91b0ca84d98391aec2d Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Wed, 4 Mar 2020 09:41:40 +0100 Subject: [PATCH] Torch multicat support (7419) --- rllib/models/catalog.py | 12 ++-- rllib/models/tests/test_distributions.py | 73 +++++++++++++++++++++++- rllib/models/tf/tf_action_dist.py | 6 +- rllib/models/torch/torch_action_dist.py | 67 +++++++++++++++++++++- 4 files changed, 144 insertions(+), 14 deletions(-) diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index 881f5299d..71daab5dc 100644 --- a/rllib/models/catalog.py +++ b/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): diff --git a/rllib/models/tests/test_distributions.py b/rllib/models/tests/test_distributions.py index 531111aff..3d7f3adf1 100644 --- a/rllib/models/tests/test_distributions.py +++ b/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(): diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index a8c6337ed..c6080fc52 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/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( diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index fd6717240..e63ea8beb 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/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."""