From c74dc58f8b38bb947b903d3f92efbd8d1cd675ce Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Fri, 5 Jun 2020 15:40:30 +0200 Subject: [PATCH] [RLlib] Fix `use_lstm` flag for ModelV2 (w/o ModelV1 wrapping) and add it for PyTorch. (#8734) --- ci/travis/ci.sh | 2 +- doc/README.md | 2 +- doc/source/rllib-models.rst | 2 +- rllib/BUILD | 36 +++++++-- rllib/agents/impala/tests/test_impala.py | 7 +- rllib/agents/impala/vtrace_tf_policy.py | 2 +- rllib/agents/impala/vtrace_torch_policy.py | 6 +- rllib/examples/env/repeat_after_me_env.py | 2 +- rllib/models/catalog.py | 37 ++++++--- rllib/models/tf/fcnet.py | 35 +++++---- rllib/models/tf/recurrent_net.py | 77 +++++++++++++++++++ rllib/models/torch/fcnet.py | 50 ++++++++---- rllib/models/torch/recurrent_net.py | 71 +++++++++++++++-- rllib/models/torch/torch_modelv2.py | 8 +- rllib/models/torch/visionnet.py | 3 +- rllib/policy/policy.py | 8 +- rllib/tests/test_nested_observation_spaces.py | 2 +- rllib/tests/test_rollout.py | 12 +-- rllib/train.py | 10 ++- .../ppo/repeatafterme-ppo-lstm.yaml | 24 ++++++ rllib/utils/test_utils.py | 20 ++++- 21 files changed, 331 insertions(+), 85 deletions(-) create mode 100644 rllib/tuned_examples/ppo/repeatafterme-ppo-lstm.yaml diff --git a/ci/travis/ci.sh b/ci/travis/ci.sh index ab8340949..689f69349 100755 --- a/ci/travis/ci.sh +++ b/ci/travis/ci.sh @@ -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 ) } diff --git a/doc/README.md b/doc/README.md index 153795a9a..98978a268 100644 --- a/doc/README.md +++ b/doc/README.md @@ -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 ``` diff --git a/doc/source/rllib-models.rst b/doc/source/rllib-models.rst index 2178573ac..d0b8e19fa 100644 --- a/doc/source/rllib-models.rst +++ b/doc/source/rllib-models.rst @@ -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): ... diff --git a/rllib/BUILD b/rllib/BUILD index 653d3b6a4..73b11ab2a 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -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"] ) diff --git a/rllib/agents/impala/tests/test_impala.py b/rllib/agents/impala/tests/test_impala.py index bbbd701da..79a753fc8 100644 --- a/rllib/agents/impala/tests/test_impala.py +++ b/rllib/agents/impala/tests/test_impala.py @@ -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() diff --git a/rllib/agents/impala/vtrace_tf_policy.py b/rllib/agents/impala/vtrace_tf_policy.py index edd78fa39..05d173738 100644 --- a/rllib/agents/impala/vtrace_tf_policy.py +++ b/rllib/agents/impala/vtrace_tf_policy.py @@ -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: diff --git a/rllib/agents/impala/vtrace_torch_policy.py b/rllib/agents/impala/vtrace_torch_policy.py index 067b40d4e..7a0a3cec1 100644 --- a/rllib/agents/impala/vtrace_torch_policy.py +++ b/rllib/agents/impala/vtrace_torch_policy.py @@ -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) diff --git a/rllib/examples/env/repeat_after_me_env.py b/rllib/examples/env/repeat_after_me_env.py index ae028034c..445320b52 100644 --- a/rllib/examples/env/repeat_after_me_env.py +++ b/rllib/examples/env/repeat_after_me_env.py @@ -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 = [] diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index 732aa13fa..cb2e7422d 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -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, diff --git a/rllib/models/tf/fcnet.py b/rllib/models/tf/fcnet.py index 745639a98..2b13eea18 100644 --- a/rllib/models/tf/fcnet.py +++ b/rllib/models/tf/fcnet.py @@ -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): diff --git a/rllib/models/tf/recurrent_net.py b/rllib/models/tf/recurrent_net.py index f1abca6b8..843efb430 100644 --- a/rllib/models/tf/recurrent_net.py +++ b/rllib/models/tf/recurrent_net.py @@ -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]) diff --git a/rllib/models/torch/fcnet.py b/rllib/models/torch/fcnet.py index 5dfa6ba86..f059eefb9 100644 --- a/rllib/models/torch/fcnet.py +++ b/rllib/models/torch/fcnet.py @@ -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) diff --git a/rllib/models/torch/recurrent_net.py b/rllib/models/torch/recurrent_net.py index 934f89402..7c5c91965 100644 --- a/rllib/models/torch/recurrent_net.py +++ b/rllib/models/torch/recurrent_net.py @@ -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]) diff --git a/rllib/models/torch/torch_modelv2.py b/rllib/models/torch/torch_modelv2.py index c232ecd83..8b50dfd1d 100644 --- a/rllib/models/torch/torch_modelv2.py +++ b/rllib/models/torch/torch_modelv2.py @@ -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): diff --git a/rllib/models/torch/visionnet.py b/rllib/models/torch/visionnet.py index 66bb04f78..253e1a35f 100644 --- a/rllib/models/torch/visionnet.py +++ b/rllib/models/torch/visionnet.py @@ -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") diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 142729f39..c5cede21a 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -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], diff --git a/rllib/tests/test_nested_observation_spaces.py b/rllib/tests/test_nested_observation_spaces.py index daf242ac5..18d7a9bda 100644 --- a/rllib/tests/test_nested_observation_spaces.py +++ b/rllib/tests/test_nested_observation_spaces.py @@ -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={ diff --git a/rllib/tests/test_rollout.py b/rllib/tests/test_rollout.py index e6433c50d..7e187a10f 100644 --- a/rllib/tests/test_rollout.py +++ b/rllib/tests/test_rollout.py @@ -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] diff --git a/rllib/train.py b/rllib/train.py index a6232e3b4..0932e08b2 100755 --- a/rllib/train.py +++ b/rllib/train.py @@ -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), diff --git a/rllib/tuned_examples/ppo/repeatafterme-ppo-lstm.yaml b/rllib/tuned_examples/ppo/repeatafterme-ppo-lstm.yaml new file mode 100644 index 000000000..e737cb09d --- /dev/null +++ b/rllib/tuned_examples/ppo/repeatafterme-ppo-lstm.yaml @@ -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] diff --git a/rllib/utils/test_utils.py b/rllib/utils/test_utils.py index b8ac00445..c88437098 100644 --- a/rllib/utils/test_utils.py +++ b/rllib/utils/test_utils.py @@ -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 "