From 2589309cf0beed938e54665e35a73bba87841f9b Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Sat, 20 Jun 2020 00:05:19 +0200 Subject: [PATCH] [RLlib] Make sure torch and tf behave the same wrt conv2d nets. (#8785) --- rllib/BUILD | 52 +------- rllib/models/tf/visionnet.py | 38 ++++-- rllib/models/torch/fcnet.py | 7 +- rllib/models/torch/misc.py | 6 +- rllib/models/torch/visionnet.py | 145 ++++++++++++++++++--- rllib/rollout.py | 1 + rllib/tests/test_attention_net_learning.py | 2 +- rllib/tests/test_supported_spaces.py | 25 ++-- rllib/train.py | 3 + rllib/tuned_examples/ars/cartpole-ars.yaml | 2 +- rllib/tuned_examples/es/cartpole-es.yaml | 2 +- 11 files changed, 179 insertions(+), 104 deletions(-) diff --git a/rllib/BUILD b/rllib/BUILD index 0674540e9..bf57d4e68 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -296,7 +296,7 @@ py_test( py_test( name = "run_regression_tests_repeat_after_me_torch", main = "tests/run_regression_tests.py", - tags = ["learning_tests_tf"], + tags = ["learning_tests_torch"], size = "medium", srcs = ["tests/run_regression_tests.py"], data = ["tuned_examples/ppo/repeatafterme-ppo-lstm.yaml"], @@ -899,19 +899,6 @@ py_test( ] ) -py_test( - name = "test_ppo_tf_cartpole_v1_lstm", - main = "train.py", srcs = ["train.py"], - tags = ["quick_train"], - args = [ - "--env", "CartPole-v1", - "--run", "PPO", - "--stop", "'{\"training_iteration\": 1}'", - "--config", "'{\"framework\": \"tf\", \"simple_optimizer\": false, \"num_sgd_iter\": 2, \"model\": {\"use_lstm\": true}}'", - "--ray-num-cpus", "4" - ] -) - py_test( name = "test_ppo_tf_cartpole_v1_lstm_simple_optimizer", main = "train.py", srcs = ["train.py"], @@ -925,19 +912,6 @@ py_test( ] ) -# TODO(sven): Fix LSTM auto-wrapping for torch models. This test case did not(!) exist in Jenkins. -#py_test( -# name = "test_ppo_torch_cartpole_v1_lstm_simple_optimizer", -# main = "train.py", srcs = ["train.py"], -# args = [ -# "--env", "CartPole-v1", -# "--run", "PPO", -# "--stop", "'{\"training_iteration\": 1}'", -# "--config", "'{\"framework\": \"torch\", \"simple_optimizer\": true, \"num_sgd_iter\": 2, \"model\": {\"use_lstm\": true}}'", -# "--ray-num-cpus", "4" -# ] -#) - py_test( name = "test_ppo_tf_cartpole_v1_complete_episode_batches", main = "train.py", srcs = ["train.py"], @@ -974,30 +948,6 @@ py_test( ] ) -py_test( - name = "test_ppo_tf_montezuma_revenge_v0", - main = "train.py", srcs = ["train.py"], - tags = ["quick_train"], - args = [ - "--env", "MontezumaRevenge-v0", - "--run", "PPO", - "--stop", "'{\"training_iteration\": 1}'", - "--config", "'{\"framework\": \"tf\", \"kl_coeff\": 1.0, \"num_sgd_iter\": 10, \"lr\": 1e-4, \"sgd_minibatch_size\": 64, \"train_batch_size\": 2000, \"num_workers\": 1, \"model\": {\"dim\": 40, \"conv_filters\": [[16, [8, 8], 4], [32, [4, 4], 2], [512, [5, 5], 1]]}}'" - ] -) - -py_test( - name = "test_ppo_torch_montezuma_revenge_v0", - main = "train.py", srcs = ["train.py"], - tags = ["quick_train"], - args = [ - "--env", "MontezumaRevenge-v0", - "--run", "PPO", - "--stop", "'{\"training_iteration\": 1}'", - "--config", "'{\"framework\": \"torch\", \"kl_coeff\": 1.0, \"num_sgd_iter\": 10, \"lr\": 1e-4, \"sgd_minibatch_size\": 64, \"train_batch_size\": 2000, \"num_workers\": 1, \"model\": {\"dim\": 40, \"conv_filters\": [[16, [8, 8], 4], [32, [4, 4], 2], [512, [5, 5], 1]]}}'" - ] -) - py_test( name = "test_appo_tf_pendulum_v0_no_gpus", main = "train.py", srcs = ["train.py"], diff --git a/rllib/models/tf/visionnet.py b/rllib/models/tf/visionnet.py index 4887d04b7..e4fd75a40 100644 --- a/rllib/models/tf/visionnet.py +++ b/rllib/models/tf/visionnet.py @@ -24,6 +24,9 @@ class VisionNetwork(TFModelV2): inputs = tf.keras.layers.Input( shape=obs_space.shape, name="observations") last_layer = inputs + # Whether the last layer is the output of a Flattened (rather than + # a n x (1,1) Conv2D). + self.last_layer_is_flattened = False # Build the action layers for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1): @@ -35,10 +38,11 @@ class VisionNetwork(TFModelV2): padding="same", data_format="channels_last", name="conv{}".format(i))(last_layer) + out_size, kernel, stride = filters[-1] # No final linear: Last layer is a Conv2D and uses num_outputs. - if no_final_linear: + if no_final_linear and num_outputs: last_layer = tf.keras.layers.Conv2D( num_outputs, kernel, @@ -59,12 +63,23 @@ class VisionNetwork(TFModelV2): padding="valid", data_format="channels_last", name="conv{}".format(i + 1))(last_layer) - conv_out = tf.keras.layers.Conv2D( - num_outputs, [1, 1], - activation=None, - padding="same", - data_format="channels_last", - name="conv_out")(last_layer) + + # num_outputs defined. Use that to create an exact + # `num_output`-sized (1,1)-Conv2D. + if num_outputs: + conv_out = tf.keras.layers.Conv2D( + num_outputs, [1, 1], + activation=None, + padding="same", + data_format="channels_last", + name="conv_out")(last_layer) + # num_outputs not known -> Flatten, then set self.num_outputs + # to the resulting number of nodes. + else: + self.last_layer_is_flattened = True + conv_out = tf.keras.layers.Flatten( + data_format="channels_last")(last_layer) + self.num_outputs = conv_out.shape[1] # Build the value layers if vf_share_layers: @@ -109,10 +124,15 @@ class VisionNetwork(TFModelV2): self.register_variables(self.base_model.variables) def forward(self, input_dict, state, seq_lens): - # explicit cast to float32 needed in eager + # Explicit cast to float32 needed in eager. model_out, self._value_out = self.base_model( tf.cast(input_dict["obs"], tf.float32)) - return tf.squeeze(model_out, axis=[1, 2]), state + # Our last layer is already flat. + if self.last_layer_is_flattened: + return model_out, state + # Last layer is a n x [1,1] Conv2D -> Flatten. + else: + return tf.squeeze(model_out, axis=[1, 2]), state 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 da100f711..109b282b5 100644 --- a/rllib/models/torch/fcnet.py +++ b/rllib/models/torch/fcnet.py @@ -90,17 +90,16 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): 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 = [] + vf_layers = [] for size in hiddens: - self._value_branch_separate.append( + vf_layers.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_separate = nn.Sequential(*vf_layers) self._value_branch = SlimFC( in_size=prev_layer_size, diff --git a/rllib/models/torch/misc.py b/rllib/models/torch/misc.py index 1679e4715..5913ff2b6 100644 --- a/rllib/models/torch/misc.py +++ b/rllib/models/torch/misc.py @@ -15,7 +15,7 @@ def normc_initializer(std=1.0): return initializer -def valid_padding(in_size, filter_size, stride_size): +def same_padding(in_size, filter_size, stride_size): """Note: Padding is added to match TF conv2d `same` padding. See www.tensorflow.org/versions/r0.12/api_docs/python/nn/convolution @@ -25,8 +25,8 @@ def valid_padding(in_size, filter_size, stride_size): filter_size (tuple): Rows (Height), Column (Width) for filter Output: - padding (tuple): For input into torch.nn.ZeroPad2d - output (tuple): Output shape after padding and convolution + padding (tuple): For input into torch.nn.ZeroPad2d. + output (tuple): Output shape after padding and convolution. """ in_height, in_width = in_size filter_height, filter_width = filter_size diff --git a/rllib/models/torch/visionnet.py b/rllib/models/torch/visionnet.py index c3aa3d51c..31e262cc7 100644 --- a/rllib/models/torch/visionnet.py +++ b/rllib/models/torch/visionnet.py @@ -1,5 +1,7 @@ +import numpy as np + from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 -from ray.rllib.models.torch.misc import normc_initializer, valid_padding, \ +from ray.rllib.models.torch.misc import normc_initializer, same_padding, \ SlimConv2d, SlimFC from ray.rllib.models.tf.visionnet_v1 import _get_filter_config from ray.rllib.utils.annotations import override @@ -22,15 +24,19 @@ class VisionNetwork(TorchModelV2, nn.Module): filters = model_config.get("conv_filters") if not filters: filters = _get_filter_config(obs_space.shape) - # no_final_linear = model_config.get("no_final_linear") - # vf_share_layers = model_config.get("vf_share_layers") + no_final_linear = model_config.get("no_final_linear") + vf_share_layers = model_config.get("vf_share_layers") + + # Whether the last layer is the output of a Flattened (rather than + # a n x (1,1) Conv2D). + self.last_layer_is_flattened = False + self._logits = None layers = [] (w, h, in_channels) = obs_space.shape in_size = [w, h] for out_channels, kernel, stride in filters[:-1]: - padding, out_size = valid_padding(in_size, kernel, - [stride, stride]) + padding, out_size = same_padding(in_size, kernel, [stride, stride]) layers.append( SlimConv2d( in_channels, @@ -43,33 +49,130 @@ class VisionNetwork(TorchModelV2, nn.Module): in_size = out_size out_channels, kernel, stride = filters[-1] - layers.append( - SlimConv2d( - in_channels, - out_channels, - kernel, - stride, - None, - activation_fn=activation)) + + # No final linear: Last layer is a Conv2D and uses num_outputs. + if no_final_linear and num_outputs: + layers.append( + SlimConv2d( + in_channels, + num_outputs, + kernel, + stride, + None, # padding=valid + activation_fn=activation)) + out_channels = num_outputs + # Finish network normally (w/o overriding last layer size with + # `num_outputs`), then add another linear one of size `num_outputs`. + else: + layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + None, # padding=valid + activation_fn=activation)) + + # num_outputs defined. Use that to create an exact + # `num_output`-sized (1,1)-Conv2D. + if num_outputs: + in_size = [ + np.ceil((in_size[0] - kernel[0]) / stride), + np.ceil((in_size[1] - kernel[1]) / stride) + ] + padding, _ = same_padding(in_size, [1, 1], [1, 1]) + self._logits = SlimConv2d( + out_channels, + num_outputs, [1, 1], + 1, + padding, + activation_fn=None) + # num_outputs not known -> Flatten, then set self.num_outputs + # to the resulting number of nodes. + else: + self.last_layer_is_flattened = True + layers.append(nn.Flatten()) + self.num_outputs = out_channels + self._convs = nn.Sequential(*layers) - self._logits = SlimFC( - out_channels, num_outputs, initializer=nn.init.xavier_uniform_) - self._value_branch = SlimFC( - out_channels, 1, initializer=normc_initializer()) + # Build the value layers + self._value_branch_separate = self._value_branch = None + if vf_share_layers: + self._value_branch = SlimFC( + out_channels, 1, initializer=normc_initializer(0.01)) + else: + vf_layers = [] + (w, h, in_channels) = obs_space.shape + in_size = [w, h] + for out_channels, kernel, stride in filters[:-1]: + padding, out_size = same_padding(in_size, kernel, + [stride, stride]) + vf_layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + padding, + activation_fn=activation)) + in_channels = out_channels + in_size = out_size + + out_channels, kernel, stride = filters[-1] + vf_layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + None, + activation_fn=activation)) + + vf_layers.append( + SlimConv2d( + in_channels=out_channels, + out_channels=1, + kernel=1, + stride=1, + padding=None)) + self._value_branch_separate = nn.Sequential(*vf_layers) + # Holds the current "base" output (before logits layer). self._features = None @override(TorchModelV2) def forward(self, input_dict, state, seq_lens): - self._features = self._hidden_layers(input_dict["obs"].float()) - logits = self._logits(self._features) - return logits, state + self._features = input_dict["obs"].float().permute(0, 3, 1, 2) + conv_out = self._convs(self._features) + # Store features to save forward pass when getting value_function out. + if not self._value_branch_separate: + self._features = conv_out + + if not self.last_layer_is_flattened: + if self._logits: + conv_out = self._logits(conv_out) + logits = conv_out.squeeze(3) + logits = logits.squeeze(2) + return logits, state + else: + return conv_out, state @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: + value = self._value_branch_separate(self._features) + value = value.squeeze(3) + value = value.squeeze(2) + return value.squeeze(1) + else: + if not self.last_layer_is_flattened: + features = self._features.squeeze(3) + features = features.squeeze(2) + else: + features = self._features + return self._value_branch(features).squeeze(1) def _hidden_layers(self, obs): res = self._convs(obs.permute(0, 3, 1, 2)) # switch to channel-major diff --git a/rllib/rollout.py b/rllib/rollout.py index 7f0b8dc48..87ff510d6 100755 --- a/rllib/rollout.py +++ b/rllib/rollout.py @@ -306,6 +306,7 @@ def run(args, parser): save_info=args.save_info) as saver: rollout(agent, args.env, num_steps, num_episodes, saver, args.no_render, video_dir) + agent.stop() class DefaultMapping(collections.defaultdict): diff --git a/rllib/tests/test_attention_net_learning.py b/rllib/tests/test_attention_net_learning.py index f63514f35..23e6ac96e 100644 --- a/rllib/tests/test_attention_net_learning.py +++ b/rllib/tests/test_attention_net_learning.py @@ -17,7 +17,7 @@ class TestAttentionNetLearning(unittest.TestCase): } stop = { - "episode_reward_mean": 180.0, + "episode_reward_mean": 150.0, "timesteps_total": 5000000, } diff --git a/rllib/tests/test_supported_spaces.py b/rllib/tests/test_supported_spaces.py index b7e128924..313339104 100644 --- a/rllib/tests/test_supported_spaces.py +++ b/rllib/tests/test_supported_spaces.py @@ -66,11 +66,14 @@ def check_support(alg, config, train=True, check_bounds=False, tfe=False): p_done=1.0, check_action_bounds=check_bounds))) stat = "ok" - a = None + if alg == "SAC": + config["use_state_preprocessor"] = o_name in ["atari", "image"] + try: - if alg == "SAC": - config["use_state_preprocessor"] = o_name in ["atari", "image"] a = get_agent_class(alg)(config=config, env=RandomEnv) + except UnsupportedSpaceException: + stat = "unsupported" + else: if alg not in ["DDPG", "ES", "ARS", "SAC"]: if o_name in ["atari", "image"]: if fw == "torch": @@ -85,18 +88,14 @@ def check_support(alg, config, train=True, check_bounds=False, tfe=False): assert isinstance(a.get_policy().model, FCNetV2) if train: a.train() - except UnsupportedSpaceException: - stat = "unsupported" - finally: - if a: - try: - a.stop() - except Exception as e: - print("Ignoring error stopping agent", e) - pass + try: + a.stop() + except Exception as e: + print("Ignoring error stopping agent", e) + pass print(stat) - frameworks = ("tf", "torch") + frameworks = ("torch", "tf") if tfe: frameworks += ("tfe", ) for _ in framework_iterator(config, frameworks=frameworks): diff --git a/rllib/train.py b/rllib/train.py index 4aa80a8a4..c7885f5a0 100755 --- a/rllib/train.py +++ b/rllib/train.py @@ -217,6 +217,7 @@ def run(args, parser): num_cpus=args.ray_num_cpus, num_gpus=args.ray_num_gpus, local_mode=args.local_mode) + run_experiments( experiments, scheduler=_make_scheduler(args), @@ -225,6 +226,8 @@ def run(args, parser): verbose=verbose, concurrent=True) + ray.shutdown() + if __name__ == "__main__": parser = create_parser() diff --git a/rllib/tuned_examples/ars/cartpole-ars.yaml b/rllib/tuned_examples/ars/cartpole-ars.yaml index 67dc4a88e..a22496322 100644 --- a/rllib/tuned_examples/ars/cartpole-ars.yaml +++ b/rllib/tuned_examples/ars/cartpole-ars.yaml @@ -3,7 +3,7 @@ cartpole-ars: run: ARS stop: episode_reward_mean: 150 - timesteps_total: 500000 + timesteps_total: 1000000 config: # Works for both torch and tf. framework: tf diff --git a/rllib/tuned_examples/es/cartpole-es.yaml b/rllib/tuned_examples/es/cartpole-es.yaml index 57d5ec3a3..dcafa00ba 100644 --- a/rllib/tuned_examples/es/cartpole-es.yaml +++ b/rllib/tuned_examples/es/cartpole-es.yaml @@ -3,7 +3,7 @@ cartpole-es: run: ES stop: episode_reward_mean: 150 - timesteps_total: 500000 + timesteps_total: 1000000 config: # Works for both torch and tf. framework: tf