From b197c0c4044f66628c6672fe78581768a54d0e59 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Wed, 27 Jun 2018 22:51:04 -0700 Subject: [PATCH] [rllib] General RNN support (#2299) * wip * cls * re * wip * wip * a3c working * torch support * pg works * lint * rm v2 * consumer id * clean up pg * clean up more * fix python 2.7 * tf session management * docs * dqn wip * fix compile * dqn * apex runs * up * impotrs * ddpg * quotes * fix tests * fix last r * fix tests * lint * pass checkpoint restore * kwar * nits * policy graph * fix yapf * com * class * pyt * vectorization * update * test cpe * unit test * fix ddpg2 * changes * wip * args * faster test * common * fix * add alg option * batch mode and policy serving * multi serving test * todo * wip * serving test * doc async env * num envs * comments * thread * remove init hook * update * fix ppo * comments1 * fix * updates * add jenkins tests * fix * fix pytorch * fix * fixes * fix a3c policy * fix squeeze * fix trunc on apex * fix squeezing for real * update * remove horizon test for now * multiagent wip * update * fix race condition * fix ma * t * doc * st * wip * example * wip * working * cartpole * wip * batch wip * fix bug * make other_batches None default * working * debug * nit * warn * comments * fix ppo * fix obs filter * update * wip * tf * update * fix * cleanup * cleanup * spacing * model * fix * dqn * fix ddpg * doc * keep names * update * fix * com * docs * clarify model outputs * Update torch_policy_graph.py * fix obs filter * pass thru worker index * fix * rename * vlad torch comments * fix log action * debug name * fix lstm * remove unused ddpg net * remove conv net * revert lstm * wip * wip * cast * wip * works * fix a3c * works * lstm util test * doc * clean up * update * fix lstm check * move to end * fix sphinx * fix cmd * remove bad doc * clarify * copy * async sa * fix * comments * fix a3c conf * tune lstm * fix reshape * fix * back to 16 * tuned a3c update * update * tuned * optional * fix catalog * remove prep --- python/ray/rllib/a3c/a3c.py | 6 +- python/ray/rllib/a3c/a3c_tf_policy.py | 4 +- python/ray/rllib/models/catalog.py | 37 ++-- python/ray/rllib/models/fcnet.py | 10 +- python/ray/rllib/models/lstm.py | 183 +++++++++++++++--- python/ray/rllib/models/misc.py | 4 + python/ray/rllib/models/model.py | 34 ++-- python/ray/rllib/models/multiagentfcnet.py | 3 +- python/ray/rllib/models/pytorch/fcnet.py | 2 +- python/ray/rllib/models/pytorch/model.py | 4 +- python/ray/rllib/models/pytorch/visionnet.py | 2 +- python/ray/rllib/models/visionnet.py | 43 +++- python/ray/rllib/pg/pg.py | 2 +- python/ray/rllib/pg/pg_policy_graph.py | 22 ++- python/ray/rllib/test/test_catalog.py | 5 +- python/ray/rllib/test/test_lstm.py | 43 ++++ python/ray/rllib/test/test_multi_agent_env.py | 10 +- .../tuned_examples/pong-a3c-pytorch.yaml | 1 - python/ray/rllib/tuned_examples/pong-a3c.yaml | 14 +- python/ray/rllib/tuned_examples/pong-dqn.yaml | 5 - python/ray/rllib/utils/tf_policy_graph.py | 57 ++++-- test/jenkins_tests/run_multi_node_tests.sh | 13 ++ 22 files changed, 388 insertions(+), 116 deletions(-) create mode 100644 python/ray/rllib/test/test_lstm.py diff --git a/python/ray/rllib/a3c/a3c.py b/python/ray/rllib/a3c/a3c.py index 8af5eb82c..e50a1a04c 100644 --- a/python/ray/rllib/a3c/a3c.py +++ b/python/ray/rllib/a3c/a3c.py @@ -24,8 +24,6 @@ DEFAULT_CONFIG = { "use_pytorch": False, # Which observation filter to apply to the observation "observation_filter": "NoFilter", - # Which reward filter to apply to the reward - "reward_filter": "NoFilter", # Discount factor of MDP "gamma": 0.99, # GAE(gamma) parameter @@ -44,8 +42,10 @@ DEFAULT_CONFIG = { "summarize": False, # Model and preprocessor options "model": { - # Use LSTM model - only applicable for image states. Requires TF. + # Use LSTM model. Requires TF. "use_lstm": False, + # Max seq length for LSTM training. + "max_seq_len": 20, # (Image statespace) - Converts image to Channels = 1 "grayscale": True, # (Image statespace) - Each pixel diff --git a/python/ray/rllib/a3c/a3c_tf_policy.py b/python/ray/rllib/a3c/a3c_tf_policy.py index cc7605965..9657d9b05 100644 --- a/python/ray/rllib/a3c/a3c_tf_policy.py +++ b/python/ray/rllib/a3c/a3c_tf_policy.py @@ -83,7 +83,9 @@ class A3CPolicyGraph(TFPolicyGraph): obs_input=self.observations, action_sampler=action_dist.sample(), loss=self.loss.total_loss, loss_inputs=loss_in, is_training=is_training, state_inputs=self.state_in, - state_outputs=self.state_out) + state_outputs=self.state_out, + seq_lens=self.model.seq_lens, + max_seq_len=self.config["model"]["max_seq_len"]) if self.config.get("summarize"): bs = tf.to_float(tf.shape(self.observations)[0]) diff --git a/python/ray/rllib/models/catalog.py b/python/ray/rllib/models/catalog.py index 7d67e591e..61b31f708 100644 --- a/python/ray/rllib/models/catalog.py +++ b/python/ray/rllib/models/catalog.py @@ -22,17 +22,24 @@ from ray.rllib.models.multiagentfcnet import MultiAgentFullyConnectedNetwork MODEL_CONFIGS = [ # === Built-in options === - "conv_filters", # Number of filters + "conv_filters", # Filter configuration + "conv_activation", # Nonlinearity for built-in convnet + + "fcnet_activation", # Nonlinearity for fully connected net (tanh, relu) + "fcnet_hiddens", # Number of hidden layers for fully connected net + "dim", # Dimension for ATARI "grayscale", # Converts ATARI frame to 1 Channel Grayscale image "zero_mean", # Changes frame to range from [-1, 1] if true "extra_frameskip", # (int) for number of frames to skip - "fcnet_activation", # Nonlinearity for fully connected net (tanh, relu) - "fcnet_hiddens", # Number of hidden layers for fully connected net + "free_log_std", # Documented in ray.rllib.models.Model "channel_major", # Pytorch conv requires images to be channel-major "squash_to_range", # Whether to squash the action output to space range - "use_lstm", # Whether to use a LSTM model + + "use_lstm", # Whether to wrap the model with a LSTM + "max_seq_len", # Max seq len for training the LSTM, defaults to 20 + "lstm_cell_size", # Size of the LSTM cell # === Options for custom models === "custom_preprocessor", # Name of a custom preprocessor to use @@ -113,9 +120,9 @@ class ModelCatalog(object): if isinstance(action_space, gym.spaces.Box): return tf.placeholder( - tf.float32, shape=(None, action_space.shape[0])) + tf.float32, shape=(None, action_space.shape[0]), name="action") elif isinstance(action_space, gym.spaces.Discrete): - return tf.placeholder(tf.int64, shape=(None,)) + return tf.placeholder(tf.int64, shape=(None,), name="action") elif isinstance(action_space, gym.spaces.Tuple): size = 0 all_discrete = True @@ -126,13 +133,14 @@ class ModelCatalog(object): all_discrete = False size += np.product(action_space.spaces[i].shape) return tf.placeholder( - tf.int64 if all_discrete else tf.float32, shape=(None, size)) + tf.int64 if all_discrete else tf.float32, shape=(None, size), + name="action") else: raise NotImplementedError("action space {}" " not supported".format(action_space)) @staticmethod - def get_model(inputs, num_outputs, options={}): + def get_model(inputs, num_outputs, options=None): """Returns a suitable model conforming to given input and output specs. Args: @@ -144,15 +152,22 @@ class ModelCatalog(object): model (Model): Neural network model. """ + options = options or {} + model = ModelCatalog._get_model(inputs, num_outputs, options) + + if options.get("use_lstm"): + model = LSTM(model.last_layer, num_outputs, options) + + return model + + @staticmethod + def _get_model(inputs, num_outputs, options): if "custom_model" in options: model = options["custom_model"] print("Using custom model {}".format(model)) return _global_registry.get(RLLIB_MODEL, model)( inputs, num_outputs, options) - if options.get("use_lstm"): - return LSTM(inputs, num_outputs, options) - obs_rank = len(inputs.shape) - 1 # num_outputs > 1 used to avoid hitting this with the value function diff --git a/python/ray/rllib/models/fcnet.py b/python/ray/rllib/models/fcnet.py index ab40a6c6b..ce516f1c7 100644 --- a/python/ray/rllib/models/fcnet.py +++ b/python/ray/rllib/models/fcnet.py @@ -6,20 +6,16 @@ import tensorflow as tf import tensorflow.contrib.slim as slim from ray.rllib.models.model import Model -from ray.rllib.models.misc import normc_initializer +from ray.rllib.models.misc import normc_initializer, get_activation_fn class FullyConnectedNetwork(Model): """Generic fully connected network.""" - def _init(self, inputs, num_outputs, options): + def _build_layers(self, inputs, num_outputs, options): hiddens = options.get("fcnet_hiddens", [256, 256]) - fcnet_activation = options.get("fcnet_activation", "tanh") - if fcnet_activation == "tanh": - activation = tf.nn.tanh - elif fcnet_activation == "relu": - activation = tf.nn.relu + activation = get_activation_fn(options.get("fcnet_activation", "tanh")) with tf.name_scope("fc_net"): i = 1 diff --git a/python/ray/rllib/models/lstm.py b/python/ray/rllib/models/lstm.py index 14d8a9371..304f3470e 100644 --- a/python/ray/rllib/models/lstm.py +++ b/python/ray/rllib/models/lstm.py @@ -2,56 +2,175 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +"""LSTM support for RLlib. + +The main trick here is that we add the time dimension at the last moment. +The non-LSTM layers of the model see their inputs as one flat batch. Before +the LSTM cell, we reshape the input to add the expected time dimension. During +postprocessing, we dynamically pad the experience batches so that this +reshaping is possible. + +See the add_time_dimension() and chop_into_sequences() functions below for +more info. +""" + + import numpy as np import tensorflow as tf import tensorflow.contrib.rnn as rnn import distutils.version -from ray.rllib.models.misc import (conv2d, linear, flatten, - normc_initializer) +from ray.rllib.models.misc import linear, normc_initializer from ray.rllib.models.model import Model -class LSTM(Model): - """Vision LSTM network based here: - https://github.com/openai/universe-starter-agent""" +def add_time_dimension(padded_inputs, seq_lens): + """Adds a time dimension to padded inputs. - # TODO(rliaw): Add LSTM code for other algorithms - def _init(self, inputs, num_outputs, options): + Arguments: + padded_inputs (Tensor): a padded batch of sequences. That is, + for seq_lens=[1, 2, 2], then inputs=[A, *, B, B, C, C], where + A, B, C are sequence elements and * denotes padding. + seq_lens (Tensor): the sequence lengths within the input batch, + suitable for passing to tf.nn.dynamic_rnn(). + + Returns: + Reshaped tensor of shape [NUM_SEQUENCES, MAX_SEQ_LEN, ...]. + """ + + # Sequence lengths have to be specified for LSTM batch inputs. The + # input batch must be padded to the max seq length given here. That is, + # batch_size == len(seq_lens) * max(seq_lens) + max_seq_len = tf.reduce_max(seq_lens) + padded_batch_size = tf.shape(padded_inputs)[0] + + # Dynamically reshape the padded batch to introduce a time dimension. + new_batch_size = padded_batch_size // max_seq_len + new_shape = ( + [new_batch_size, max_seq_len] + + padded_inputs.get_shape().as_list()[1:]) + return tf.reshape(padded_inputs, new_shape) + + +def chop_into_sequences( + time_column, feature_columns, state_columns, max_seq_len): + """Truncate and pad experiences into fixed-length sequences. + + Arguments: + time_column (list): Timesteps per feature / state. This contains + sequences of monotonically increasing step values, e.g., + [0, 1, 2, 0, 1, 2, 3, 4, 5, 0, 1, 2]. + feature_columns (list): List of arrays containing features. + state_columns (list): List of arrays containing LSTM state values. + max_seq_len (int): Max length of sequences before truncation. + + Returns: + f_pad (list): Padded feature columns. These will be of shape + [NUM_SEQUENCES * MAX_SEQ_LEN, ...]. + s_init (list): Initial states for each sequence, of shape + [NUM_SEQUENCES, ...]. + seq_lens (list): List of sequence lengths, of shape [NUM_SEQUENCES]. + + Examples: + >>> f_pad, s_init, seq_lens = chop_into_sequences( + time_column=[0, 1, 0, 1, 2, 3], + feature_columns=[[4, 4, 8, 8, 8, 8], + [1, 1, 0, 1, 1, 0]], + state_columns=[[4, 5, 4, 5, 5, 5]], + max_seq_len=3) + >>> print(f_pad) + [[4, 4, 0, 8, 8, 8, 8, 0, 0], + [1, 1, 0, 0, 1, 1, 0, 0, 0]] + >>> print(s_init) + [[4, 4, 5]] + >>> print(seq_lens) + [2, 3, 1] + """ + + prev_t = -1 + seq_lens = [] + seq_len = 0 + for t in time_column: + if t <= prev_t or seq_len >= max_seq_len: + seq_lens.append(seq_len) + seq_len = 0 + seq_len += 1 + prev_t = t + if seq_len: + seq_lens.append(seq_len) + assert sum(seq_lens) == len(time_column) + + # Dynamically shrink max len as needed to optimize memory usage + max_seq_len = max(seq_lens) + + feature_sequences = [] + for f in feature_columns: + f = np.array(f) + f_pad = np.zeros((len(seq_lens) * max_seq_len,) + np.shape(f)[1:]) + seq_base = 0 + i = 0 + for l in seq_lens: + for seq_offset in range(l): + f_pad[seq_base + seq_offset] = f[i] + i += 1 + seq_base += max_seq_len + assert i == len(time_column), f + feature_sequences.append(f_pad) + + initial_states = [] + for s in state_columns: + s = np.array(s) + s_init = [] + i = 0 + for l in seq_lens: + s_init.append(s[i]) + i += l + initial_states.append(np.array(s_init)) + + return feature_sequences, initial_states, np.array(seq_lens) + + +class LSTM(Model): + """Adds a LSTM cell on top of some other model output. + + Uses a linear layer at the end for output. + + Important: we assume inputs is a padded batch of sequences denoted by + self.seq_lens. See add_time_dimension() for more information. + """ + + def _build_layers(self, inputs, num_outputs, options): + cell_size = options.get("lstm_cell_size", 256) use_tf100_api = (distutils.version.LooseVersion(tf.VERSION) >= distutils.version.LooseVersion("1.0.0")) + last_layer = add_time_dimension(inputs, self.seq_lens) - self.x = x = inputs - for i in range(4): - x = tf.nn.elu(conv2d(x, 32, "l{}".format(i + 1), [3, 3], [2, 2])) - # Introduce a "fake" batch dimension of 1 after flatten so that we can - # do LSTM over the time dim. - x = tf.expand_dims(flatten(x), [0]) - - size = 256 + # Setup the LSTM cell if use_tf100_api: - lstm = rnn.BasicLSTMCell(size, state_is_tuple=True) + lstm = rnn.BasicLSTMCell(cell_size, state_is_tuple=True) else: - lstm = rnn.rnn_cell.BasicLSTMCell(size, state_is_tuple=True) - step_size = tf.shape(self.x)[:1] + lstm = rnn.rnn_cell.BasicLSTMCell(cell_size, state_is_tuple=True) + self.state_init = [ + np.zeros(lstm.state_size.c, np.float32), + np.zeros(lstm.state_size.h, np.float32)] - c_init = np.zeros(lstm.state_size.c, np.float32) - h_init = np.zeros(lstm.state_size.h, np.float32) - self.state_init = [c_init, h_init] - c_in = tf.placeholder(tf.float32, [1, lstm.state_size.c]) - h_in = tf.placeholder(tf.float32, [1, lstm.state_size.h]) + # Setup LSTM inputs + c_in = tf.placeholder(tf.float32, [None, lstm.state_size.c], name="c") + h_in = tf.placeholder(tf.float32, [None, lstm.state_size.h], name="h") self.state_in = [c_in, h_in] + # Setup LSTM outputs if use_tf100_api: state_in = rnn.LSTMStateTuple(c_in, h_in) else: state_in = rnn.rnn_cell.LSTMStateTuple(c_in, h_in) - lstm_out, lstm_state = tf.nn.dynamic_rnn(lstm, x, - initial_state=state_in, - sequence_length=step_size, - time_major=False) - lstm_c, lstm_h = lstm_state - x = tf.reshape(lstm_out, [-1, size]) - logits = linear(x, num_outputs, "action", normc_initializer(0.01)) - self.state_out = [lstm_c[:1, :], lstm_h[:1, :]] - return logits, x + lstm_out, lstm_state = tf.nn.dynamic_rnn( + lstm, last_layer, initial_state=state_in, + sequence_length=self.seq_lens, time_major=False) + self.state_out = list(lstm_state) + + # Compute outputs + last_layer = tf.reshape(lstm_out, [-1, cell_size]) + logits = linear( + last_layer, num_outputs, "action", normc_initializer(0.01)) + return logits, last_layer diff --git a/python/ray/rllib/models/misc.py b/python/ray/rllib/models/misc.py index a531bc07b..461296ecd 100644 --- a/python/ray/rllib/models/misc.py +++ b/python/ray/rllib/models/misc.py @@ -14,6 +14,10 @@ def normc_initializer(std=1.0): return _initializer +def get_activation_fn(name): + return getattr(tf.nn, name) + + def conv2d(x, num_filters, name, filter_size=(3, 3), stride=(1, 1), pad="SAME", dtype=tf.float32, collections=None): with tf.variable_scope(name): diff --git a/python/ray/rllib/models/model.py b/python/ray/rllib/models/model.py index f9081751c..278fd887f 100644 --- a/python/ray/rllib/models/model.py +++ b/python/ray/rllib/models/model.py @@ -15,6 +15,19 @@ class Model(object): The last layer of the network can also be retrieved if the algorithm needs to further post-processing (e.g. Actor and Critic networks in A3C). + Attributes: + inputs (Tensor): The input placeholder for this model, of shape + [BATCH_SIZE, ...]. + outputs (Tensor): The output vector of this model, of shape + [BATCH_SIZE, num_outputs]. + last_layer (Tensor): The network layer right before the model output, + of shape [BATCH_SIZE, N]. + state_init (list): List of initial recurrent state tensors (if any). + state_in (list): List of input recurrent state tensors (if any). + state_out (list): List of output recurrent state tensors (if any). + seq_lens (Tensor): The tensor input for RNN sequence lengths. This + defaults to a Tensor of [1] * len(batch) in the non-RNN case. + If `options["free_log_std"]` is True, the last half of the output layer will be free variables that are not dependent on inputs. This is often used if the output of the network is used @@ -22,25 +35,24 @@ class Model(object): first half of the parameters can be interpreted as a location parameter (like a mean) and the second half can be interpreted as a scale parameter (like a standard deviation). - - Attributes: - inputs (Tensor): The input placeholder for this model. - outputs (Tensor): The output vector of this model. - last_layer (Tensor): The network layer right before the model output. - state_init (list): List of initial recurrent state tensors (if any). - state_in (list): List of input recurrent state tensors (if any). - state_out (list): List of output recurrent state tensors (if any). """ def __init__(self, inputs, num_outputs, options): + self.inputs = inputs + + # Default attribute values for the non-RNN case self.state_init = [] self.state_in = [] self.state_out = [] - self.inputs = inputs + self.seq_lens = tf.placeholder_with_default( + tf.ones( # reshape needed for older tf versions + tf.reshape(tf.shape(inputs)[0], [1]), dtype=tf.int32), + [None], name="seq_lens") + if options.get("free_log_std", False): assert num_outputs % 2 == 0 num_outputs = num_outputs // 2 - self.outputs, self.last_layer = self._init( + self.outputs, self.last_layer = self._build_layers( inputs, num_outputs, options) if options.get("free_log_std", False): log_std = tf.get_variable(name="log_std", shape=[num_outputs], @@ -48,6 +60,6 @@ class Model(object): self.outputs = tf.concat( [self.outputs, 0.0 * self.outputs + log_std], 1) - def _init(self): + def _build_layers(self): """Builds and returns the output and last layer of the network.""" raise NotImplementedError diff --git a/python/ray/rllib/models/multiagentfcnet.py b/python/ray/rllib/models/multiagentfcnet.py index daf88a7f2..81d9c8d15 100644 --- a/python/ray/rllib/models/multiagentfcnet.py +++ b/python/ray/rllib/models/multiagentfcnet.py @@ -12,8 +12,7 @@ from ray.rllib.models.action_dist import Reshaper class MultiAgentFullyConnectedNetwork(Model): """Multiagent fully connected network.""" - def _init(self, inputs, num_outputs, options): - + def _build_layers(self, inputs, num_outputs, options): # Split the input and output tensors input_shapes = options["custom_options"]["multiagent_obs_shapes"] output_shapes = options["custom_options"]["multiagent_act_shapes"] diff --git a/python/ray/rllib/models/pytorch/fcnet.py b/python/ray/rllib/models/pytorch/fcnet.py index 8f5223b8c..e8f50da2f 100644 --- a/python/ray/rllib/models/pytorch/fcnet.py +++ b/python/ray/rllib/models/pytorch/fcnet.py @@ -10,7 +10,7 @@ import torch.nn as nn class FullyConnectedNetwork(Model): """TODO(rliaw): Logits, Value should both be contained here""" - def _init(self, inputs, num_outputs, options): + def _build_layers(self, inputs, num_outputs, options): assert type(inputs) is int hiddens = options.get("fcnet_hiddens", [256, 256]) fcnet_activation = options.get("fcnet_activation", "tanh") diff --git a/python/ray/rllib/models/pytorch/model.py b/python/ray/rllib/models/pytorch/model.py index 876196741..b25428a09 100644 --- a/python/ray/rllib/models/pytorch/model.py +++ b/python/ray/rllib/models/pytorch/model.py @@ -8,9 +8,9 @@ import torch.nn as nn class Model(nn.Module): def __init__(self, obs_space, ac_space, options): super(Model, self).__init__() - self._init(obs_space, ac_space, options) + self._build_layers(obs_space, ac_space, options) - def _init(self, inputs, num_outputs, options): + def _build_layers(self, inputs, num_outputs, options): raise NotImplementedError def forward(self, obs): diff --git a/python/ray/rllib/models/pytorch/visionnet.py b/python/ray/rllib/models/pytorch/visionnet.py index 825a5a69f..067e1659e 100644 --- a/python/ray/rllib/models/pytorch/visionnet.py +++ b/python/ray/rllib/models/pytorch/visionnet.py @@ -11,7 +11,7 @@ from ray.rllib.models.pytorch.misc import normc_initializer, valid_padding class VisionNetwork(Model): """Generic vision network""" - def _init(self, inputs, num_outputs, options): + def _build_layers(self, inputs, num_outputs, options): """TF visionnet in PyTorch. Params: diff --git a/python/ray/rllib/models/visionnet.py b/python/ray/rllib/models/visionnet.py index 198f40762..893f7acd2 100644 --- a/python/ray/rllib/models/visionnet.py +++ b/python/ray/rllib/models/visionnet.py @@ -6,25 +6,50 @@ import tensorflow as tf import tensorflow.contrib.slim as slim from ray.rllib.models.model import Model +from ray.rllib.models.misc import get_activation_fn, flatten class VisionNetwork(Model): """Generic vision network.""" - def _init(self, inputs, num_outputs, options): - filters = options.get("conv_filters", [ - [16, [8, 8], 4], - [32, [4, 4], 2], - [512, [10, 10], 1], - ]) + def _build_layers(self, inputs, num_outputs, options): + filters = options.get("conv_filters") + if not filters: + filters = get_filter_config(options) + + activation = get_activation_fn(options.get("conv_activation", "relu")) + with tf.name_scope("vision_net"): for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1): inputs = slim.conv2d( inputs, out_size, kernel, stride, - scope="conv{}".format(i)) + activation_fn=activation, scope="conv{}".format(i)) out_size, kernel, stride = filters[-1] fc1 = slim.conv2d( - inputs, out_size, kernel, stride, padding="VALID", scope="fc1") + inputs, out_size, kernel, stride, + activation_fn=activation, padding="VALID", scope="fc1") fc2 = slim.conv2d(fc1, num_outputs, [1, 1], activation_fn=None, normalizer_fn=None, scope="fc2") - return tf.squeeze(fc2, [1, 2]), tf.squeeze(fc1, [1, 2]) + return flatten(fc2), flatten(fc1) + + +def get_filter_config(options): + filters_80x80 = [ + [16, [8, 8], 4], + [32, [4, 4], 2], + [512, [10, 10], 1], + ] + filters_42x42 = [ + [16, [4, 4], 2], + [32, [4, 4], 2], + [512, [11, 11], 1], + ] + dim = options.get("dim", 80) + if dim == 80: + return filters_80x80 + elif dim == 42: + return filters_42x42 + else: + raise ValueError( + "No default configuration for image size={}".format(dim) + + ", you must specify `conv_filters` manually as a model option.") diff --git a/python/ray/rllib/pg/pg.py b/python/ray/rllib/pg/pg.py index 1926d4c1c..e64af9fe6 100644 --- a/python/ray/rllib/pg/pg.py +++ b/python/ray/rllib/pg/pg.py @@ -26,7 +26,7 @@ DEFAULT_CONFIG = { # Arguments to pass to the rllib optimizer "optimizer": {}, # Model parameters - "model": {"fcnet_hiddens": [128, 128]}, + "model": {"fcnet_hiddens": [128, 128], "max_seq_len": 20}, # Arguments to pass to the env creator "env_config": {}, diff --git a/python/ray/rllib/pg/pg_policy_graph.py b/python/ray/rllib/pg/pg_policy_graph.py index 6a095ee85..2fec360f2 100644 --- a/python/ray/rllib/pg/pg_policy_graph.py +++ b/python/ray/rllib/pg/pg_policy_graph.py @@ -21,12 +21,12 @@ class PGPolicyGraph(TFPolicyGraph): self.config = config # Setup policy - obs = tf.placeholder(tf.float32, shape=[None]+list(obs_space.shape)) + obs = tf.placeholder(tf.float32, shape=[None] + list(obs_space.shape)) dist_class, self.logit_dim = ModelCatalog.get_action_dist( action_space, self.config["model"]) - model = ModelCatalog.get_model( + self.model = ModelCatalog.get_model( obs, self.logit_dim, options=self.config["model"]) - action_dist = dist_class(model.outputs) # logit for each action + action_dist = dist_class(self.model.outputs) # logit for each action # Setup policy loss actions = ModelCatalog.get_action_placeholder(action_space) @@ -40,13 +40,25 @@ class PGPolicyGraph(TFPolicyGraph): ("actions", actions), ("advantages", advantages), ] - self.is_training = tf.placeholder_with_default(True, ()) + + # LSTM support + for i, ph in enumerate(self.model.state_in): + loss_in.append(("state_in_{}".format(i), ph)) + + is_training = tf.placeholder_with_default(True, ()) TFPolicyGraph.__init__( self, obs_space, action_space, sess, obs_input=obs, action_sampler=action_dist.sample(), loss=loss, - loss_inputs=loss_in, is_training=self.is_training) + loss_inputs=loss_in, is_training=is_training, + state_inputs=self.model.state_in, + state_outputs=self.model.state_out, + seq_lens=self.model.seq_lens, + max_seq_len=config["model"]["max_seq_len"]) sess.run(tf.global_variables_initializer()) def postprocess_trajectory(self, sample_batch, other_agent_batches=None): return compute_advantages( sample_batch, 0.0, self.config["gamma"], use_gae=False) + + def get_initial_state(self): + return self.model.state_init diff --git a/python/ray/rllib/test/test_catalog.py b/python/ray/rllib/test/test_catalog.py index e975b1b4c..e3865c3e7 100644 --- a/python/ray/rllib/test/test_catalog.py +++ b/python/ray/rllib/test/test_catalog.py @@ -23,7 +23,7 @@ class CustomPreprocessor2(Preprocessor): class CustomModel(Model): - def _init(self, *args): + def _build_layers(self, *args): return None, None @@ -78,7 +78,8 @@ class ModelCatalogTest(unittest.TestCase): def testCustomModel(self): ray.init() ModelCatalog.register_custom_model("foo", CustomModel) - p1 = ModelCatalog.get_model(1, 5, {"custom_model": "foo"}) + p1 = ModelCatalog.get_model( + tf.constant([1, 2, 3]), 5, {"custom_model": "foo"}) self.assertEqual(str(type(p1)), str(CustomModel)) diff --git a/python/ray/rllib/test/test_lstm.py b/python/ray/rllib/test/test_lstm.py new file mode 100644 index 000000000..0e92901fd --- /dev/null +++ b/python/ray/rllib/test/test_lstm.py @@ -0,0 +1,43 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import unittest + +from ray.rllib.models.lstm import chop_into_sequences + + +class LSTMUtilsTest(unittest.TestCase): + def testBasic(self): + t = [1, 2, 3, 1, 2, 3, 4, 5] + f = [ + [101, 102, 103, 201, 202, 203, 204, 205], + [[101], [102], [103], [201], [202], [203], [204], [205]] + ] + s = [[209, 208, 207, 109, 108, 107, 106, 105]] + f_pad, s_init, seq_lens = chop_into_sequences(t, f, s, 4) + self.assertEqual( + [f.tolist() for f in f_pad], + [ + [101, 102, 103, 0, + 201, 202, 203, 204, + 205, 0, 0, 0], + [[101], [102], [103], [0], + [201], [202], [203], [204], + [205], [0], [0], [0]], + ]) + self.assertEqual([s.tolist() for s in s_init], [[209, 109, 105]]) + self.assertEqual(seq_lens.tolist(), [3, 4, 1]) + + def testDynamicMaxLen(self): + t = [1, 1, 2] + f = [[1, 1, 1]] + s = [[1, 1, 1]] + f_pad, s_init, seq_lens = chop_into_sequences(t, f, s, 4) + self.assertEqual([f.tolist() for f in f_pad], [[1, 0, 1, 1]]) + self.assertEqual([s.tolist() for s in s_init], [[1, 1]]) + self.assertEqual(seq_lens.tolist(), [1, 2]) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/python/ray/rllib/test/test_multi_agent_env.py b/python/ray/rllib/test/test_multi_agent_env.py index 7424d0eba..c058e7714 100644 --- a/python/ray/rllib/test/test_multi_agent_env.py +++ b/python/ray/rllib/test/test_multi_agent_env.py @@ -275,12 +275,12 @@ class TestMultiAgentEnv(unittest.TestCase): # happen since the replay buffer doesn't encode extra fields like # "advantages" that PG uses. policies = { - "p1": (DQNPolicyGraph, obs_space, act_space, {}), + "p1": (DQNPolicyGraph, obs_space, act_space, dqn_config), "p2": (DQNPolicyGraph, obs_space, act_space, dqn_config), } else: policies = { - "p1": (PGPolicyGraph, obs_space, act_space, dqn_config), + "p1": (PGPolicyGraph, obs_space, act_space, {}), "p2": (DQNPolicyGraph, obs_space, act_space, dqn_config), } ev = CommonPolicyEvaluator( @@ -297,10 +297,10 @@ class TestMultiAgentEnv(unittest.TestCase): else: remote_evs = [] optimizer = optimizer_cls({}, ev, remote_evs) - ev.foreach_policy( - lambda p, _: p.set_epsilon(0.02) - if isinstance(p, DQNPolicyGraph) else None) for i in range(200): + ev.foreach_policy( + lambda p, _: p.set_epsilon(max(0.02, 1 - i * .02)) + if isinstance(p, DQNPolicyGraph) else None) optimizer.step() result = collect_metrics(ev, remote_evs) if i % 20 == 0: diff --git a/python/ray/rllib/tuned_examples/pong-a3c-pytorch.yaml b/python/ray/rllib/tuned_examples/pong-a3c-pytorch.yaml index 4d85eb688..46d84e6ae 100644 --- a/python/ray/rllib/tuned_examples/pong-a3c-pytorch.yaml +++ b/python/ray/rllib/tuned_examples/pong-a3c-pytorch.yaml @@ -12,7 +12,6 @@ pong-a3c-pytorch-cnn: lambda: 1.0 lr: 0.0001 observation_filter: NoFilter - reward_filter: NoFilter model: use_lstm: false channel_major: true diff --git a/python/ray/rllib/tuned_examples/pong-a3c.yaml b/python/ray/rllib/tuned_examples/pong-a3c.yaml index 400401b6d..4cb868bb5 100644 --- a/python/ray/rllib/tuned_examples/pong-a3c.yaml +++ b/python/ray/rllib/tuned_examples/pong-a3c.yaml @@ -1,8 +1,10 @@ +# This gets to ~19-20 reward in ~30 minutes / 4m steps on a m4.10xl instance +# TODO(rliaw): this has regressed in performance pong-a3c: env: PongDeterministic-v4 run: A3C config: - num_workers: 1 + num_workers: 16 batch_size: 20 use_pytorch: false vf_loss_coeff: 0.5 @@ -12,12 +14,18 @@ pong-a3c: lambda: 1.0 lr: 0.0001 observation_filter: NoFilter - reward_filter: NoFilter model: use_lstm: true - channel_major: false + conv_activation: elu dim: 42 grayscale: true zero_mean: false + # Reduced channel depth and kernel size from default + conv_filters: [ + [32, [3, 3], 2], + [32, [3, 3], 2], + [32, [3, 3], 2], + [32, [3, 3], 2], + ] optimizer: grads_per_step: 1000 diff --git a/python/ray/rllib/tuned_examples/pong-dqn.yaml b/python/ray/rllib/tuned_examples/pong-dqn.yaml index c22020c67..e36efe06a 100644 --- a/python/ray/rllib/tuned_examples/pong-dqn.yaml +++ b/python/ray/rllib/tuned_examples/pong-dqn.yaml @@ -19,8 +19,3 @@ pong-deterministic-dqn: grayscale: True zero_mean: False dim: 42 - conv_filters: [ - [16, [4, 4], 2], - [32, [4, 4], 2], - [512, [11, 11], 1], - ] diff --git a/python/ray/rllib/utils/tf_policy_graph.py b/python/ray/rllib/utils/tf_policy_graph.py index f881478e3..3f72bcb22 100644 --- a/python/ray/rllib/utils/tf_policy_graph.py +++ b/python/ray/rllib/utils/tf_policy_graph.py @@ -5,6 +5,7 @@ from __future__ import print_function import tensorflow as tf import ray +from ray.rllib.models.lstm import chop_into_sequences from ray.rllib.utils.policy_graph import PolicyGraph from ray.rllib.utils.tf_run_builder import TFRunBuilder @@ -16,7 +17,7 @@ class TFPolicyGraph(PolicyGraph): optimizations on the policy graph, e.g., parallelization across gpus or fusing multiple graphs together in the multi-agent setting. - All input and output tensors are of shape [BATCH_DIM, ...]. + Input tensors are typically shaped like [BATCH_SIZE, ...]. Attributes: observation_space (gym.Space): observation space of the policy. @@ -35,24 +36,32 @@ class TFPolicyGraph(PolicyGraph): def __init__( self, observation_space, action_space, sess, obs_input, - action_sampler, loss, loss_inputs, - is_training, state_inputs=None, state_outputs=None): + action_sampler, loss, loss_inputs, is_training, + state_inputs=None, state_outputs=None, seq_lens=None, + max_seq_len=20): """Initialize the policy graph. Arguments: observation_space (gym.Space): Observation space of the env. action_space (gym.Space): Action space of the env. sess (Session): TensorFlow session to use. - obs_input (Tensor): input placeholder for observations. - action_sampler (Tensor): Tensor for sampling an action. + obs_input (Tensor): input placeholder for observations, of shape + [BATCH_SIZE, obs...]. + action_sampler (Tensor): Tensor for sampling an action, of shape + [BATCH_SIZE, action...] loss (Tensor): scalar policy loss output tensor. loss_inputs (list): a (name, placeholder) tuple for each loss input argument. Each placeholder name must correspond to a - SampleBatch column key returned by postprocess_trajectory(). + SampleBatch column key returned by postprocess_trajectory(), + and has shape [BATCH_SIZE, data...]. is_training (Tensor): input placeholder for whether we are currently training the policy. state_inputs (list): list of RNN state output Tensors. state_outputs (list): list of initial state values. + seq_lens (Tensor): placeholder for RNN sequence lengths, of shape + [NUM_SEQUENCES]. Note that NUM_SEQUENCES << BATCH_SIZE. See + models/lstm.py for more information. + max_seq_len (int): max sequence length for LSTM training. """ self.observation_space = observation_space @@ -62,9 +71,12 @@ class TFPolicyGraph(PolicyGraph): self._sampler = action_sampler self._loss = loss self._loss_inputs = loss_inputs + self._loss_input_dict = dict(self._loss_inputs) self._is_training = is_training self._state_inputs = state_inputs or [] self._state_outputs = state_outputs or [] + self._seq_lens = seq_lens + self._max_seq_len = max_seq_len self._optimizer = self.optimizer() self._grads_and_vars = [ (g, v) for (g, v) in self.gradients(self._optimizer) @@ -77,6 +89,8 @@ class TFPolicyGraph(PolicyGraph): assert len(self._state_inputs) == len(self._state_outputs) == \ len(self.get_initial_state()), \ (self._state_inputs, self._state_outputs, self.get_initial_state()) + if self._state_inputs: + assert self._seq_lens is not None def build_compute_actions( self, builder, obs_batch, state_batches=None, is_training=False): @@ -99,15 +113,30 @@ class TFPolicyGraph(PolicyGraph): builder, obs_batch, state_batches, is_training) return builder.get(fetches) - def _get_loss_inputs_dict(self, postprocessed_batch): + def _get_loss_inputs_dict(self, batch): feed_dict = {} - for key, ph in self._loss_inputs: - # TODO(ekl) fix up handling of RNN inputs so that we can batch - # across multiple rollouts - if key.startswith("state_in_"): - feed_dict[ph] = postprocessed_batch[key][:1] # in state only - else: - feed_dict[ph] = postprocessed_batch[key] + + # Simple case + if not self._state_inputs: + for k, ph in self._loss_inputs: + feed_dict[ph] = batch[k] + return feed_dict + + # RNN case + feature_keys = [ + k for k, v in self._loss_inputs if not k.startswith("state_in_")] + state_keys = [ + k for k, v in self._loss_inputs if k.startswith("state_in_")] + feature_sequences, initial_states, seq_lens = chop_into_sequences( + batch["t"], + [batch[k] for k in feature_keys], + [batch[k] for k in state_keys], + self._max_seq_len) + for k, v in zip(feature_keys, feature_sequences): + feed_dict[self._loss_input_dict[k]] = v + for k, v in zip(state_keys, initial_states): + feed_dict[self._loss_input_dict[k]] = v + feed_dict[self._seq_lens] = seq_lens return feed_dict def build_compute_gradients(self, builder, postprocessed_batch): diff --git a/test/jenkins_tests/run_multi_node_tests.sh b/test/jenkins_tests/run_multi_node_tests.sh index 2324b3f60..15cb2beb7 100755 --- a/test/jenkins_tests/run_multi_node_tests.sh +++ b/test/jenkins_tests/run_multi_node_tests.sh @@ -128,6 +128,13 @@ docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ --stop '{"training_iteration": 2}' \ --config '{"batch_size": 500, "num_workers": 1}' +docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ + python /ray/python/ray/rllib/train.py \ + --env CartPole-v0 \ + --run PG \ + --stop '{"training_iteration": 2}' \ + --config '{"batch_size": 500, "num_workers": 1, "model": {"use_lstm": true, "max_seq_len": 100}}' + docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/rllib/train.py \ --env CartPole-v0 \ @@ -175,6 +182,12 @@ docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/rllib/test/test_serving_env.py +docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ + python /ray/python/ray/rllib/test/test_lstm.py + +docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ + python /ray/python/ray/rllib/test/test_multi_agent_env.py + docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/rllib/test/test_supported_spaces.py