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

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)