[rllib] Pull out multi-gpu optimizer as a generic class (#1313)

doc/source/rllib.rst

@@ -13,7 +13,7 @@ Ray RLlib is a reinforcement learning library that aims to provide both performa
     - Scalable primitives for developing new algorithms
     - Shared models between algorithms
 
-You can find the code for RLlib `here on GitHub <https://github.com/ray-project/ray/tree/master/python/ray/rllib>`__.
+You can find the code for RLlib `here on GitHub <https://github.com/ray-project/ray/tree/master/python/ray/rllib>`__, and the NIPS symposium paper `here <https://drive.google.com/open?id=1lDMOFLMUQXn8qGtuahOBUwjmFb2iASxu>`__.
 
 RLlib currently provides the following algorithms:
 
@@ -30,11 +30,6 @@ RLlib currently provides the following algorithms:
 
 - `Deep Q Network (DQN) <https://arxiv.org/abs/1312.5602>`__.
 
-Proximal Policy Optimization scales to hundreds of cores and several GPUs,
-Evolution Strategies to clusters with thousands of cores and
-the Asynchronous Advantage Actor-Critic scales to dozens of cores
-on a single node.
-
 These algorithms can be run on any `OpenAI Gym MDP <https://github.com/openai/gym>`__,
 including custom ones written and registered by the user.
 
@@ -119,16 +114,6 @@ and renders its behavior in the environment specified by ``--env``.
 Checkpoints are be found within the experiment directory,
 specified by ``--local-dir`` and ``--experiment-name`` when running ``train.py``.
 
-
-The ``eval.py`` script has a number of options you can show by running
-
-::
-    python ray/python/ray/rllib/eval.py --help
-
-The most important argument is the checkpoint positional argument from which
-the script reconstructs the agent. The options ``--env`` and ``--run``
-must match the values chosen while running ``train.py``. 
-
 Tuned Examples
 --------------
 
@@ -248,10 +233,19 @@ The Developer API
 This part of the API will be useful if you need to change existing RL algorithms
 or implement new ones. Note that the API is not considered to be stable yet.
 
+Optimizers and Evaluators
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ray.rllib.optimizers.optimizer.Optimizer
+    :members:
+
+.. autoclass:: ray.rllib.optimizers.evaluator.Evaluator
+    :members:
+
 Models
 ~~~~~~
 
-Models are subclasses of the Model class:
+Algorithms share neural network models which inherit from the following class:
 
 .. autoclass:: ray.rllib.models.Model
 

python/ray/rllib/README.rst

@@ -1,7 +1,7 @@
 Ray RLlib: A Composable and Scalable Reinforcement Learning Library
 ===================================================================
 
-This README provides a brief technical overview of RLlib. See also the `user documentation <http://ray.readthedocs.io/en/latest/rllib.html>`__.
+This README provides a brief technical overview of RLlib. See also the `user documentation <http://ray.readthedocs.io/en/latest/rllib.html>`__ and `NIPS symposium paper <https://drive.google.com/open?id=1lDMOFLMUQXn8qGtuahOBUwjmFb2iASxu>`__.
 
 RLlib currently provides the following algorithms:
 
@@ -18,11 +18,8 @@ RLlib currently provides the following algorithms:
 
 - `Deep Q Network (DQN) <https://arxiv.org/abs/1312.5602>`__.
 
-Proximal Policy Optimization scales to hundreds of cores and several GPUs, Evolution Strategies to clusters with thousands of cores and the Asynchronous Advantage Actor-Critic scales to dozens of cores on a single node.
-
 These algorithms can be run on any OpenAI Gym MDP, including custom ones written and registered by the user.
 
-For more detailed usage information, see the `user documentation <http://ray.readthedocs.io/en/latest/rllib.html>`__.
 
 Training API
 ------------

python/ray/rllib/a3c/runner.py

@@ -4,7 +4,7 @@ from __future__ import print_function
 
 import ray
 from ray.rllib.envs import create_and_wrap
-from ray.rllib.evaluator import Evaluator
+from ray.rllib.optimizers import Evaluator
 from ray.rllib.a3c.common import get_policy_cls
 from ray.rllib.utils.filter import get_filter
 from ray.rllib.utils.sampler import AsyncSampler

python/ray/rllib/dqn/base_evaluator.py

@@ -9,7 +9,7 @@ import ray
 from ray.rllib.dqn import models
 from ray.rllib.dqn.common.wrappers import wrap_dqn
 from ray.rllib.dqn.common.schedules import LinearSchedule
-from ray.rllib.evaluator import TFMultiGPUSupport
+from ray.rllib.optimizers import SampleBatch, TFMultiGPUSupport
 
 
 class DQNEvaluator(TFMultiGPUSupport):
@@ -55,20 +55,23 @@ class DQNEvaluator(TFMultiGPUSupport):
         self.dqn_graph.update_target(self.sess)
 
     def sample(self):
-        output = []
+        obs, actions, rewards, new_obs, dones = [], [], [], [], []
         for _ in range(self.config["sample_batch_size"]):
-            result = self._step(self.global_timestep)
-            output.append(result)
-        return output
+            ob, act, rew, ob1, done = self._step(self.global_timestep)
+            obs.append(ob)
+            actions.append(act)
+            rewards.append(rew)
+            new_obs.append(ob1)
+            dones.append(done)
+        return SampleBatch({
+            "obs": obs, "actions": actions, "rewards": rewards,
+            "new_obs": new_obs, "dones": dones,
+            "weights": np.ones_like(rewards)})
 
     def compute_gradients(self, samples):
-        if self.config["prioritized_replay"]:
-            obses_t, actions, rewards, obses_tp1, dones, _ = samples
-        else:
-            obses_t, actions, rewards, obses_tp1, dones = samples
         _, grad = self.dqn_graph.compute_gradients(
-            self.sess, obses_t, actions, rewards, obses_tp1, dones,
-            np.ones_like(rewards))
+            self.sess, samples["obs"], samples["actions"], samples["rewards"],
+            samples["new_obs"], samples["dones"], samples["weights"])
         return grad
 
     def apply_gradients(self, grads):

python/ray/rllib/dqn/dqn.py

@@ -102,11 +102,7 @@ DEFAULT_CONFIG = dict(
     async_updates=False,
     # (Experimental) Whether to use multiple GPUs for SGD optimization.
     # Note that this only helps performance if the SGD batch size is large.
-    multi_gpu_optimize=False,
-    # Number of SGD iterations over the data. Only applies in multi-gpu mode.
-    num_sgd_iter=1,
-    # Devices to use for parallel SGD. Only applies in multi-gpu mode.
-    devices=["/gpu:0"])
+    multi_gpu=False)
 
 
 class DQNAgent(Agent):
@@ -136,7 +132,7 @@ class DQNAgent(Agent):
             # will internally create more workers for parallelism. This means
             # there is only one replay buffer regardless of num_workers.
             self.remote_evaluators = []
-            if self.config["multi_gpu_optimize"]:
+            if self.config["multi_gpu"]:
                 optimizer_cls = LocalMultiGPUOptimizer
             else:
                 optimizer_cls = LocalSyncOptimizer

python/ray/rllib/dqn/models.py

@@ -6,7 +6,7 @@ import tensorflow as tf
 import tensorflow.contrib.layers as layers
 
 from ray.rllib.models import ModelCatalog
-from ray.rllib.parallel import LocalSyncParallelOptimizer, TOWER_SCOPE_NAME
+from ray.rllib.parallel import TOWER_SCOPE_NAME
 
 
 def _build_q_network(inputs, num_actions, config):
@@ -159,10 +159,7 @@ class DQNGraph(object):
             tf.float32, shape=(None,) + env.observation_space.shape)
 
         # Action Q network
-        if config["multi_gpu_optimize"]:
-            q_scope_name = TOWER_SCOPE_NAME + "/q_func"
-        else:
-            q_scope_name = "q_func"
+        q_scope_name = TOWER_SCOPE_NAME + "/q_func"
         with tf.variable_scope(q_scope_name) as scope:
             q_values = _build_q_network(
                 self.cur_observations, num_actions, config)
@@ -194,26 +191,21 @@ class DQNGraph(object):
                 obs_t, act_t, rew_t, obs_tp1, done_mask, importance_weights)
 
         self.loss_inputs = [
-            self.obs_t, self.act_t, self.rew_t, self.obs_tp1, self.done_mask,
-            self.importance_weights]
-        self.build_loss = build_loss
+            ("obs", self.obs_t),
+            ("actions", self.act_t),
+            ("rewards", self.rew_t),
+            ("new_obs", self.obs_tp1),
+            ("dones", self.done_mask),
+            ("weights", self.importance_weights),
+        ]
 
-        if config["multi_gpu_optimize"]:
-            self.multi_gpu_optimizer = LocalSyncParallelOptimizer(
-                optimizer,
-                config["devices"],
-                [self.obs_t, self.act_t, self.rew_t, self.obs_tp1,
-                 self.done_mask, self.importance_weights],
-                int(config["sgd_batch_size"] / len(config["devices"])),
-                build_loss,
-                logdir,
-                grad_norm_clipping=config["grad_norm_clipping"])
-            loss_obj = self.multi_gpu_optimizer.get_common_loss()
-        else:
+        with tf.variable_scope(TOWER_SCOPE_NAME):
             loss_obj = build_loss(
                 self.obs_t, self.act_t, self.rew_t, self.obs_tp1,
                 self.done_mask, self.importance_weights)
 
+        self.build_loss = build_loss
+
         weighted_error = loss_obj.loss
         target_q_func_vars = loss_obj.target_q_func_vars
         self.q_t = loss_obj.q_t

python/ray/rllib/dqn/replay_evaluator.py

@@ -8,6 +8,7 @@ import ray
 from ray.rllib.dqn.base_evaluator import DQNEvaluator
 from ray.rllib.dqn.common.schedules import LinearSchedule
 from ray.rllib.dqn.replay_buffer import ReplayBuffer, PrioritizedReplayBuffer
+from ray.rllib.optimizers import SampleBatch
 
 
 class DQNReplayEvaluator(DQNEvaluator):
@@ -63,38 +64,44 @@ class DQNReplayEvaluator(DQNEvaluator):
             samples = [DQNEvaluator.sample(self)]
 
         for s in samples:
-            for obs, action, rew, new_obs, done in s:
-                self.replay_buffer.add(obs, action, rew, new_obs, done)
+            for row in s.rows():
+                self.replay_buffer.add(
+                    row["obs"], row["actions"], row["rewards"], row["new_obs"],
+                    row["dones"])
 
         if no_replay:
             return samples
 
         # Then return a batch sampled from the buffer
         if self.config["prioritized_replay"]:
-            experience = self.replay_buffer.sample(
-                self.config["train_batch_size"],
-                beta=self.beta_schedule.value(self.global_timestep))
             (obses_t, actions, rewards, obses_tp1,
-                dones, _, batch_idxes) = experience
+                dones, weights, batch_indexes) = self.replay_buffer.sample(
+                    self.config["train_batch_size"],
+                    beta=self.beta_schedule.value(self.global_timestep))
             self._update_priorities_if_needed()
-            self.samples_to_prioritize = (
-                obses_t, actions, rewards, obses_tp1, dones, batch_idxes)
+            batch = SampleBatch({
+                "obs": obses_t, "actions": actions, "rewards": rewards,
+                "new_obs": obses_tp1, "dones": dones, "weights": weights,
+                "batch_indexes": batch_indexes})
+            self.samples_to_prioritize = batch
         else:
             obses_t, actions, rewards, obses_tp1, dones = \
                 self.replay_buffer.sample(self.config["train_batch_size"])
-            batch_idxes = None
-
-        return self.samples_to_prioritize
+            batch = SampleBatch({
+                "obs": obses_t, "actions": actions, "rewards": rewards,
+                "new_obs": obses_tp1, "dones": dones,
+                "weights": np.ones_like(rewards)})
+        return batch
 
     def compute_gradients(self, samples):
-        obses_t, actions, rewards, obses_tp1, dones, batch_indxes = samples
         td_errors, grad = self.dqn_graph.compute_gradients(
-            self.sess, obses_t, actions, rewards, obses_tp1, dones,
-            np.ones_like(rewards))
+            self.sess, samples["obs"], samples["actions"], samples["rewards"],
+            samples["new_obs"], samples["dones"], samples["weights"])
         if self.config["prioritized_replay"]:
             new_priorities = (
                 np.abs(td_errors) + self.config["prioritized_replay_eps"])
-            self.replay_buffer.update_priorities(batch_indxes, new_priorities)
+            self.replay_buffer.update_priorities(
+                samples["batch_indexes"], new_priorities)
             self.samples_to_prioritize = None
         return grad
 
@@ -109,14 +116,15 @@ class DQNReplayEvaluator(DQNEvaluator):
         if not self.samples_to_prioritize:
             return
 
-        obses_t, actions, rewards, obses_tp1, dones, batch_idxes = \
-            self.samples_to_prioritize
+        batch = self.samples_to_prioritize
         td_errors = self.dqn_graph.compute_td_error(
-            self.sess, obses_t, actions, rewards, obses_tp1, dones,
-            np.ones_like(rewards))
+            self.sess, batch["obs"], batch["actions"], batch["rewards"],
+            batch["new_obs"], batch["dones"], batch["weights"])
+
         new_priorities = (
             np.abs(td_errors) + self.config["prioritized_replay_eps"])
-        self.replay_buffer.update_priorities(batch_idxes, new_priorities)
+        self.replay_buffer.update_priorities(
+            batch["batch_indexes"], new_priorities)
         self.samples_to_prioritize = None
 
     def stats(self):

python/ray/rllib/evaluator.py

@@ -1,50 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-
-class Evaluator(object):
-    """RLlib optimizers require RL algorithms to implement this interface.
-
-    Any algorithm that implements Evaluator can plug in any RLLib optimizer,
-    e.g. async SGD, local multi-GPU SGD, etc.
-    """
-
-    def sample(self):
-        """Returns experience samples from this Evaluator."""
-
-        raise NotImplementedError
-
-    def compute_gradients(self, samples):
-        """Returns a gradient computed w.r.t the specified samples."""
-
-        raise NotImplementedError
-
-    def apply_gradients(self, grads):
-        """Applies the given gradients to this Evaluator's weights."""
-
-        raise NotImplementedError
-
-    def get_weights(self):
-        """Returns the model weights of this Evaluator."""
-
-        raise NotImplementedError
-
-    def set_weights(self, weights):
-        """Sets the model weights of this Evaluator."""
-
-        raise NotImplementedError
-
-
-class TFMultiGPUSupport(Evaluator):
-    """The multi-GPU TF optimizer requires this additional interface."""
-
-    def tf_loss_inputs(self):
-        """Returns a list of the input placeholders required for the loss."""
-
-        raise NotImplementedError
-
-    def build_tf_loss(self, input_placeholders):
-        """Returns a new loss tensor graph for the specified inputs."""
-
-        raise NotImplementedError

python/ray/rllib/optimizers/__init__.py

@@ -1,6 +1,10 @@
 from ray.rllib.optimizers.async import AsyncOptimizer
 from ray.rllib.optimizers.local_sync import LocalSyncOptimizer
 from ray.rllib.optimizers.multi_gpu import LocalMultiGPUOptimizer
+from ray.rllib.optimizers.sample_batch import SampleBatch
+from ray.rllib.optimizers.evaluator import Evaluator, TFMultiGPUSupport
 
 
-__all__ = ["AsyncOptimizer", "LocalSyncOptimizer", "LocalMultiGPUOptimizer"]
+__all__ = [
+    "AsyncOptimizer", "LocalSyncOptimizer", "LocalMultiGPUOptimizer",
+    "SampleBatch", "Evaluator", "TFMultiGPUSupport"]

python/ray/rllib/optimizers/evaluator.py

@@ -0,0 +1,104 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+class Evaluator(object):
+    """Algorithms implement this interface to leverage RLlib optimizers.
+
+    Any algorithm that implements Evaluator can plug in any RLLib optimizer,
+    e.g. async SGD, local multi-GPU SGD, etc.
+    """
+
+    def sample(self):
+        """Returns experience samples from this Evaluator.
+
+        Returns:
+            SampleBatch: A columnar batch of experiences.
+
+        Examples:
+            >>> print(ev.sample())
+            SampleBatch({"a": [1, 2, 3], "b": [4, 5, 6]})
+        """
+
+        raise NotImplementedError
+
+    def compute_gradients(self, samples):
+        """Returns a gradient computed w.r.t the specified samples.
+
+        Returns:
+            object: A gradient that can be applied on a compatible evaluator.
+        """
+
+        raise NotImplementedError
+
+    def apply_gradients(self, grads):
+        """Applies the given gradients to this Evaluator's weights.
+
+        Examples:
+            >>> samples = ev1.sample()
+            >>> grads = ev2.compute_gradients(samples)
+            >>> ev1.apply_gradients(grads)
+        """
+
+        raise NotImplementedError
+
+    def get_weights(self):
+        """Returns the model weights of this Evaluator.
+
+        Returns:
+            object: weights that can be set on a compatible evaluator.
+        """
+
+        raise NotImplementedError
+
+    def set_weights(self, weights):
+        """Sets the model weights of this Evaluator.
+
+        Examples:
+            >>> weights = ev1.get_weights()
+            >>> ev2.set_weights(weights)
+        """
+
+        raise NotImplementedError
+
+
+class TFMultiGPUSupport(Evaluator):
+    """The multi-GPU TF optimizer requires additional TF-specific supportt.
+
+    Attributes:
+        sess (Session) the tensorflow session associated with this evaluator
+    """
+
+    def tf_loss_inputs(self):
+        """Returns a list of the input placeholders required for the loss.
+
+        For example, the following calls should work:
+
+        Returns:
+            list: a (name, placeholder) tuple for each loss input argument.
+                Each placeholder name must correspond to one of the SampleBatch
+                column keys returned by sample().
+
+        Examples:
+            >>> print(ev.tf_loss_inputs())
+            [("action", action_placeholder), ("reward", reward_placeholder)]
+
+            >>> print(ev.sample().data.keys())
+            ["action", "reward"]
+        """
+
+        raise NotImplementedError
+
+    def build_tf_loss(self, input_placeholders):
+        """Returns a new loss tensor graph for the specified inputs.
+
+        The graph must share vars with this Evaluator's policy model, so that
+        the multi-gpu optimizer can update the weights.
+
+        Examples:
+            >>> loss_inputs = ev.tf_loss_inputs()
+            >>> ev.build_tf_loss([ph for _, ph in loss_inputs])
+        """
+
+        raise NotImplementedError

python/ray/rllib/optimizers/local_sync.py

@@ -4,6 +4,7 @@ from __future__ import print_function
 
 import ray
 from ray.rllib.optimizers.optimizer import Optimizer
+from ray.rllib.optimizers.sample_batch import SampleBatch
 from ray.rllib.utils.timer import TimerStat
 
 
@@ -29,7 +30,7 @@ class LocalSyncOptimizer(Optimizer):
 
         with self.sample_timer:
             if self.remote_evaluators:
-                samples = _concat(
+                samples = SampleBatch.concat_samples(
                     ray.get(
                         [e.sample.remote() for e in self.remote_evaluators]))
             else:
@@ -45,11 +46,3 @@ class LocalSyncOptimizer(Optimizer):
             "grad_time_ms": round(1000 * self.grad_timer.mean, 3),
             "update_time_ms": round(1000 * self.update_weights_timer.mean, 3),
         }
-
-
-# TODO(ekl) this should be implemented by some sample batch class
-def _concat(samples):
-    result = []
-    for s in samples:
-        result.extend(s)
-    return result

python/ray/rllib/optimizers/multi_gpu.py

@@ -2,8 +2,104 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+import os
+import tensorflow as tf
+
+import ray
+from ray.rllib.optimizers.evaluator import TFMultiGPUSupport
 from ray.rllib.optimizers.optimizer import Optimizer
+from ray.rllib.optimizers.sample_batch import SampleBatch
+from ray.rllib.parallel import LocalSyncParallelOptimizer
+from ray.rllib.utils.timer import TimerStat
 
 
 class LocalMultiGPUOptimizer(Optimizer):
-    pass  # TODO(ekl)
+    """A synchronous optimizer that uses multiple local GPUs.
+
+    Samples are pulled synchronously from multiple remote evaluators,
+    concatenated, and then split across the memory of multiple local GPUs.
+    A number of SGD passes are then taken over the in-memory data. For more
+    details, see `ray.rllib.parallel.LocalSyncParallelOptimizer`.
+
+    This optimizer is Tensorflow-specific and require evaluators to implement
+    the TFMultiGPUSupport API.
+    """
+
+    def _init(self):
+        assert isinstance(self.local_evaluator, TFMultiGPUSupport)
+        self.batch_size = self.config.get("sgd_batch_size", 128)
+        gpu_ids = ray.get_gpu_ids()
+        if not gpu_ids:
+            self.devices = ["/cpu:0"]
+        else:
+            self.devices = ["/gpu:{}".format(i) for i in range(len(gpu_ids))]
+        assert self.batch_size > len(self.devices), "batch size too small"
+        self.per_device_batch_size = self.batch_size // len(self.devices)
+        self.sample_timer = TimerStat()
+        self.load_timer = TimerStat()
+        self.grad_timer = TimerStat()
+        self.update_weights_timer = TimerStat()
+
+        print("LocalMultiGPUOptimizer devices", self.devices)
+        print("LocalMultiGPUOptimizer batch size", self.batch_size)
+
+        # List of (feature name, feature placeholder) tuples
+        self.loss_inputs = self.local_evaluator.tf_loss_inputs()
+
+        # per-GPU graph copies created below must share vars with the policy
+        tf.get_variable_scope().reuse_variables()
+
+        self.par_opt = LocalSyncParallelOptimizer(
+            tf.train.AdamOptimizer(self.config.get("sgd_stepsize", 5e-5)),
+            self.devices,
+            [ph for _, ph in self.loss_inputs],
+            self.per_device_batch_size,
+            lambda *ph: self.local_evaluator.build_tf_loss(ph),
+            self.config.get("logdir", os.getcwd()))
+
+        self.sess = self.local_evaluator.sess
+        self.sess.run(tf.global_variables_initializer())
+
+    def step(self):
+        with self.update_weights_timer:
+            if self.remote_evaluators:
+                weights = ray.put(self.local_evaluator.get_weights())
+                for e in self.remote_evaluators:
+                    e.set_weights.remote(weights)
+
+        with self.sample_timer:
+            if self.remote_evaluators:
+                samples = SampleBatch.concat_samples(
+                    ray.get(
+                        [e.sample.remote() for e in self.remote_evaluators]))
+            else:
+                samples = self.local_evaluator.sample()
+            assert isinstance(samples, SampleBatch)
+
+        with self.load_timer:
+            tuples_per_device = self.par_opt.load_data(
+                self.local_evaluator.sess,
+                samples.columns([key for key, _ in self.loss_inputs]))
+
+        with self.grad_timer:
+            for i in range(self.config.get("num_sgd_iter", 10)):
+                batch_index = 0
+                num_batches = (
+                    int(tuples_per_device) // int(self.per_device_batch_size))
+                permutation = np.random.permutation(num_batches)
+                while batch_index < num_batches:
+                    # TODO(ekl) support ppo's debugging features, e.g.
+                    # printing the current loss and tracing
+                    self.par_opt.optimize(
+                        self.sess,
+                        permutation[batch_index] * self.per_device_batch_size)
+                    batch_index += 1
+
+    def stats(self):
+        return {
+            "sample_time_ms": round(1000 * self.sample_timer.mean, 3),
+            "load_time_ms": round(1000 * self.load_timer.mean, 3),
+            "grad_time_ms": round(1000 * self.grad_timer.mean, 3),
+            "update_time_ms": round(1000 * self.update_weights_timer.mean, 3),
+        }

python/ray/rllib/optimizers/sample_batch.py

@@ -0,0 +1,87 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from functools import reduce
+
+import numpy as np
+
+
+class SampleBatch(object):
+    """Wrapper around a dictionary with string keys and array-like values.
+
+    For example, {"obs": [1, 2, 3], "reward": [0, -1, 1]} is a batch of three
+    samples, each with an "obs" and "reward" attribute.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Constructs a sample batch (same params as dict constructor)."""
+
+        self.data = dict(*args, **kwargs)
+        lengths = []
+        for k, v in self.data.copy().items():
+            assert type(k) == str, self
+            lengths.append(len(v))
+        assert len(set(lengths)) == 1, "data columns must be same length"
+
+    @staticmethod
+    def concat_samples(samples):
+        return reduce(lambda a, b: a.concat(b), samples)
+
+    def concat(self, other):
+        """Returns a new SampleBatch with each data column concatenated.
+
+        Examples:
+            >>> b1 = SampleBatch({"a": [1, 2]})
+            >>> b2 = SampleBatch({"a": [3, 4, 5]})
+            >>> print(b1.concat(b2))
+            {"a": [1, 2, 3, 4, 5]}
+        """
+
+        assert self.data.keys() == other.data.keys(), "must have same columns"
+        out = {}
+        for k in self.data.keys():
+            out[k] = np.concatenate([self.data[k], other.data[k]])
+        return SampleBatch(out)
+
+    def rows(self):
+        """Returns an iterator over data rows, i.e. dicts with column values.
+
+        Examples:
+            >>> batch = SampleBatch({"a": [1, 2, 3], "b": [4, 5, 6]})
+            >>> for row in batch.rows():
+                   print(row)
+            {"a": 1, "b": 4}
+            {"a": 2, "b": 5}
+            {"a": 3, "b": 6}
+        """
+
+        num_rows = len(list(self.data.values())[0])
+        for i in range(num_rows):
+            row = {}
+            for k in self.data.keys():
+                row[k] = self[k][i]
+            yield row
+
+    def columns(self, keys):
+        """Returns a list of just the specified columns.
+
+        Examples:
+            >>> batch = SampleBatch({"a": [1], "b": [2], "c": [3]})
+            >>> print(batch.columns(["a", "b"]))
+            [[1], [2]]
+        """
+
+        out = []
+        for k in keys:
+            out.append(self.data[k])
+        return out
+
+    def __getitem__(self, key):
+        return self.data[key]
+
+    def __str__(self):
+        return str(self.data)
+
+    def __repr__(self):
+        return str(self.data)

python/ray/rllib/parallel.py

@@ -127,6 +127,11 @@ class LocalSyncParallelOptimizer(object):
             trace_file.write(trace.generate_chrome_trace_format())
 
         tuples_per_device = truncated_len / len(self.devices)
+        assert tuples_per_device > 0, \
+            "Too few tuples per batch, trying increasing the training " \
+            "batch size or decreasing the sgd batch size. Tried to split up " \
+            "{} rows {}-ways in batches of {} (total across devices).".format(
+                len(arr), len(self.devices), self.batch_size)
         assert tuples_per_device % self.per_device_batch_size == 0
         return tuples_per_device
 

python/ray/rllib/tuned_examples/pong-dqn.yaml

@@ -1,3 +1,4 @@
+# You can expect ~20 reward within 1.1m timesteps / 2.1 hours on a K80 GPU
 pong-deterministic-dqn:
     env: PongDeterministic-v4
     run: DQN
@@ -26,31 +27,3 @@ pong-deterministic-dqn:
               [32, [4, 4], 2],
               [512, [11, 11], 1],
           ]
-pong-noframeskip-dqn:
-    env: PongNoFrameskip-v4
-    run: DQN
-    resources:
-       cpu: 1
-       gpu: 1
-    stop:
-        episode_reward_mean: 20
-        time_total_s: 7200
-    config:
-        gamma: 0.99
-        lr: .0001
-        learning_starts: 10000
-        buffer_size: 50000
-        sample_batch_size: 4
-        train_batch_size: 32
-        schedule_max_timesteps: 2000000
-        exploration_final_eps: .01
-        exploration_fraction: .1
-        model:
-          grayscale: True
-          zero_mean: False
-          dim: 42
-          conv_filters: [
-              [16, [4, 4], 2],
-              [32, [4, 4], 2],
-              [512, [11, 11], 1],
-          ]

test/jenkins_tests/run_multi_node_tests.sh

@@ -112,6 +112,13 @@ docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \
     --stop '{"training_iteration": 2}' \
     --config '{"async_updates": true, "num_workers": 2}'
 
+docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \
+    python /ray/python/ray/rllib/train.py \
+    --env CartPole-v0 \
+    --run DQN \
+    --stop '{"training_iteration": 2}' \
+    --config '{"multi_gpu": true, "optimizer": {"sgd_batch_size": 4}}'
+
 docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \
     python /ray/python/ray/rllib/train.py \
     --env FrozenLake-v0 \