[rllib] Refactor DQN to use an Evaluator abstraction (#1276)

This introduces rllib.Evaluator and rllib.Optimizer classes. Optimizers encapsulate a particular distributed optimization strategy for RL. Evaluators encapsulate the model graph, and once implemented, any Optimizer may be "plugged in" to any algorithm that implements the Evaluator interface.

python/ray/rllib/optimizers/__init__.py

@@ -0,0 +1,6 @@
+from ray.rllib.optimizers.async import AsyncOptimizer
+from ray.rllib.optimizers.local_sync import LocalSyncOptimizer
+from ray.rllib.optimizers.multi_gpu import LocalMultiGPUOptimizer
+
+
+__all__ = ["AsyncOptimizer", "LocalSyncOptimizer", "LocalMultiGPUOptimizer"]

python/ray/rllib/optimizers/async.py

@@ -0,0 +1,58 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import ray
+from ray.rllib.optimizers.optimizer import Optimizer
+from ray.rllib.utils.timer import TimerStat
+
+
+class AsyncOptimizer(Optimizer):
+    """An asynchronous RL optimizer, e.g. for implementing A3C.
+
+    This optimizer asynchronously pulls and applies gradients from remote
+    evaluators, sending updated weights back as needed. This pipelines the
+    gradient computations on the remote workers.
+    """
+    def _init(self):
+        self.apply_timer = TimerStat()
+        self.wait_timer = TimerStat()
+        self.dispatch_timer = TimerStat()
+        self.grads_per_step = self.config.get("grads_per_step", 100)
+
+    def step(self):
+        weights = ray.put(self.local_evaluator.get_weights())
+        gradient_queue = []
+        num_gradients = 0
+
+        # Kick off the first wave of async tasks
+        for e in self.remote_evaluators:
+            e.set_weights.remote(weights)
+            fut = e.compute_gradients.remote(e.sample.remote())
+            gradient_queue.append((fut, e))
+            num_gradients += 1
+
+        # Note: can't use wait: https://github.com/ray-project/ray/issues/1128
+        while gradient_queue:
+            with self.wait_timer:
+                fut, e = gradient_queue[0]
+                gradient_queue = gradient_queue[1:]
+                gradient = ray.get(fut)
+
+            if gradient is not None:
+                with self.apply_timer:
+                    self.local_evaluator.apply_gradients(gradient)
+
+            if num_gradients < self.grads_per_step:
+                with self.dispatch_timer:
+                    e.set_weights.remote(self.local_evaluator.get_weights())
+                    fut = e.compute_gradients.remote(e.sample.remote())
+                    gradient_queue.append((fut, e))
+                    num_gradients += 1
+
+    def stats(self):
+        return {
+            "wait_time_ms": round(1000 * self.wait_timer.mean, 3),
+            "apply_time_ms": round(1000 * self.apply_timer.mean, 3),
+            "dispatch_time_ms": round(1000 * self.dispatch_timer.mean, 3),
+        }

python/ray/rllib/optimizers/local_sync.py

@@ -0,0 +1,55 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import ray
+from ray.rllib.optimizers.optimizer import Optimizer
+from ray.rllib.utils.timer import TimerStat
+
+
+class LocalSyncOptimizer(Optimizer):
+    """A simple synchronous RL optimizer.
+
+    In each step, this optimizer pulls samples from a number of remote
+    evaluators, concatenates them, and then updates a local model. The updated
+    model weights are then broadcast to all remote evaluators.
+    """
+
+    def _init(self):
+        self.update_weights_timer = TimerStat()
+        self.sample_timer = TimerStat()
+        self.grad_timer = TimerStat()
+
+    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 = _concat(
+                    ray.get(
+                        [e.sample.remote() for e in self.remote_evaluators]))
+            else:
+                samples = self.local_evaluator.sample()
+
+        with self.grad_timer:
+            grad = self.local_evaluator.compute_gradients(samples)
+            self.local_evaluator.apply_gradients(grad)
+
+    def stats(self):
+        return {
+            "sample_time_ms": round(1000 * self.sample_timer.mean, 3),
+            "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

@@ -0,0 +1,9 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.optimizers.optimizer import Optimizer
+
+
+class LocalMultiGPUOptimizer(Optimizer):
+    pass  # TODO(ekl)

python/ray/rllib/optimizers/optimizer.py

@@ -0,0 +1,45 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+class Optimizer(object):
+    """RLlib optimizers encapsulate distributed RL optimization strategies.
+
+    For example, AsyncOptimizer is used for A3C, and LocalMultiGPUOptimizer is
+    used for PPO. These optimizers are all pluggable however, it is possible
+    to mix as match as needed.
+
+    In order for an algorithm to use an RLlib optimizer, it must implement
+    the Evaluator interface and pass a number of Evaluators to its Optimizer
+    of choice. The Optimizer uses these Evaluators to sample from the
+    environment and compute model gradient updates.
+    """
+
+    def __init__(self, config, local_evaluator, remote_evaluators):
+        """Create an optimizer instance.
+
+        Args:
+            config (dict): Optimizer-specific configuration data.
+            local_evaluator (Evaluator): Local evaluator instance, required.
+            remote_evaluators (list): A list of handles to remote evaluators.
+                if empty, the optimizer should fall back to to using only the
+                local evaluator.
+        """
+        self.config = config
+        self.local_evaluator = local_evaluator
+        self.remote_evaluators = remote_evaluators
+        self._init()
+
+    def _init(self):
+        pass
+
+    def step(self):
+        """Takes a logical optimization step."""
+
+        raise NotImplementedError
+
+    def stats(self):
+        """Returns a dictionary of internal performance statistics."""
+
+        return {}

python/ray/rllib/optimizers/util.py

@@ -0,0 +1,33 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+
+# TODO(ekl) move this to a common location
+from ray.rllib.ppo.filter import RunningStat
+
+
+class TimerStat(RunningStat):
+    """A running stat for conveniently logging the duration of a code block.
+
+    Example:
+        wait_timer = TimeStat()
+        with wait_timer:
+            ray.wait(...)
+
+    Note that this class is *not* thread-safe.
+    """
+
+    def __init__(self):
+        RunningStat.__init__(self, ())
+        self._start_time = None
+
+    def __enter__(self):
+        assert self._start_time is None, "concurrent updates not supported"
+        self._start_time = time.monotonic()
+
+    def __exit__(self, type, value, tb):
+        assert self._start_time is not None
+        self.push(time.monotonic() - self._start_time)
+        self._start_time = None