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

[RLlib] Fix use_lstm flag for ModelV2 (w/o ModelV1 wrapping) and add it for PyTorch. (#8734)

Browse Source
This commit is contained in:
Sven Mika
2020-06-05 15:40:30 +02:00
committed by GitHub
parent d78757623d
commit c74dc58f8b
21 changed files with 331 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ci/travis/ci.sh
+1 -1
View File
@@ -168,7 +168,7 @@ build_sphinx_docs() {
if [ "${OSTYPE}" = msys ]; then
echo "WARNING: Documentation not built on Windows due to currently-unresolved issues"
else
sphinx-build -q -W -E -T -b html source _build/html
sphinx-build -q -E -T -b html source _build/html
fi
)
}
doc/README.md
+1 -1
View File
@@ -12,5 +12,5 @@ open _build/html/index.html
To test if there are any build errors with the documentation, do the following.
```
sphinx-build -W -b html -d _build/doctrees source _build/html
sphinx-build -b html -d _build/doctrees source _build/html
```
doc/source/rllib-models.rst
+1 -1
View File
@@ -121,7 +121,7 @@ Once implemented, the model can then be registered and used in place of a built-
from ray.rllib.models import ModelCatalog
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
class CustomTorchModel(nn.Module, TorchModelV2):
class CustomTorchModel(TorchModelV2):
def __init__(self, obs_space, action_space, num_outputs, model_config, name): ...
def forward(self, input_dict, state, seq_lens): ...
def value_function(self): ...
rllib/BUILD
+28 -8
View File
@@ -283,6 +283,26 @@ py_test(
args = ["--torch", "--yaml-dir=tuned_examples/ppo"]
)
py_test(
name = "run_regression_tests_repeat_after_me_tf",
main = "tests/run_regression_tests.py",
tags = ["learning_tests_tf"],
size = "medium",
srcs = ["tests/run_regression_tests.py"],
data = ["tuned_examples/ppo/repeatafterme-ppo-lstm.yaml"],
args = ["--yaml-dir=tuned_examples/ppo"]
)
py_test(
name = "run_regression_tests_repeat_after_me_torch",
main = "tests/run_regression_tests.py",
tags = ["learning_tests_tf"],
size = "medium",
srcs = ["tests/run_regression_tests.py"],
data = ["tuned_examples/ppo/repeatafterme-ppo-lstm.yaml"],
args = ["--torch", "--yaml-dir=tuned_examples/ppo"]
)
# SAC
py_test(
name = "run_regression_tests_cartpole_sac_tf",
@@ -1285,7 +1305,7 @@ py_test(
name = "test_rollout",
main = "tests/test_rollout.py",
tags = ["tests_dir", "tests_dir_R"],
size = "enormous",
size = "large",
data = ["train.py", "rollout.py"],
srcs = ["tests/test_rollout.py"]
)
@@ -1417,7 +1437,7 @@ py_test(
name = "examples/cartpole_lstm_impala_tf",
main = "examples/cartpole_lstm.py",
tags = ["examples", "examples_C"],
size = "medium",
size = "large",
srcs = ["examples/cartpole_lstm.py"],
args = ["--as-test", "--run=IMPALA", "--stop-reward=40", "--num-cpus=4"]
)
@@ -1426,7 +1446,7 @@ py_test(
name = "examples/cartpole_lstm_impala_torch",
main = "examples/cartpole_lstm.py",
tags = ["examples", "examples_C"],
size = "medium",
size = "large",
srcs = ["examples/cartpole_lstm.py"],
args = ["--as-test", "--torch", "--run=IMPALA", "--stop-reward=40", "--num-cpus=4"]
)
@@ -1435,7 +1455,7 @@ py_test(
name = "examples/cartpole_lstm_ppo_tf",
main = "examples/cartpole_lstm.py",
tags = ["examples", "examples_C"],
size = "medium",
size = "large",
srcs = ["examples/cartpole_lstm.py"],
args = ["--as-test", "--run=PPO", "--stop-reward=40", "--num-cpus=4"]
)
@@ -1444,7 +1464,7 @@ py_test(
name = "examples/cartpole_lstm_ppo_torch",
main = "examples/cartpole_lstm.py",
tags = ["examples", "examples_C"],
size = "medium",
size = "large",
srcs = ["examples/cartpole_lstm.py"],
args = ["--as-test", "--torch", "--run=PPO", "--stop-reward=40", "--num-cpus=4"]
)
@@ -1453,7 +1473,7 @@ py_test(
name = "examples/cartpole_lstm_ppo_tf_with_prev_a_and_r",
main = "examples/cartpole_lstm.py",
tags = ["examples", "examples_C"],
size = "medium",
size = "large",
srcs = ["examples/cartpole_lstm.py"],
args = ["--as-test", "--run=PPO", "--stop-reward=40", "--use-prev-action-reward", "--num-cpus=4"]
)
@@ -1462,7 +1482,7 @@ py_test(
name = "examples/centralized_critic_tf",
main = "examples/centralized_critic.py",
tags = ["examples", "examples_C"],
size = "medium",
size = "large",
srcs = ["examples/centralized_critic.py"],
args = ["--as-test", "--stop-reward=7.2"]
)
@@ -1471,7 +1491,7 @@ py_test(
name = "examples/centralized_critic_torch",
main = "examples/centralized_critic.py",
tags = ["examples", "examples_C"],
size = "medium",
size = "large",
srcs = ["examples/centralized_critic.py"],
args = ["--as-test", "--torch", "--stop-reward=7.2"]
)
rllib/agents/impala/tests/test_impala.py
+4 -3
View File
@@ -26,8 +26,9 @@ class TestIMPALA(unittest.TestCase):
local_cfg = config.copy()
for env in ["Pendulum-v0", "CartPole-v0"]:
print("Env={}".format(env))
print("w/ LSTM")
print("w/o LSTM")
# Test w/o LSTM.
local_cfg["model"]["use_lstm"] = False
local_cfg["num_aggregation_workers"] = 0
trainer = impala.ImpalaTrainer(config=local_cfg, env=env)
for i in range(num_iterations):
@@ -36,13 +37,13 @@ class TestIMPALA(unittest.TestCase):
trainer.stop()
# Test w/ LSTM.
print("w/o LSTM")
print("w/ LSTM")
local_cfg["model"]["use_lstm"] = True
local_cfg["num_aggregation_workers"] = 2
trainer = impala.ImpalaTrainer(config=local_cfg, env=env)
for i in range(num_iterations):
print(trainer.train())
check_compute_action(trainer)
check_compute_action(trainer, include_state=True)
trainer.stop()
rllib/agents/impala/vtrace_tf_policy.py
+1 -1
View File
@@ -177,7 +177,7 @@ def build_vtrace_loss(policy, model, dist_class, train_batch):
values = model.value_function()
if policy.is_recurrent():
max_seq_len = tf.reduce_max(train_batch["seq_lens"]) - 1
max_seq_len = tf.reduce_max(train_batch["seq_lens"])
mask = tf.sequence_mask(train_batch["seq_lens"], max_seq_len)
mask = tf.reshape(mask, [-1])
else:
rllib/agents/impala/vtrace_torch_policy.py
+3 -3
View File
@@ -147,9 +147,9 @@ def build_vtrace_loss(policy, model, dist_class, train_batch):
values = model.value_function()
if policy.is_recurrent():
max_seq_len = torch.max(train_batch["seq_lens"]) - 1
mask = sequence_mask(train_batch["seq_lens"], max_seq_len)
mask = torch.reshape(mask, [-1])
max_seq_len = torch.max(train_batch["seq_lens"])
mask_orig = sequence_mask(train_batch["seq_lens"], max_seq_len)
mask = torch.reshape(mask_orig, [-1])
else:
mask = torch.ones_like(rewards)
rllib/examples/env/repeat_after_me_env.py
Vendored
+1 -1
View File
@@ -9,7 +9,7 @@ class RepeatAfterMeEnv(gym.Env):
def __init__(self, config):
self.observation_space = Discrete(2)
self.action_space = Discrete(2)
self.delay = config["repeat_delay"]
self.delay = config.get("repeat_delay", 1)
assert self.delay >= 1, "`repeat_delay` must be at least 1!"
self.history = []
rllib/models/catalog.py
+28 -9
View File
@@ -11,15 +11,18 @@ from ray.rllib.models.preprocessors import get_preprocessor
from ray.rllib.models.tf.fcnet_v1 import FullyConnectedNetwork
from ray.rllib.models.tf.lstm_v1 import LSTM
from ray.rllib.models.tf.modelv1_compat import make_v1_wrapper
from ray.rllib.models.tf.recurrent_net import LSTMWrapper
from ray.rllib.models.tf.tf_action_dist import Categorical, \
Deterministic, DiagGaussian, Dirichlet, \
MultiActionDistribution, MultiCategorical
from ray.rllib.models.tf.tf_modelv2 import TFModelV2
from ray.rllib.models.tf.visionnet_v1 import VisionNetwork
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.models.torch.recurrent_net import LSTMWrapper as \
TorchLSTMWrapper
from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \
TorchDeterministic, TorchDiagGaussian, \
TorchMultiActionDistribution, TorchMultiCategorical
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.utils import try_import_tree
from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI
from ray.rllib.utils.deprecation import deprecation_warning, DEPRECATED_VALUE
@@ -57,13 +60,13 @@ MODEL_DEFAULTS = {
"vf_share_layers": True,
# == LSTM ==
# Whether to wrap the model with a LSTM
# Whether to wrap the model with an LSTM.
"use_lstm": False,
# Max seq len for training the LSTM, defaults to 20
# Max seq len for training the LSTM, defaults to 20.
"max_seq_len": 20,
# Size of the LSTM cell
# Size of the LSTM cell.
"lstm_cell_size": 256,
# Whether to feed a_{t-1}, r_{t-1} to LSTM
# Whether to feed a_{t-1}, r_{t-1} to LSTM.
"lstm_use_prev_action_reward": False,
# When using modelv1 models with a modelv2 algorithm, you may have to
# define the state shape here (e.g., [256, 256]).
@@ -107,8 +110,9 @@ class ModelCatalog:
>>> observation = prep.transform(raw_observation)
>>> dist_class, dist_dim = ModelCatalog.get_action_dist(
env.action_space, {})
>>> model = ModelCatalog.get_model(inputs, dist_dim, options)
... env.action_space, {})
>>> model = ModelCatalog.get_model_v2(
... obs_space, action_space, num_outputs, options)
>>> dist = dist_class(model.outputs, model)
>>> action = dist.sample()
"""
@@ -307,6 +311,7 @@ class ModelCatalog:
else:
model_cls = _global_registry.get(RLLIB_MODEL,
model_config["custom_model"])
# TODO(sven): Hard-deprecate Model(V1).
if issubclass(model_cls, ModelV2):
logger.info("Wrapping {} as {}".format(model_cls,
@@ -374,10 +379,18 @@ class ModelCatalog:
if framework in ["tf", "tfe"]:
v2_class = None
# try to get a default v2 model
# Try to get a default v2 model.
if not model_config.get("custom_model"):
v2_class = default_model or ModelCatalog._get_v2_model_class(
obs_space, model_config, framework=framework)
if model_config.get("use_lstm"):
wrapped_cls = v2_class
forward = wrapped_cls.forward
v2_class = ModelCatalog._wrap_if_needed(
wrapped_cls, LSTMWrapper)
v2_class._wrapped_forward = forward
# fallback to a default v1 model
if v2_class is None:
if tf.executing_eagerly():
@@ -387,7 +400,7 @@ class ModelCatalog:
"observation space: {}, use_lstm={}".format(
obs_space, model_config.get("use_lstm")))
v2_class = make_v1_wrapper(ModelCatalog.get_model)
# wrap in the requested interface
# Wrap in the requested interface.
wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface)
return wrapper(obs_space, action_space, num_outputs, model_config,
name, **model_kwargs)
@@ -395,6 +408,12 @@ class ModelCatalog:
v2_class = \
default_model or ModelCatalog._get_v2_model_class(
obs_space, model_config, framework=framework)
if model_config.get("use_lstm"):
wrapped_cls = v2_class
forward = wrapped_cls.forward
v2_class = ModelCatalog._wrap_if_needed(
wrapped_cls, TorchLSTMWrapper)
v2_class._wrapped_forward = forward
# Wrap in the requested interface.
wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface)
return wrapper(obs_space, action_space, num_outputs, model_config,
rllib/models/tf/fcnet.py
+21 -14
View File
@@ -21,7 +21,7 @@ class FullyConnectedNetwork(TFModelV2):
vf_share_layers = model_config.get("vf_share_layers")
free_log_std = model_config.get("free_log_std")
# Maybe generate free-floating bias variables for the second half of
# Generate free-floating bias variables for the second half of
# the outputs.
if free_log_std:
assert num_outputs % 2 == 0, (
@@ -34,7 +34,10 @@ class FullyConnectedNetwork(TFModelV2):
# 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
# Last hidden layer output (before logits outputs).
last_layer = inputs
# The action distribution outputs.
logits_out = None
i = 1
# Create layers 0 to second-last.
@@ -49,7 +52,7 @@ 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 num_outputs:
layer_out = tf.keras.layers.Dense(
logits_out = tf.keras.layers.Dense(
num_outputs,
name="fc_out",
activation=activation,
@@ -64,7 +67,7 @@ class FullyConnectedNetwork(TFModelV2):
activation=activation,
kernel_initializer=normc_initializer(1.0))(last_layer)
if num_outputs:
layer_out = tf.keras.layers.Dense(
logits_out = tf.keras.layers.Dense(
num_outputs,
name="fc_out",
activation=None,
@@ -72,38 +75,42 @@ class FullyConnectedNetwork(TFModelV2):
# 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]
[np.product(obs_space.shape)] + hiddens[-1:])[-1]
# Concat the log std vars to the end of the state-dependent means.
if free_log_std:
if free_log_std and logits_out is not None:
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])
logits_out = tf.keras.layers.Concatenate(axis=1)(
[logits_out, log_std_out])
last_vf_layer = None
if not vf_share_layers:
# build a parallel set of hidden layers for the value net
last_layer = inputs
# Build a parallel set of hidden layers for the value net.
last_vf_layer = inputs
i = 1
for size in hiddens:
last_layer = tf.keras.layers.Dense(
last_vf_layer = tf.keras.layers.Dense(
size,
name="fc_value_{}".format(i),
activation=activation,
kernel_initializer=normc_initializer(1.0))(last_layer)
kernel_initializer=normc_initializer(1.0))(last_vf_layer)
i += 1
value_out = tf.keras.layers.Dense(
1,
name="value_out",
activation=None,
kernel_initializer=normc_initializer(0.01))(last_layer)
kernel_initializer=normc_initializer(0.01))(
last_vf_layer if last_vf_layer is not None else last_layer)
self.base_model = tf.keras.Model(inputs, [layer_out, value_out])
self.base_model = tf.keras.Model(
inputs, [(logits_out
if logits_out is not None else last_layer), value_out])
self.register_variables(self.base_model.variables)
def forward(self, input_dict, state, seq_lens):
rllib/models/tf/recurrent_net.py
+77
View File
@@ -1,3 +1,5 @@
import numpy as np
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.models.tf.tf_modelv2 import TFModelV2
from ray.rllib.policy.rnn_sequencing import add_time_dimension
@@ -94,3 +96,78 @@ class RecurrentNetwork(TFModelV2):
]
"""
raise NotImplementedError("You must implement this for a RNN model")
class LSTMWrapper(RecurrentNetwork):
"""An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.
"""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super(LSTMWrapper, self).__init__(obs_space, action_space, None,
model_config, name)
self.cell_size = model_config["lstm_cell_size"]
# Define input layers.
input_layer = tf.keras.layers.Input(
shape=(None, self.num_outputs), name="inputs")
self.num_outputs = num_outputs
state_in_h = tf.keras.layers.Input(shape=(self.cell_size, ), name="h")
state_in_c = tf.keras.layers.Input(shape=(self.cell_size, ), name="c")
seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32)
# Preprocess observation with a hidden layer and send to LSTM cell
lstm_out, state_h, state_c = tf.keras.layers.LSTM(
self.cell_size,
return_sequences=True,
return_state=True,
name="lstm")(
inputs=input_layer,
mask=tf.sequence_mask(seq_in),
initial_state=[state_in_h, state_in_c])
# Postprocess LSTM output with another hidden layer and compute values
logits = tf.keras.layers.Dense(
self.num_outputs,
activation=tf.keras.activations.linear,
name="logits")(lstm_out)
values = tf.keras.layers.Dense(
1, activation=None, name="values")(lstm_out)
# Create the RNN model
self._rnn_model = tf.keras.Model(
inputs=[input_layer, seq_in, state_in_h, state_in_c],
outputs=[logits, values, state_h, state_c])
self.register_variables(self._rnn_model.variables)
self._rnn_model.summary()
@override(RecurrentNetwork)
def forward(self, input_dict, state, seq_lens):
assert seq_lens is not None
# Push obs through "unwrapped" net's `forward()` first.
wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
# Then through our LSTM.
input_dict["obs_flat"] = wrapped_out
return super().forward(input_dict, state, seq_lens)
@override(RecurrentNetwork)
def forward_rnn(self, inputs, state, seq_lens):
model_out, self._value_out, h, c = self._rnn_model([inputs, seq_lens] +
state)
return model_out, [h, c]
@override(ModelV2)
def get_initial_state(self):
return [
np.zeros(self.cell_size, np.float32),
np.zeros(self.cell_size, np.float32),
]
@override(ModelV2)
def value_function(self):
return tf.reshape(self._value_out, [-1])
rllib/models/torch/fcnet.py
+34 -16
View File
@@ -13,36 +13,32 @@ torch, nn = try_import_torch()
logger = logging.getLogger(__name__)
class FullyConnectedNetwork(TorchModelV2, nn.Module):
class FullyConnectedNetwork(TorchModelV2):
"""Generic fully connected network."""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
model_config, name)
nn.Module.__init__(self)
activation = get_activation_fn(
model_config.get("fcnet_activation"), framework="torch")
hiddens = model_config.get("fcnet_hiddens")
no_final_linear = model_config.get("no_final_linear")
self.vf_share_layers = model_config.get("vf_share_layers")
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")
logger.debug("Constructing fcnet {} {}".format(hiddens, activation))
layers = []
prev_layer_size = int(np.product(obs_space.shape))
self._logits = None
# Maybe generate free-floating bias variables for the second half of
# 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
layers = []
prev_layer_size = int(np.product(obs_space.shape))
self._logits = None
# Create layers 0 to second-last.
for size in hiddens[:-1]:
layers.append(
@@ -82,15 +78,30 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
activation_fn=None)
else:
self.num_outputs = (
[np.product(obs_space.shape)] + hiddens[-1:-1])[-1]
[np.product(obs_space.shape)] + hiddens[-1:])[-1]
# Layer to add the log std vars to the state-dependent means.
if self.free_log_std:
if self.free_log_std and self._logits:
self._append_free_log_std = AppendBiasLayer(num_outputs)
self._hidden_layers = nn.Sequential(*layers)
# TODO(sven): Implement non-shared value branch.
self._value_branch_separate = None
if not self.vf_share_layers:
# Build a parallel set of hidden layers for the value net.
prev_vf_layer_size = int(np.product(obs_space.shape))
self._value_branch_separate = []
for size in hiddens:
self._value_branch_separate.append(
SlimFC(
in_size=prev_vf_layer_size,
out_size=size,
activation_fn=activation,
initializer=normc_initializer(1.0)))
prev_vf_layer_size = size
self._value_branch_separate = nn.Sequential(
*self._value_branch_separate)
self._value_branch = SlimFC(
in_size=prev_layer_size,
out_size=1,
@@ -98,11 +109,14 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
activation_fn=None)
# Holds the current "base" output (before logits layer).
self._features = None
# Holds the last input, in case value branch is separate.
self._last_flat_in = None
@override(TorchModelV2)
def forward(self, input_dict, state, seq_lens):
obs = input_dict["obs_flat"].float()
self._features = self._hidden_layers(obs.reshape(obs.shape[0], -1))
self._last_flat_in = obs.reshape(obs.shape[0], -1)
self._features = self._hidden_layers(self._last_flat_in)
logits = self._logits(self._features) if self._logits else \
self._features
if self.free_log_std:
@@ -112,4 +126,8 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
@override(TorchModelV2)
def value_function(self):
assert self._features is not None, "must call forward() first"
return self._value_branch(self._features).squeeze(1)
if self._value_branch_separate:
return self._value_branch(
self._value_branch_separate(self._last_flat_in)).squeeze(1)
else:
return self._value_branch(self._features).squeeze(1)
rllib/models/torch/recurrent_net.py
+64 -7
View File
@@ -1,6 +1,7 @@
import numpy as np
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.models.torch.misc import SlimFC
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.policy.rnn_sequencing import add_time_dimension
from ray.rllib.utils.annotations import override, DeveloperAPI
@@ -10,7 +11,7 @@ torch, nn = try_import_torch()
@DeveloperAPI
class RecurrentNetwork(TorchModelV2, nn.Module):
class RecurrentNetwork(TorchModelV2):
"""Helper class to simplify implementing RNN models with TorchModelV2.
Instead of implementing forward(), you can implement forward_rnn() which
@@ -51,12 +52,6 @@ class RecurrentNetwork(TorchModelV2, nn.Module):
return q, [h]
"""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
model_config, name)
nn.Module.__init__(self)
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
"""Adds time dimension to batch before sending inputs to forward_rnn().
@@ -90,3 +85,65 @@ class RecurrentNetwork(TorchModelV2, nn.Module):
return model_out, [h, c]
"""
raise NotImplementedError("You must implement this for an RNN model")
class LSTMWrapper(RecurrentNetwork):
"""An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.
"""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super(LSTMWrapper, self).__init__(obs_space, action_space, None,
model_config, name)
self.cell_size = model_config["lstm_cell_size"]
self.lstm = nn.LSTM(self.num_outputs, self.cell_size, batch_first=True)
self.num_outputs = num_outputs
# Postprocess LSTM output with another hidden layer and compute values.
self._logits_branch = SlimFC(
in_size=self.cell_size,
out_size=self.num_outputs,
activation_fn=None,
initializer=torch.nn.init.xavier_uniform_)
self._value_branch = SlimFC(
in_size=self.cell_size,
out_size=1,
activation_fn=None,
initializer=torch.nn.init.xavier_uniform_)
@override(RecurrentNetwork)
def forward(self, input_dict, state, seq_lens):
assert seq_lens is not None
# Push obs through "unwrapped" net's `forward()` first.
wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
# Then through our LSTM.
input_dict["obs_flat"] = wrapped_out
return super().forward(input_dict, state, seq_lens)
@override(RecurrentNetwork)
def forward_rnn(self, inputs, state, seq_lens):
self._features, [h, c] = self.lstm(
inputs,
[torch.unsqueeze(state[0], 0),
torch.unsqueeze(state[1], 0)])
model_out = self._logits_branch(self._features)
return model_out, [torch.squeeze(h, 0), torch.squeeze(c, 0)]
@override(ModelV2)
def get_initial_state(self):
# Place hidden states on same device as model.
linear = next(self._logits_branch._model.children())
h = [
linear.weight.new(1, self.cell_size).zero_().squeeze(0),
linear.weight.new(1, self.cell_size).zero_().squeeze(0)
]
return h
@override(ModelV2)
def value_function(self):
assert self._features is not None, "must call forward() first"
return torch.reshape(self._value_branch(self._features), [-1])
rllib/models/torch/torch_modelv2.py
+2 -6
View File
@@ -6,7 +6,7 @@ _, nn = try_import_torch()
@PublicAPI
class TorchModelV2(ModelV2):
class TorchModelV2(ModelV2, nn.Module):
"""Torch version of ModelV2.
Note that this class by itself is not a valid model unless you
@@ -27,11 +27,6 @@ class TorchModelV2(ModelV2):
self._value_branch = ...
"""
if not isinstance(self, nn.Module):
raise ValueError(
"Subclasses of TorchModelV2 must also inherit from "
"nn.Module, e.g., MyModel(TorchModel, nn.Module)")
ModelV2.__init__(
self,
obs_space,
@@ -40,6 +35,7 @@ class TorchModelV2(ModelV2):
model_config,
name,
framework="torch")
nn.Module.__init__(self)
@override(ModelV2)
def variables(self, as_dict=False):
rllib/models/torch/visionnet.py
+1 -2
View File
@@ -9,14 +9,13 @@ from ray.rllib.utils import try_import_torch
_, nn = try_import_torch()
class VisionNetwork(TorchModelV2, nn.Module):
class VisionNetwork(TorchModelV2):
"""Generic vision network."""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
model_config, name)
nn.Module.__init__(self)
activation = get_activation_fn(
model_config.get("conv_activation"), framework="torch")
rllib/policy/policy.py
+7 -1
View File
@@ -6,10 +6,12 @@ from typing import Any
from ray.rllib.utils import try_import_tree
from ray.rllib.utils.annotations import DeveloperAPI
from ray.rllib.utils.exploration.exploration import Exploration
from ray.rllib.utils.framework import try_import_torch
from ray.rllib.utils.from_config import from_config
from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space, \
unbatch
torch, _ = try_import_torch()
tree = try_import_tree()
# By convention, metrics from optimizing the loss can be reported in the
@@ -157,7 +159,11 @@ class Policy(metaclass=ABCMeta):
if episode is not None:
episodes = [episode]
if state is not None:
state_batch = [[s] for s in state]
state_batch = [
s.unsqueeze(0)
if torch and isinstance(s, torch.Tensor) else [s]
for s in state
]
batched_action, state_out, info = self.compute_actions(
[obs],
rllib/tests/test_nested_observation_spaces.py
+1 -1
View File
@@ -228,7 +228,7 @@ class NestedSpacesTest(unittest.TestCase):
ModelCatalog.register_custom_model("invalid", InvalidModel)
self.assertRaisesRegexp(
ValueError,
"Subclasses of TorchModelV2 must also inherit from",
"optimizer got an empty parameter list",
lambda: PGTrainer(
env="CartPole-v0",
config={
rllib/tests/test_rollout.py
+7 -5
View File
@@ -8,7 +8,9 @@ from ray.rllib.utils.test_utils import framework_iterator
def rollout_test(algo, env="CartPole-v0", test_episode_rollout=False):
extra_config = ""
if algo == "ES":
if algo == "ARS":
extra_config = ",\"train_batch_size\": 10, \"noise_size\": 250000"
elif algo == "ES":
extra_config = ",\"episodes_per_batch\": 1,\"train_batch_size\": 10, "\
"\"noise_size\": 250000"
@@ -28,10 +30,10 @@ def rollout_test(algo, env="CartPole-v0", test_episode_rollout=False):
"--checkpoint-freq=1 ".format(rllib_dir, tmp_dir, algo) +
"--config='{" + "\"num_workers\": 1, \"num_gpus\": 0{}{}".
format(fw_, extra_config) +
", \"model\": {\"fcnet_hiddens\": [10]}"
"}' --stop='{\"training_iteration\": 1, "
"\"timesteps_per_iter\": 5, "
"\"min_iter_time_s\": 0.1}'" + " --env={}".format(env))
", \"timesteps_per_iteration\": 5,\"min_iter_time_s\": 0.1, "
"\"model\": {\"fcnet_hiddens\": [10]}"
"}' --stop='{\"training_iteration\": 1}'" +
" --env={}".format(env))
checkpoint_path = os.popen("ls {}/default/*/checkpoint_1/"
"checkpoint-1".format(tmp_dir)).read()[:-1]
rllib/train.py
+9 -1
View File
@@ -49,6 +49,11 @@ def create_parser(parser_creator=None):
"--no-ray-ui",
action="store_true",
help="Whether to disable the Ray web ui.")
parser.add_argument(
"--local-mode",
action="store_true",
help="Whether to run ray with `local_mode=True`. "
"Only if --ray-num-nodes is not used.")
parser.add_argument(
"--ray-num-cpus",
default=None,
@@ -178,6 +183,8 @@ def run(args, parser):
exp["config"]["framework"] = "tfe"
elif args.torch:
exp["config"]["framework"] = "torch"
else:
exp["config"]["framework"] = "tf"
if args.v:
exp["config"]["log_level"] = "INFO"
verbose = 2
@@ -207,7 +214,8 @@ def run(args, parser):
memory=args.ray_memory,
redis_max_memory=args.ray_redis_max_memory,
num_cpus=args.ray_num_cpus,
num_gpus=args.ray_num_gpus)
num_gpus=args.ray_num_gpus,
local_mode=args.local_mode)
run_experiments(
experiments,
scheduler=_make_scheduler(args),
rllib/tuned_examples/ppo/repeatafterme-ppo-lstm.yaml
+24
View File
@@ -0,0 +1,24 @@
repeat-after-me-ppo-w-lstm:
env: "ray.rllib.examples.env.repeat_after_me_env.RepeatAfterMeEnv"
run: PPO
stop:
episode_reward_mean: 50
timesteps_total: 100000
config:
# Works for both torch and tf.
framework: tf
env_config:
config:
repeat_delay: 2
gamma: 0.9
lr: 0.0003
num_workers: 0
num_envs_per_worker: 20
num_sgd_iter: 5
vf_share_layers: true
entropy_coeff: 0.00001
model:
use_lstm: true
lstm_cell_size: 64
max_seq_len: 20
fcnet_hiddens: [64]
rllib/utils/test_utils.py
+16 -4
View File
@@ -248,7 +248,9 @@ def check_learning_achieved(tune_results, min_reward):
print("ok")
def check_compute_action(trainer, include_prev_action_reward=False):
def check_compute_action(trainer,
include_state=False,
include_prev_action_reward=False):
"""Tests different combinations of arguments for trainer.compute_action.
Args:
@@ -269,17 +271,27 @@ def check_compute_action(trainer, include_prev_action_reward=False):
for explore in [True, False]:
for full_fetch in [True, False]:
obs = np.clip(obs_space.sample(), -1.0, 1.0)
state_in = None
if include_state:
state_in = pol.model.get_initial_state()
action_in = action_space.sample() \
if include_prev_action_reward else None
reward_in = 1.0 if include_prev_action_reward else None
action = trainer.compute_action(
out = trainer.compute_action(
obs,
state=state_in,
prev_action=action_in,
prev_reward=reward_in,
explore=explore,
full_fetch=full_fetch)
if full_fetch:
action, _, _ = action
state_out = None
if state_in or full_fetch:
action, state_out, _ = out
if state_out:
for si, so in zip(state_in, state_out):
check(list(si.shape), so.shape)
if not action_space.contains(action):
raise ValueError(
"Returned action ({}) of trainer {} not in Env's "