wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-06 13:31:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] Beta distribution. (#8229)

Browse Source
This commit is contained in:
Sven Mika
2020-04-30 11:09:33 -07:00
committed by GitHub
parent b23b6addfc
commit eea75ac623
4 changed files with 139 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/agents/sac/sac_tf_policy.py
+8 -9
View File
@@ -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,
rllib/models/tests/test_distributions.py
+81 -24
View File
@@ -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)
rllib/models/tf/tf_action_dist.py
+46
View File
@@ -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.
rllib/tests/test_supported_spaces.py
+4 -1
View File
@@ -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(