[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
2026-07-09 09:19:50 +08:00 · 2018-06-27 22:51:04 -07:00
parent d3f81d5aad
commit b197c0c404
22 changed files with 388 additions and 116 deletions
@@ -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
@@ -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])
@@ -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
@@ -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
@@ -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
@@ -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):
@@ -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
@@ -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"]
@@ -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")
@@ -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):
@@ -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:
@@ -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.")
@@ -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": {},

@@ -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
@@ -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))


@@ -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)
@@ -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:
@@ -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
@@ -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
@@ -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],
-          ]
@@ -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):