From eea75ac62370e404c3505e69e03709050be8b6b6 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Thu, 30 Apr 2020 20:09:33 +0200 Subject: [PATCH] [RLlib] Beta distribution. (#8229) --- rllib/agents/sac/sac_tf_policy.py | 17 ++-- rllib/models/tests/test_distributions.py | 105 +++++++++++++++++------ rllib/models/tf/tf_action_dist.py | 46 ++++++++++ rllib/tests/test_supported_spaces.py | 5 +- 4 files changed, 139 insertions(+), 34 deletions(-) diff --git a/rllib/agents/sac/sac_tf_policy.py b/rllib/agents/sac/sac_tf_policy.py index 4a4b68b24..63ba0946d 100644 --- a/rllib/agents/sac/sac_tf_policy.py +++ b/rllib/agents/sac/sac_tf_policy.py @@ -9,8 +9,8 @@ from ray.rllib.agents.dqn.dqn_tf_policy import postprocess_nstep_and_prio from ray.rllib.agents.sac.sac_tf_model import SACTFModel from ray.rllib.agents.sac.sac_torch_model import SACTorchModel from ray.rllib.models import ModelCatalog -from ray.rllib.models.tf.tf_action_dist import (Categorical, SquashedGaussian, - DiagGaussian) +from ray.rllib.models.tf.tf_action_dist import Beta, Categorical, \ + DiagGaussian, SquashedGaussian from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.tf_policy_template import build_tf_policy from ray.rllib.utils.error import UnsupportedSpaceException @@ -101,15 +101,14 @@ def postprocess_trajectory(policy, def get_dist_class(config, action_space): - assert config["_use_beta_distribution"] is False, \ - "Beta-distr. not supported for tf!" - if isinstance(action_space, Discrete): - action_dist_class = Categorical + return Categorical else: - action_dist_class = (SquashedGaussian - if config["normalize_actions"] else DiagGaussian) - return action_dist_class + if config["normalize_actions"]: + return SquashedGaussian if \ + not config["_use_beta_distribution"] else Beta + else: + return DiagGaussian def get_distribution_inputs_and_class(policy, diff --git a/rllib/models/tests/test_distributions.py b/rllib/models/tests/test_distributions.py index b9dddcaf3..2bd119326 100644 --- a/rllib/models/tests/test_distributions.py +++ b/rllib/models/tests/test_distributions.py @@ -1,14 +1,15 @@ +from functools import partial import numpy as np -from gym.spaces import Box, Tuple -from scipy.stats import norm, beta +from gym.spaces import Box, Dict, Tuple +from scipy.stats import beta, norm import unittest -from ray.rllib.models.tf.tf_action_dist import Categorical, \ +from ray.rllib.models.tf.tf_action_dist import Beta, Categorical, \ DiagGaussian, GumbelSoftmax, MultiActionDistribution, MultiCategorical, \ SquashedGaussian -from ray.rllib.models.torch.torch_action_dist import TorchMultiCategorical, \ - TorchSquashedGaussian, TorchBeta, TorchCategorical, \ - TorchMultiActionDistribution, TorchDiagGaussian +from ray.rllib.models.torch.torch_action_dist import TorchBeta, \ + TorchCategorical, TorchDiagGaussian, TorchMultiActionDistribution, \ + TorchMultiCategorical, TorchSquashedGaussian from ray.rllib.utils import try_import_tree from ray.rllib.utils.framework import try_import_tf, try_import_torch from ray.rllib.utils.numpy import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, \ @@ -241,13 +242,17 @@ class TestDistributions(unittest.TestCase): low, high = -1.0, 2.0 plain_beta_value_space = Box(0.0, 1.0, shape=(200, 5)) - for fw, sess in framework_iterator(frameworks="torch", session=True): - cls = TorchBeta + for fw, sess in framework_iterator(session=True): + cls = TorchBeta if fw == "torch" else Beta inputs = input_space.sample() beta_distribution = cls(inputs, {}, low=low, high=high) inputs = beta_distribution.inputs - alpha, beta_ = np.split(inputs.numpy(), 2, axis=-1) + if sess: + inputs = sess.run(inputs) + else: + inputs = inputs.numpy() + alpha, beta_ = np.split(inputs, 2, axis=-1) # Mean for a Beta distribution: 1 / [1 + (beta/alpha)] expected = (1.0 / (1.0 + beta_ / alpha)) * (high - low) + low @@ -270,11 +275,17 @@ class TestDistributions(unittest.TestCase): inputs = input_space.sample() beta_distribution = cls(inputs, {}, low=low, high=high) inputs = beta_distribution.inputs - alpha, beta_ = np.split(inputs.numpy(), 2, axis=-1) + if sess: + inputs = sess.run(inputs) + else: + inputs = inputs.numpy() + alpha, beta_ = np.split(inputs, 2, axis=-1) values = plain_beta_value_space.sample() values_scaled = values * (high - low) + low - out = beta_distribution.logp(torch.Tensor(values_scaled)) + if fw == "torch": + values_scaled = torch.Tensor(values_scaled) + out = beta_distribution.logp(values_scaled) check( out, np.sum(np.log(beta.pdf(values, alpha, beta_)), -1), @@ -310,14 +321,17 @@ class TestDistributions(unittest.TestCase): check(np.mean(np.argmax(outs, -1)), expected_mean, rtol=0.08) def test_multi_action_distribution(self): - """Tests the MultiActionDistribution (only torch so far).""" + """Tests the MultiActionDistribution (across all frameworks).""" batch_size = 1000 input_space = Tuple([ Box(-10.0, 10.0, shape=(batch_size, 4)), Box(-2.0, 2.0, shape=( batch_size, 6, - )) + )), + Dict({ + "a": Box(-1.0, 1.0, shape=(batch_size, 4)) + }), ]) std_space = Box( -0.05, 0.05, shape=( @@ -325,37 +339,57 @@ class TestDistributions(unittest.TestCase): 3, )) + low, high = -1.0, 1.0 value_space = Tuple([ Box(0, 3, shape=(batch_size, ), dtype=np.int32), - Box(-2.0, 2.0, shape=(batch_size, 3), dtype=np.float32) + Box(-2.0, 2.0, shape=(batch_size, 3), dtype=np.float32), + Dict({ + "a": Box(0.0, 1.0, shape=(batch_size, 2), dtype=np.float32) + }) ]) - for fw, sess in framework_iterator(frameworks="torch", session=True): + for fw, sess in framework_iterator(session=True): if fw == "torch": cls = TorchMultiActionDistribution - child_distr_cls = [TorchCategorical, TorchDiagGaussian] + child_distr_cls = [ + TorchCategorical, TorchDiagGaussian, + partial(TorchBeta, low=low, high=high) + ] else: cls = MultiActionDistribution - child_distr_cls = [Categorical, DiagGaussian] + child_distr_cls = [ + Categorical, + DiagGaussian, + partial(Beta, low=low, high=high), + ] inputs = list(input_space.sample()) distr = cls( - np.concatenate([inputs[0], inputs[1]], axis=1), + np.concatenate([inputs[0], inputs[1], inputs[2]["a"]], axis=1), model={}, action_space=value_space, child_distributions=child_distr_cls, - input_lens=[4, 6]) + input_lens=[4, 6, 4]) + # Adjust inputs for the Beta distr just as Beta itself does. + inputs[2]["a"] = np.clip(inputs[2]["a"], np.log(SMALL_NUMBER), + -np.log(SMALL_NUMBER)) + inputs[2]["a"] = np.log(np.exp(inputs[2]["a"]) + 1.0) + 1.0 # Sample deterministically. expected_det = [ np.argmax(inputs[0], axis=-1), inputs[1][:, :3], # [:3]=Mean values. + # Mean for a Beta distribution: + # 1 / [1 + (beta/alpha)] * range + low + (1.0 / (1.0 + inputs[2]["a"][:, 2:] / inputs[2]["a"][:, 0:2])) + * (high - low) + low, ] out = distr.deterministic_sample() if sess: out = sess.run(out) check(out[0], expected_det[0]) check(out[1], expected_det[1]) + check(out[2]["a"], expected_det[2]) # Stochastic sampling -> expect roughly the mean. inputs = list(input_space.sample()) @@ -364,15 +398,23 @@ class TestDistributions(unittest.TestCase): inputs[0] = softmax(inputs[0], -1) # Fix std inputs (shouldn't be too large for this test). inputs[1][:, 3:] = std_space.sample() + # Adjust inputs for the Beta distr just as Beta itself does. + inputs[2]["a"] = np.clip(inputs[2]["a"], np.log(SMALL_NUMBER), + -np.log(SMALL_NUMBER)) + inputs[2]["a"] = np.log(np.exp(inputs[2]["a"]) + 1.0) + 1.0 distr = cls( - np.concatenate([inputs[0], inputs[1]], axis=1), + np.concatenate([inputs[0], inputs[1], inputs[2]["a"]], axis=1), model={}, action_space=value_space, child_distributions=child_distr_cls, - input_lens=[4, 6]) + input_lens=[4, 6, 4]) expected_mean = [ np.mean(np.sum(inputs[0] * np.array([0, 1, 2, 3]), -1)), inputs[1][:, :3], # [:3]=Mean values. + # Mean for a Beta distribution: + # 1 / [1 + (beta/alpha)] * range + low + (1.0 / (1.0 + inputs[2]["a"][:, 2:] / inputs[2]["a"][:, :2])) * + (high - low) + low, ] out = distr.sample() if sess: @@ -381,21 +423,36 @@ class TestDistributions(unittest.TestCase): if fw == "torch": out[0] = out[0].numpy() out[1] = out[1].numpy() + out[2]["a"] = out[2]["a"].numpy() check(np.mean(out[0]), expected_mean[0], decimals=1) check(np.mean(out[1], 0), np.mean(expected_mean[1], 0), decimals=1) + check( + np.mean(out[2]["a"], 0), + np.mean(expected_mean[2], 0), + decimals=1) # Test log-likelihood outputs. # Make sure beta-values are within 0.0 and 1.0 for the numpy # calculation (which doesn't have scaling). inputs = list(input_space.sample()) + # Adjust inputs for the Beta distr just as Beta itself does. + inputs[2]["a"] = np.clip(inputs[2]["a"], np.log(SMALL_NUMBER), + -np.log(SMALL_NUMBER)) + inputs[2]["a"] = np.log(np.exp(inputs[2]["a"]) + 1.0) + 1.0 distr = cls( - np.concatenate([inputs[0], inputs[1]], axis=1), + np.concatenate([inputs[0], inputs[1], inputs[2]["a"]], axis=1), model={}, action_space=value_space, child_distributions=child_distr_cls, - input_lens=[4, 6]) + input_lens=[4, 6, 4]) inputs[0] = softmax(inputs[0], -1) values = list(value_space.sample()) + log_prob_beta = np.log( + beta.pdf(values[2]["a"], inputs[2]["a"][:, :2], + inputs[2]["a"][:, 2:])) + # Now do the up-scaling for [2] (beta values) to be between + # low/high. + values[2]["a"] = values[2]["a"] * (high - low) + low inputs[1][:, 3:] = np.exp(inputs[1][:, 3:]) expected_log_llh = np.sum( np.concatenate([ @@ -405,7 +462,7 @@ class TestDistributions(unittest.TestCase): for j, i in enumerate(inputs[0])]), -1), np.log( norm.pdf(values[1], inputs[1][:, :3], - inputs[1][:, 3:])) + inputs[1][:, 3:])), log_prob_beta ], -1), -1) values[0] = np.expand_dims(values[0], -1) diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index 8197b7f6e..3a60ce939 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/rllib/models/tf/tf_action_dist.py @@ -1,3 +1,4 @@ +from math import log import numpy as np import functools @@ -326,6 +327,51 @@ class SquashedGaussian(TFActionDistribution): return unsquashed +class Beta(TFActionDistribution): + """ + A Beta distribution is defined on the interval [0, 1] and parameterized by + shape parameters alpha and beta (also called concentration parameters). + + PDF(x; alpha, beta) = x**(alpha - 1) (1 - x)**(beta - 1) / Z + with Z = Gamma(alpha) Gamma(beta) / Gamma(alpha + beta) + and Gamma(n) = (n - 1)! + """ + + def __init__(self, inputs, model, low=0.0, high=1.0): + # Stabilize input parameters (possibly coming from a linear layer). + inputs = tf.clip_by_value(inputs, log(SMALL_NUMBER), + -log(SMALL_NUMBER)) + inputs = tf.math.log(tf.math.exp(inputs) + 1.0) + 1.0 + self.low = low + self.high = high + alpha, beta = tf.split(inputs, 2, axis=-1) + # Note: concentration0==beta, concentration1=alpha (!) + self.dist = tfp.distributions.Beta( + concentration1=alpha, concentration0=beta) + super().__init__(inputs, model) + + @override(ActionDistribution) + def deterministic_sample(self): + mean = self.dist.mean() + return self._squash(mean) + + @override(TFActionDistribution) + def _build_sample_op(self): + return self._squash(self.dist.sample()) + + @override(ActionDistribution) + def logp(self, x): + unsquashed_values = self._unsquash(x) + return tf.math.reduce_sum( + self.dist.log_prob(unsquashed_values), axis=-1) + + def _squash(self, raw_values): + return raw_values * (self.high - self.low) + self.low + + def _unsquash(self, values): + return (values - self.low) / (self.high - self.low) + + class Deterministic(TFActionDistribution): """Action distribution that returns the input values directly. diff --git a/rllib/tests/test_supported_spaces.py b/rllib/tests/test_supported_spaces.py index ccafffce7..57fc1c1b0 100644 --- a/rllib/tests/test_supported_spaces.py +++ b/rllib/tests/test_supported_spaces.py @@ -200,7 +200,10 @@ class ModelSupportedSpaces(unittest.TestCase): check_bounds=True) def test_dqn(self): - check_support("DQN", {"timesteps_per_iteration": 1}, self.stats) + config = {"timesteps_per_iteration": 1} + check_support("DQN", config, self.stats) + config["use_pytorch"] = True + check_support("DQN", config, self.stats) def test_es(self): check_support(