From a708ab66f5ccb81cb12b036b159b00073738781a Mon Sep 17 00:00:00 2001 From: Zekun Shi Date: Sun, 17 Feb 2019 04:44:59 +0800 Subject: [PATCH] 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 | 7 ++- python/ray/rllib/models/action_dist.py | 27 ++++++++++++ python/ray/rllib/models/catalog.py | 12 ++++-- python/ray/rllib/models/extra_spaces.py | 57 +++++++++++++++++++++++++ 4 files changed, 98 insertions(+), 5 deletions(-) create mode 100644 python/ray/rllib/models/extra_spaces.py diff --git a/python/ray/rllib/models/__init__.py b/python/ray/rllib/models/__init__.py index 52e47e807..8ddbc6ed8 100644 --- a/python/ray/rllib/models/__init__.py +++ b/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", ] diff --git a/python/ray/rllib/models/action_dist.py b/python/ray/rllib/models/action_dist.py index 5de4ea225..724e54fd1 100644 --- a/python/ray/rllib/models/action_dist.py +++ b/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() diff --git a/python/ray/rllib/models/catalog.py b/python/ray/rllib/models/catalog.py index 474a0e905..5dd650a03 100644 --- a/python/ray/rllib/models/catalog.py +++ b/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)) diff --git a/python/ray/rllib/models/extra_spaces.py b/python/ray/rllib/models/extra_spaces.py new file mode 100644 index 000000000..44ed9adcf --- /dev/null +++ b/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