Add simplex action space and dirichlet action distribution (#4070)

* add simplex action space and dirichlet action distribution

* Update and rename spaces.py to extra_spaces.py

* Update __init__.py

* Update catalog.py

* Fix python 2

* Update extra_spaces.py

* change Simplex.contains() to return False

python/ray/rllib/models/__init__.py

@@ -1,6 +1,7 @@
 from ray.rllib.models.catalog import ModelCatalog, MODEL_DEFAULTS
-from ray.rllib.models.action_dist import (ActionDistribution, Categorical,
-                                          DiagGaussian, Deterministic)
+from ray.rllib.models.extra_spaces import Simplex
+from ray.rllib.models.action_dist import (
+    ActionDistribution, Categorical, DiagGaussian, Deterministic, Dirichlet)
 from ray.rllib.models.model import Model
 from ray.rllib.models.preprocessors import Preprocessor
 from ray.rllib.models.fcnet import FullyConnectedNetwork
@@ -11,10 +12,12 @@ __all__ = [
     "Categorical",
     "DiagGaussian",
     "Deterministic",
+    "Dirichlet",
     "ModelCatalog",
     "Model",
     "Preprocessor",
     "FullyConnectedNetwork",
     "LSTM",
     "MODEL_DEFAULTS",
+    "Simplex",
 ]

python/ray/rllib/models/action_dist.py

@@ -233,3 +233,30 @@ class MultiActionDistribution(ActionDistribution):
 
 
 TupleActions = namedtuple("TupleActions", ["batches"])
+
+
+class Dirichlet(ActionDistribution):
+    """Dirichlet distribution for countinuous actions that are between
+    [0,1] and sum to 1.
+
+    e.g. actions that represent resource allocation."""
+
+    def __init__(self, inputs):
+        self.dist = tf.distributions.Dirichlet(concentration=inputs)
+        ActionDistribution.__init__(self, inputs)
+
+    @override(ActionDistribution)
+    def logp(self, x):
+        return self.dist.log_prob(x)
+
+    @override(ActionDistribution)
+    def entropy(self):
+        return self.dist.entropy()
+
+    @override(ActionDistribution)
+    def kl(self, other):
+        return self.dist.kl_divergence(other.dist)
+
+    @override(ActionDistribution)
+    def _build_sample_op(self):
+        return self.dist.sample()

python/ray/rllib/models/catalog.py

@@ -11,8 +11,10 @@ from functools import partial
 from ray.tune.registry import RLLIB_MODEL, RLLIB_PREPROCESSOR, \
     _global_registry
 
-from ray.rllib.models.action_dist import (
-    Categorical, Deterministic, DiagGaussian, MultiActionDistribution)
+from ray.rllib.models.extra_spaces import Simplex
+from ray.rllib.models.action_dist import (Categorical, Deterministic,
+                                          DiagGaussian,
+                                          MultiActionDistribution, Dirichlet)
 from ray.rllib.models.preprocessors import get_preprocessor
 from ray.rllib.models.fcnet import FullyConnectedNetwork
 from ray.rllib.models.visionnet import VisionNetwork
@@ -132,7 +134,8 @@ class ModelCatalog(object):
                 child_distributions=child_dist,
                 action_space=action_space,
                 input_lens=input_lens), sum(input_lens)
-
+        elif isinstance(action_space, Simplex):
+            return Dirichlet, action_space.shape[0]
         raise NotImplementedError("Unsupported args: {} {}".format(
             action_space, dist_type))
 
@@ -165,6 +168,9 @@ class ModelCatalog(object):
                 tf.int64 if all_discrete else tf.float32,
                 shape=(None, size),
                 name="action")
+        elif isinstance(action_space, Simplex):
+            return tf.placeholder(
+                tf.float32, shape=(None, action_space.shape[0]), name="action")
         else:
             raise NotImplementedError("action space {}"
                                       " not supported".format(action_space))

python/ray/rllib/models/extra_spaces.py

@@ -0,0 +1,57 @@
+import numpy as np
+import gym
+
+
+class Simplex(gym.Space):
+    """Represents a d - 1 dimensional Simplex in R^d.
+
+    That is, all coordinates are in [0, 1] and sum to 1.
+    The dimension d of the simplex is assumed to be shape[-1].
+
+    Additionally one can specify the underlying distribution of
+    the simplex as a Dirichlet distribution by providing concentration
+    parameters. By default, sampling is uniform, i.e. concentration is
+    all 1s.
+
+    Example usage:
+    self.action_space = spaces.Simplex(shape=(3, 4))
+        --> 3 independent 4d Dirichlet with uniform concentration
+    """
+
+    def __init__(self, shape, concentration=None, dtype=np.float32):
+        assert type(shape) in [tuple, list]
+        self.shape = shape
+        self.dtype = dtype
+        self.dim = shape[-1]
+
+        if concentration is not None:
+            assert concentration.shape == shape[:-1]
+        else:
+            self.concentration = [1] * self.dim
+
+        super().__init__(shape, dtype)
+        self.np_random = np.random.RandomState()
+
+    def seed(self, seed):
+        self.np_random.seed(seed)
+
+    def sample(self):
+        return np.random.dirichlet(
+            self.concentration, size=self.shape[:-1]).astype(self.dtype)
+
+    def contains(self, x):
+        return x.shape == self.shape and np.allclose(
+            np.sum(x, axis=-1), np.ones_like(x[..., 0]))
+
+    def to_jsonable(self, sample_n):
+        return np.array(sample_n).tolist()
+
+    def from_jsonable(self, sample_n):
+        return [np.asarray(sample) for sample in sample_n]
+
+    def __repr__(self):
+        return "Simplex({}; {})".format(self.shape, self.concentration)
+
+    def __eq__(self, other):
+        return np.allclose(self.concentration,
+                           other.concentration) and self.shape == other.shape