wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-21 12:50:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rllib] Support free_log_std in ModelV2 (#8380)

Browse Source 
* update

* factor

* update

* fix test failures

* fix torch net
This commit is contained in:
Eric Liang
2020-05-12 10:14:05 -07:00
committed by GitHub
parent ac1ed293e3
commit 7ce138a6dc
5 changed files with 155 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/agents/ppo/tests/test_ppo.py
+87 -25
View File
@@ -1,3 +1,4 @@
import copy
import numpy as np
import unittest
@@ -19,6 +20,22 @@ from ray.rllib.utils.test_utils import check, framework_iterator, \
tf = try_import_tf()
# Fake CartPole episode of n time steps.
FAKE_BATCH = {
SampleBatch.CUR_OBS: np.array(
[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]],
dtype=np.float32),
SampleBatch.ACTIONS: np.array([0, 1, 1]),
SampleBatch.PREV_ACTIONS: np.array([0, 1, 1]),
SampleBatch.REWARDS: np.array([1.0, -1.0, .5], dtype=np.float32),
SampleBatch.PREV_REWARDS: np.array([1.0, -1.0, .5], dtype=np.float32),
SampleBatch.DONES: np.array([False, False, True]),
SampleBatch.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32),
SampleBatch.ACTION_DIST_INPUTS: np.array(
[[-2., 0.5], [-3., -0.3], [-0.1, 2.5]], dtype=np.float32),
SampleBatch.ACTION_LOGP: np.array([-0.5, -0.1, -0.2], dtype=np.float32),
}
class TestPPO(unittest.TestCase):
@classmethod
@@ -31,7 +48,7 @@ class TestPPO(unittest.TestCase):
def test_ppo_compilation(self):
"""Test whether a PPOTrainer can be built with both frameworks."""
config = ppo.DEFAULT_CONFIG.copy()
config = copy.deepcopy(ppo.DEFAULT_CONFIG)
config["num_workers"] = 0 # Run locally.
num_iterations = 2
@@ -43,7 +60,7 @@ class TestPPO(unittest.TestCase):
def test_ppo_fake_multi_gpu_learning(self):
"""Test whether PPOTrainer can learn CartPole w/ faked multi-GPU."""
config = ppo.DEFAULT_CONFIG.copy()
config = copy.deepcopy(ppo.DEFAULT_CONFIG)
# Fake GPU setup.
config["num_gpus"] = 2
config["_fake_gpus"] = True
@@ -70,7 +87,7 @@ class TestPPO(unittest.TestCase):
def test_ppo_exploration_setup(self):
"""Tests, whether PPO runs with different exploration setups."""
config = ppo.DEFAULT_CONFIG.copy()
config = copy.deepcopy(ppo.DEFAULT_CONFIG)
config["num_workers"] = 0 # Run locally.
config["env_config"] = {"is_slippery": False, "map_name": "4x4"}
obs = np.array(0)
@@ -107,46 +124,91 @@ class TestPPO(unittest.TestCase):
prev_reward=np.array(1.0)))
check(np.mean(actions), 1.5, atol=0.2)
def test_ppo_free_log_std(self):
"""Tests the free log std option works."""
config = copy.deepcopy(ppo.DEFAULT_CONFIG)
config["num_workers"] = 0 # Run locally.
config["gamma"] = 0.99
config["model"]["fcnet_hiddens"] = [10]
config["model"]["fcnet_activation"] = "linear"
config["model"]["free_log_std"] = True
config["vf_share_layers"] = True
for fw, sess in framework_iterator(config, session=True):
trainer = ppo.PPOTrainer(config=config, env="CartPole-v0")
policy = trainer.get_policy()
# Check the free log std var is created.
if fw == "torch":
matching = [
v for (n, v) in policy.model.named_parameters()
if "log_std" in n
]
else:
matching = [
v for v in policy.model.trainable_variables()
if "log_std" in str(v)
]
assert len(matching) == 1, matching
log_std_var = matching[0]
def get_value():
if fw == "tf":
return policy.get_session().run(log_std_var)[0]
elif fw == "torch":
return log_std_var.detach().numpy()[0]
else:
return log_std_var.numpy()[0]
# Check the variable is initially zero.
init_std = get_value()
assert init_std == 0.0, init_std
if fw == "tf" or fw == "eager":
batch = postprocess_ppo_gae_tf(policy, FAKE_BATCH)
else:
batch = postprocess_ppo_gae_torch(policy, FAKE_BATCH)
batch = policy._lazy_tensor_dict(batch)
policy.learn_on_batch(batch)
# Check the variable is updated.
post_std = get_value()
assert post_std != 0.0, post_std
def test_ppo_loss_function(self):
"""Tests the PPO loss function math."""
config = ppo.DEFAULT_CONFIG.copy()
config = copy.deepcopy(ppo.DEFAULT_CONFIG)
config["num_workers"] = 0 # Run locally.
config["gamma"] = 0.99
config["model"]["fcnet_hiddens"] = [10]
config["model"]["fcnet_activation"] = "linear"
config["vf_share_layers"] = True
# Fake CartPole episode of n time steps.
train_batch = {
SampleBatch.CUR_OBS: np.array(
[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8],
[0.9, 1.0, 1.1, 1.2]],
dtype=np.float32),
SampleBatch.ACTIONS: np.array([0, 1, 1]),
SampleBatch.PREV_ACTIONS: np.array([0, 1, 1]),
SampleBatch.REWARDS: np.array([1.0, -1.0, .5], dtype=np.float32),
SampleBatch.PREV_REWARDS: np.array(
[1.0, -1.0, .5], dtype=np.float32),
SampleBatch.DONES: np.array([False, False, True]),
SampleBatch.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32),
SampleBatch.ACTION_DIST_INPUTS: np.array(
[[-2., 0.5], [-3., -0.3], [-0.1, 2.5]], dtype=np.float32),
SampleBatch.ACTION_LOGP: np.array(
[-0.5, -0.1, -0.2], dtype=np.float32),
}
for fw, sess in framework_iterator(config, session=True):
trainer = ppo.PPOTrainer(config=config, env="CartPole-v0")
policy = trainer.get_policy()
# Check no free log std var by default.
if fw == "torch":
matching = [
v for (n, v) in policy.model.named_parameters()
if "log_std" in n
]
else:
matching = [
v for v in policy.model.trainable_variables()
if "log_std" in str(v)
]
assert len(matching) == 0, matching
# Post-process (calculate simple (non-GAE) advantages) and attach
# to train_batch dict.
# A = [0.99^2 * 0.5 + 0.99 * -1.0 + 1.0, 0.99 * 0.5 - 1.0, 0.5] =
# [0.50005, -0.505, 0.5]
if fw == "tf" or fw == "eager":
train_batch = postprocess_ppo_gae_tf(policy, train_batch)
train_batch = postprocess_ppo_gae_tf(policy, FAKE_BATCH)
else:
train_batch = postprocess_ppo_gae_torch(policy, train_batch)
train_batch = postprocess_ppo_gae_torch(policy, FAKE_BATCH)
train_batch = policy._lazy_tensor_dict(train_batch)
# Check Advantage values.
rllib/models/catalog.py
+3 -1
View File
@@ -44,7 +44,9 @@ MODEL_DEFAULTS = {
"fcnet_activation": "tanh",
# Number of hidden layers for fully connected net
"fcnet_hiddens": [256, 256],
# For control envs, documented in ray.rllib.models.Model
# For DiagGaussian action distributions, make the second half of the model
# outputs floating bias variables instead of state-dependent. This only
# has an effect is using the default fully connected net.
"free_log_std": False,
# Whether to skip the final linear layer used to resize the hidden layer
# outputs to size `num_outputs`. If True, then the last hidden layer
rllib/models/tf/fcnet_v2.py
+28 -7
View File
@@ -19,8 +19,19 @@ class FullyConnectedNetwork(TFModelV2):
hiddens = model_config.get("fcnet_hiddens", [])
no_final_linear = model_config.get("no_final_linear")
vf_share_layers = model_config.get("vf_share_layers")
free_log_std = model_config.get("free_log_std")
# we are using obs_flat, so take the flattened shape as input
# Maybe generate free-floating bias variables for the second half of
# the outputs.
if free_log_std:
assert num_outputs % 2 == 0, (
"num_outputs must be divisible by two", num_outputs)
num_outputs = num_outputs // 2
self.log_std_var = tf.Variable(
[0.0] * num_outputs, dtype=tf.float32, name="log_std")
self.register_variables([self.log_std_var])
# We are using obs_flat, so take the flattened shape as input.
inputs = tf.keras.layers.Input(
shape=(np.product(obs_space.shape), ), name="observations")
last_layer = layer_out = inputs
@@ -37,9 +48,9 @@ class FullyConnectedNetwork(TFModelV2):
# The last layer is adjusted to be of size num_outputs, but it's a
# layer with activation.
if no_final_linear and self.num_outputs:
if no_final_linear and num_outputs:
layer_out = tf.keras.layers.Dense(
self.num_outputs,
num_outputs,
name="fc_out",
activation=activation,
kernel_initializer=normc_initializer(1.0))(last_layer)
@@ -52,18 +63,28 @@ class FullyConnectedNetwork(TFModelV2):
name="fc_{}".format(i),
activation=activation,
kernel_initializer=normc_initializer(1.0))(last_layer)
if self.num_outputs:
if num_outputs:
layer_out = tf.keras.layers.Dense(
self.num_outputs,
num_outputs,
name="fc_out",
activation=None,
kernel_initializer=normc_initializer(0.01))(last_layer)
# Adjust self.num_outputs to be the number of nodes in the last
# layer.
# Adjust num_outputs to be the number of nodes in the last layer.
else:
self.num_outputs = (
[np.product(obs_space.shape)] + hiddens[-1:-1])[-1]
# Concat the log std vars to the end of the state-dependent means.
if free_log_std:
def tiled_log_std(x):
return tf.tile(
tf.expand_dims(self.log_std_var, 0), [tf.shape(x)[0], 1])
log_std_out = tf.keras.layers.Lambda(tiled_log_std)(inputs)
layer_out = tf.keras.layers.Concatenate(axis=1)(
[layer_out, log_std_out])
if not vf_share_layers:
# build a parallel set of hidden layers for the value net
last_layer = inputs
rllib/models/torch/fcnet.py
+22 -7
View File
@@ -2,12 +2,13 @@ import logging
import numpy as np
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.models.torch.misc import SlimFC, normc_initializer
from ray.rllib.models.torch.misc import SlimFC, AppendBiasLayer, \
normc_initializer
from ray.rllib.utils.annotations import override
from ray.rllib.utils.framework import get_activation_fn
from ray.rllib.utils import try_import_torch
_, nn = try_import_torch()
torch, nn = try_import_torch()
logger = logging.getLogger(__name__)
@@ -25,6 +26,7 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
model_config.get("fcnet_activation"), framework="torch")
hiddens = model_config.get("fcnet_hiddens")
no_final_linear = model_config.get("no_final_linear")
self.free_log_std = model_config.get("free_log_std")
# TODO(sven): implement case: vf_shared_layers = False.
# vf_share_layers = model_config.get("vf_share_layers")
@@ -34,6 +36,13 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
prev_layer_size = int(np.product(obs_space.shape))
self._logits = None
# Maybe generate free-floating bias variables for the second half of
# the outputs.
if self.free_log_std:
assert num_outputs % 2 == 0, (
"num_outputs must be divisible by two", num_outputs)
num_outputs = num_outputs // 2
# Create layers 0 to second-last.
for size in hiddens[:-1]:
layers.append(
@@ -46,14 +55,14 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
# The last layer is adjusted to be of size num_outputs, but it's a
# layer with activation.
if no_final_linear and self.num_outputs:
if no_final_linear and num_outputs:
layers.append(
SlimFC(
in_size=prev_layer_size,
out_size=self.num_outputs,
out_size=num_outputs,
initializer=normc_initializer(1.0),
activation_fn=activation))
prev_layer_size = self.num_outputs
prev_layer_size = num_outputs
# Finish the layers with the provided sizes (`hiddens`), plus -
# iff num_outputs > 0 - a last linear layer of size num_outputs.
else:
@@ -65,16 +74,20 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
initializer=normc_initializer(1.0),
activation_fn=activation))
prev_layer_size = hiddens[-1]
if self.num_outputs:
if num_outputs:
self._logits = SlimFC(
in_size=prev_layer_size,
out_size=self.num_outputs,
out_size=num_outputs,
initializer=normc_initializer(0.01),
activation_fn=None)
else:
self.num_outputs = (
[np.product(obs_space.shape)] + hiddens[-1:-1])[-1]
# Layer to add the log std vars to the state-dependent means.
if self.free_log_std:
self._append_free_log_std = AppendBiasLayer(num_outputs)
self._hidden_layers = nn.Sequential(*layers)
# TODO(sven): Implement non-shared value branch.
@@ -91,6 +104,8 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
obs = input_dict["obs_flat"].float()
features = self._hidden_layers(obs.reshape(obs.shape[0], -1))
logits = self._logits(features) if self._logits else features
if self.free_log_std:
logits = self._append_free_log_std(logits)
self._cur_value = self._value_branch(features).squeeze(1)
return logits, state
rllib/models/torch/misc.py
+15
View File
@@ -106,3 +106,18 @@ class SlimFC(nn.Module):
def forward(self, x):
return self._model(x)
class AppendBiasLayer(nn.Module):
"""Simple bias appending layer for free_log_std."""
def __init__(self, num_bias_vars):
super().__init__()
self.log_std = torch.nn.Parameter(
torch.as_tensor([0.0] * num_bias_vars))
self.register_parameter("log_std", self.log_std)
def forward(self, x):
out = torch.cat(
[x, self.log_std.unsqueeze(0).repeat([len(x), 1])], axis=1)
return out