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

[RLlib] Make sure torch and tf behave the same wrt conv2d nets. (#8785)

Browse Source
This commit is contained in:
Sven Mika
2020-06-20 00:05:19 +02:00
committed by GitHub
parent ca66f88b96
commit 2589309cf0
11 changed files with 179 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/BUILD
+1 -51
View File
@@ -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"],
rllib/models/tf/visionnet.py
+29 -9
View File
@@ -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])
rllib/models/torch/fcnet.py
+3 -4
View File
@@ -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,
rllib/models/torch/misc.py
+3 -3
View File
@@ -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
rllib/models/torch/visionnet.py
+124 -21
View File
@@ -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
rllib/rollout.py
+1
View File
@@ -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):
rllib/tests/test_attention_net_learning.py
+1 -1
View File
@@ -17,7 +17,7 @@ class TestAttentionNetLearning(unittest.TestCase):
}
stop = {
"episode_reward_mean": 180.0,
"episode_reward_mean": 150.0,
"timesteps_total": 5000000,
}
rllib/tests/test_supported_spaces.py
+12 -13
View File
@@ -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):
rllib/train.py
+3
View File
@@ -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()
rllib/tuned_examples/ars/cartpole-ars.yaml
+1 -1
View File
@@ -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
rllib/tuned_examples/es/cartpole-es.yaml
+1 -1
View File
@@ -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