[rllib] Ape-X implementation and DQN refactor to handle replay in policy optimizer (#1604)

* minimal apex checkin

* cleanup dqn options

* actor utils

* Sun Feb 25 17:39:54 PST 2018

* update

* compression refactor

* fix

* add test

* fix models

* Sun Feb 25 21:46:27 PST 2018

* Wed Feb 28 10:26:34 PST 2018

* Wed Feb 28 10:28:09 PST 2018

* Wed Feb 28 10:42:59 PST 2018

* refactor

* Wed Feb 28 11:17:19 PST 2018

* Wed Feb 28 11:42:08 PST 2018

* Wed Feb 28 11:42:13 PST 2018

* Wed Feb 28 11:59:02 PST 2018

* Wed Feb 28 11:59:58 PST 2018

* Wed Feb 28 12:00:08 PST 2018

* Wed Feb 28 12:02:19 PST 2018

* Wed Feb 28 13:44:31 PST 2018

* Wed Feb 28 17:01:20 PST 2018

* Sat Mar  3 14:55:59 PST 2018

* make optimizer construction explicit

* Sat Mar  3 18:23:08 PST 2018

* Sat Mar  3 18:24:28 PST 2018

* Sat Mar  3 18:49:28 PST 2018

* Sat Mar  3 18:50:42 PST 2018

* Sat Mar  3 18:56:10 PST 2018

python/ray/rllib/optimizers/__init__.py

@@ -1,10 +1,13 @@
+from ray.rllib.optimizers.apex_optimizer import ApexOptimizer
 from ray.rllib.optimizers.async import AsyncOptimizer
 from ray.rllib.optimizers.local_sync import LocalSyncOptimizer
+from ray.rllib.optimizers.local_sync_replay import LocalSyncReplayOptimizer
 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",
-    "SampleBatch", "Evaluator", "TFMultiGPUSupport"]
+    "ApexOptimizer", "AsyncOptimizer", "LocalSyncOptimizer",
+    "LocalSyncReplayOptimizer", "LocalMultiGPUOptimizer", "SampleBatch",
+    "Evaluator", "TFMultiGPUSupport"]

python/ray/rllib/optimizers/apex_optimizer.py

@@ -0,0 +1,265 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import queue
+import random
+import time
+import threading
+
+import numpy as np
+
+import ray
+from ray.rllib.optimizers.optimizer import Optimizer
+from ray.rllib.optimizers.replay_buffer import PrioritizedReplayBuffer
+from ray.rllib.optimizers.sample_batch import SampleBatch
+from ray.rllib.utils.actors import TaskPool, create_colocated
+from ray.rllib.utils.timer import TimerStat
+from ray.rllib.utils.window_stat import WindowStat
+
+SAMPLE_QUEUE_DEPTH = 2
+REPLAY_QUEUE_DEPTH = 4
+LEARNER_QUEUE_MAX_SIZE = 16
+
+
+@ray.remote
+class ReplayActor(object):
+    def __init__(
+            self, num_shards, learning_starts, buffer_size, train_batch_size,
+            prioritized_replay_alpha, prioritized_replay_beta,
+            prioritized_replay_eps):
+        self.replay_starts = learning_starts // num_shards
+        self.buffer_size = buffer_size // num_shards
+        self.train_batch_size = train_batch_size
+        self.prioritized_replay_beta = prioritized_replay_beta
+        self.prioritized_replay_eps = prioritized_replay_eps
+
+        self.replay_buffer = PrioritizedReplayBuffer(
+            buffer_size, alpha=prioritized_replay_alpha)
+
+        # Metrics
+        self.add_batch_timer = TimerStat()
+        self.replay_timer = TimerStat()
+        self.update_priorities_timer = TimerStat()
+
+    def get_host(self):
+        return os.uname()[1]
+
+    def add_batch(self, batch):
+        with self.add_batch_timer:
+            for row in batch.rows():
+                self.replay_buffer.add(
+                    row["obs"], row["actions"], row["rewards"], row["new_obs"],
+                    row["dones"], row["weights"])
+
+    def replay(self):
+        with self.replay_timer:
+            if len(self.replay_buffer) < self.replay_starts:
+                return None
+
+            (obses_t, actions, rewards, obses_tp1,
+                dones, weights, batch_indexes) = self.replay_buffer.sample(
+                    self.train_batch_size,
+                    beta=self.prioritized_replay_beta)
+
+            batch = SampleBatch({
+                "obs": obses_t, "actions": actions, "rewards": rewards,
+                "new_obs": obses_tp1, "dones": dones, "weights": weights,
+                "batch_indexes": batch_indexes})
+            return batch
+
+    def update_priorities(self, batch, td_errors):
+        with self.update_priorities_timer:
+            new_priorities = (
+                np.abs(td_errors) + self.prioritized_replay_eps)
+            self.replay_buffer.update_priorities(
+                batch["batch_indexes"], new_priorities)
+
+    def stats(self):
+        stat = {
+            "add_batch_time_ms": round(
+                1000 * self.add_batch_timer.mean, 3),
+            "replay_time_ms": round(
+                1000 * self.replay_timer.mean, 3),
+            "update_priorities_time_ms": round(
+                1000 * self.update_priorities_timer.mean, 3),
+        }
+        stat.update(self.replay_buffer.stats())
+        return stat
+
+
+class GenericLearner(threading.Thread):
+    def __init__(self, local_evaluator):
+        threading.Thread.__init__(self)
+        self.learner_queue_size = WindowStat("size", 50)
+        self.local_evaluator = local_evaluator
+        self.inqueue = queue.Queue(maxsize=LEARNER_QUEUE_MAX_SIZE)
+        self.outqueue = queue.Queue()
+        self.queue_timer = TimerStat()
+        self.grad_timer = TimerStat()
+        self.daemon = True
+
+    def run(self):
+        while True:
+            self.step()
+
+    def step(self):
+        with self.queue_timer:
+            ra, replay = self.inqueue.get()
+        if replay is not None:
+            with self.grad_timer:
+                td_error = self.local_evaluator.compute_apply(replay)
+            self.outqueue.put((ra, replay, td_error))
+        self.learner_queue_size.push(self.inqueue.qsize())
+
+
+class ApexOptimizer(Optimizer):
+
+    def _init(
+            self, learning_starts=1000, buffer_size=10000,
+            prioritized_replay=True, prioritized_replay_alpha=0.6,
+            prioritized_replay_beta=0.4, prioritized_replay_eps=1e-6,
+            train_batch_size=512, sample_batch_size=50,
+            num_replay_buffer_shards=1, max_weight_sync_delay=400):
+
+        self.replay_starts = learning_starts
+        self.prioritized_replay_beta = prioritized_replay_beta
+        self.prioritized_replay_eps = prioritized_replay_eps
+        self.train_batch_size = train_batch_size
+        self.sample_batch_size = sample_batch_size
+        self.max_weight_sync_delay = max_weight_sync_delay
+
+        self.learner = GenericLearner(self.local_evaluator)
+        self.learner.start()
+
+        self.replay_actors = create_colocated(
+            ReplayActor,
+            [num_replay_buffer_shards, learning_starts, buffer_size,
+             train_batch_size, prioritized_replay_alpha,
+             prioritized_replay_beta, prioritized_replay_eps],
+            num_replay_buffer_shards)
+        assert len(self.remote_evaluators) > 0
+
+        # Stats
+        self.timers = {k: TimerStat() for k in [
+            "put_weights", "get_samples", "enqueue", "sample_processing",
+            "replay_processing", "update_priorities", "train", "sample"]}
+        self.meters = {k: WindowStat(k, 10) for k in [
+            "samples_per_loop", "replays_per_loop", "reprios_per_loop",
+            "reweights_per_loop"]}
+        self.num_weight_syncs = 0
+        self.learning_started = False
+
+        # Number of worker steps since the last weight update
+        self.steps_since_update = {}
+
+        # Otherwise kick of replay tasks for local gradient updates
+        self.replay_tasks = TaskPool()
+        for ra in self.replay_actors:
+            for _ in range(REPLAY_QUEUE_DEPTH):
+                self.replay_tasks.add(ra, ra.replay.remote())
+
+        # Kick off async background sampling
+        self.sample_tasks = TaskPool()
+        weights = self.local_evaluator.get_weights()
+        for ev in self.remote_evaluators:
+            ev.set_weights.remote(weights)
+            self.steps_since_update[ev] = 0
+            for _ in range(SAMPLE_QUEUE_DEPTH):
+                self.sample_tasks.add(ev, ev.sample.remote())
+
+    def step(self):
+        start = time.time()
+        sample_timesteps, train_timesteps = self._step()
+        time_delta = time.time() - start
+        self.timers["sample"].push(time_delta)
+        self.timers["sample"].push_units_processed(sample_timesteps)
+        if train_timesteps > 0:
+            self.learning_started = True
+        if self.learning_started:
+            self.timers["train"].push(time_delta)
+            self.timers["train"].push_units_processed(train_timesteps)
+        self.num_steps_sampled += sample_timesteps
+        self.num_steps_trained += train_timesteps
+
+    def _step(self):
+        sample_timesteps, train_timesteps = 0, 0
+        weights = None
+
+        with self.timers["sample_processing"]:
+            i = 0
+            num_weight_syncs = 0
+            for ev, sample_batch in self.sample_tasks.completed():
+                i += 1
+                sample_timesteps += self.sample_batch_size
+
+                # Send the data to the replay buffer
+                random.choice(self.replay_actors).add_batch.remote(
+                    sample_batch)
+
+                # Update weights if needed
+                self.steps_since_update[ev] += self.sample_batch_size
+                if self.steps_since_update[ev] >= self.max_weight_sync_delay:
+                    if weights is None:
+                        with self.timers["put_weights"]:
+                            weights = ray.put(
+                                self.local_evaluator.get_weights())
+                    ev.set_weights.remote(weights)
+                    self.num_weight_syncs += 1
+                    num_weight_syncs += 1
+                    self.steps_since_update[ev] = 0
+
+                # Kick off another sample request
+                self.sample_tasks.add(ev, ev.sample.remote())
+            self.meters["samples_per_loop"].push(i)
+            self.meters["reweights_per_loop"].push(num_weight_syncs)
+
+        with self.timers["replay_processing"]:
+            i = 0
+            for ra, replay in self.replay_tasks.completed():
+                i += 1
+                self.replay_tasks.add(ra, ra.replay.remote())
+                with self.timers["get_samples"]:
+                    samples = ray.get(replay)
+                with self.timers["enqueue"]:
+                    self.learner.inqueue.put((ra, samples))
+            self.meters["replays_per_loop"].push(i)
+
+        with self.timers["update_priorities"]:
+            i = 0
+            while not self.learner.outqueue.empty():
+                i += 1
+                ra, replay, td_error = self.learner.outqueue.get()
+                ra.update_priorities.remote(replay, td_error)
+                train_timesteps += self.train_batch_size
+            self.meters["reprios_per_loop"].push(i)
+
+        return sample_timesteps, train_timesteps
+
+    def stats(self):
+        replay_stats = ray.get(self.replay_actors[0].stats.remote())
+        timing = {
+            "{}_time_ms".format(k): round(1000 * self.timers[k].mean, 3)
+            for k in self.timers
+        }
+        timing["learner_grad_time_ms"] = round(
+            1000 * self.learner.grad_timer.mean, 3)
+        timing["learner_dequeue_time_ms"] = round(
+            1000 * self.learner.queue_timer.mean, 3)
+        stats = {
+            "replay_shard_0": replay_stats,
+            "timing_breakdown": timing,
+            "sample_throughput": round(
+                self.timers["sample"].mean_throughput, 3),
+            "train_throughput": round(self.timers["train"].mean_throughput, 3),
+            "num_weight_syncs": self.num_weight_syncs,
+            "pending_sample_tasks": self.sample_tasks.count,
+            "pending_replay_tasks": self.replay_tasks.count,
+            "learner_queue": self.learner.learner_queue_size.stats(),
+            "samples": self.meters["samples_per_loop"].stats(),
+            "replays": self.meters["replays_per_loop"].stats(),
+            "reprios": self.meters["reprios_per_loop"].stats(),
+            "reweights": self.meters["reweights_per_loop"].stats(),
+        }
+        return dict(Optimizer.stats(self), **stats)

python/ray/rllib/optimizers/async.py

@@ -14,11 +14,12 @@ class AsyncOptimizer(Optimizer):
     evaluators, sending updated weights back as needed. This pipelines the
     gradient computations on the remote workers.
     """
-    def _init(self):
+    def _init(self, grads_per_step=100, batch_size=10):
         self.apply_timer = TimerStat()
         self.wait_timer = TimerStat()
         self.dispatch_timer = TimerStat()
-        self.grads_per_step = self.config.get("grads_per_step", 100)
+        self.grads_per_step = grads_per_step
+        self.batch_size = batch_size
 
     def step(self):
         weights = ray.put(self.local_evaluator.get_weights())
@@ -49,9 +50,12 @@ class AsyncOptimizer(Optimizer):
                     gradient_queue.append((fut, e))
                     num_gradients += 1
 
+        self.num_steps_sampled += self.grads_per_step * self.batch_size
+        self.num_steps_trained += self.grads_per_step * self.batch_size
+
     def stats(self):
-        return {
+        return dict(Optimizer.stats(), **{
             "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/evaluator.py

@@ -2,6 +2,8 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
+
 
 class Evaluator(object):
     """Algorithms implement this interface to leverage RLlib optimizers.
@@ -62,6 +64,22 @@ class Evaluator(object):
 
         raise NotImplementedError
 
+    def compute_apply(self, samples):
+        """Fused compute and apply gradients on given samples.
+
+        Returns:
+            The result of calling compute_gradients(samples)
+        """
+
+        grads = self.compute_gradients(samples)
+        self.apply_gradients(grads)
+        return grads
+
+    def get_host(self):
+        """Returns hostname of actor."""
+
+        return os.uname()[1]
+
 
 class TFMultiGPUSupport(Evaluator):
     """The multi-GPU TF optimizer requires additional TF-specific supportt.

python/ray/rllib/optimizers/local_sync.py

@@ -5,6 +5,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.filter import RunningStat
 from ray.rllib.utils.timer import TimerStat
 
 
@@ -16,10 +17,12 @@ class LocalSyncOptimizer(Optimizer):
     model weights are then broadcast to all remote evaluators.
     """
 
-    def _init(self):
+    def _init(self, batch_size=32):
         self.update_weights_timer = TimerStat()
         self.sample_timer = TimerStat()
         self.grad_timer = TimerStat()
+        self.throughput = RunningStat()
+        self.batch_size = batch_size
 
     def step(self):
         with self.update_weights_timer:
@@ -39,10 +42,16 @@ class LocalSyncOptimizer(Optimizer):
         with self.grad_timer:
             grad = self.local_evaluator.compute_gradients(samples)
             self.local_evaluator.apply_gradients(grad)
+            self.grad_timer.push_units_processed(samples.count)
+
+        self.num_steps_sampled += samples.count
+        self.num_steps_trained += samples.count
 
     def stats(self):
-        return {
+        return dict(Optimizer.stats(self), **{
             "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),
-        }
+            "opt_peak_throughput": round(self.grad_timer.mean_throughput, 3),
+            "opt_samples": round(self.grad_timer.mean_units_processed, 3),
+        })

python/ray/rllib/optimizers/local_sync_replay.py

@@ -0,0 +1,109 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+import ray
+from ray.rllib.optimizers.replay_buffer import ReplayBuffer, \
+    PrioritizedReplayBuffer
+from ray.rllib.optimizers.optimizer import Optimizer
+from ray.rllib.optimizers.sample_batch import SampleBatch
+from ray.rllib.utils.filter import RunningStat
+from ray.rllib.utils.timer import TimerStat
+
+
+class LocalSyncReplayOptimizer(Optimizer):
+    """Variant of the local sync optimizer that supports replay (for DQN)."""
+
+    def _init(
+            self, learning_starts=1000, buffer_size=10000,
+            prioritized_replay=True, prioritized_replay_alpha=0.6,
+            prioritized_replay_beta=0.4, prioritized_replay_eps=1e-6,
+            train_batch_size=32, sample_batch_size=4):
+
+        self.replay_starts = learning_starts
+        self.prioritized_replay_beta = prioritized_replay_beta
+        self.prioritized_replay_eps = prioritized_replay_eps
+        self.train_batch_size = train_batch_size
+
+        # Stats
+        self.update_weights_timer = TimerStat()
+        self.sample_timer = TimerStat()
+        self.replay_timer = TimerStat()
+        self.grad_timer = TimerStat()
+        self.throughput = RunningStat()
+
+        # Set up replay buffer
+        if prioritized_replay:
+            self.replay_buffer = PrioritizedReplayBuffer(
+                buffer_size,
+                alpha=prioritized_replay_alpha)
+        else:
+            self.replay_buffer = ReplayBuffer(buffer_size)
+
+        assert buffer_size >= self.replay_starts
+
+    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:
+                batch = SampleBatch.concat_samples(
+                    ray.get(
+                        [e.sample.remote() for e in self.remote_evaluators]))
+            else:
+                batch = self.local_evaluator.sample()
+            for row in batch.rows():
+                self.replay_buffer.add(
+                    row["obs"], row["actions"], row["rewards"], row["new_obs"],
+                    row["dones"], row["weights"])
+
+        if len(self.replay_buffer) >= self.replay_starts:
+            self._optimize()
+
+        self.num_steps_sampled += batch.count
+
+    def _optimize(self):
+        with self.replay_timer:
+            if isinstance(self.replay_buffer, PrioritizedReplayBuffer):
+                (obses_t, actions, rewards, obses_tp1,
+                    dones, weights, batch_indexes) = self.replay_buffer.sample(
+                        self.train_batch_size,
+                        beta=self.prioritized_replay_beta)
+            else:
+                (obses_t, actions, rewards, obses_tp1,
+                    dones) = self.replay_buffer.sample(
+                        self.train_batch_size)
+                weights = np.ones_like(rewards)
+                batch_indexes = - np.ones_like(rewards)
+
+            samples = SampleBatch({
+                "obs": obses_t, "actions": actions, "rewards": rewards,
+                "new_obs": obses_tp1, "dones": dones, "weights": weights,
+                "batch_indexes": batch_indexes})
+
+        with self.grad_timer:
+            td_error = self.local_evaluator.compute_apply(samples)
+            new_priorities = (
+                np.abs(td_error) + self.prioritized_replay_eps)
+            if isinstance(self.replay_buffer, PrioritizedReplayBuffer):
+                self.replay_buffer.update_priorities(
+                    samples["batch_indexes"], new_priorities)
+            self.grad_timer.push_units_processed(samples.count)
+
+        self.num_steps_trained += samples.count
+
+    def stats(self):
+        return dict(Optimizer.stats(self), **{
+            "sample_time_ms": round(1000 * self.sample_timer.mean, 3),
+            "replay_time_ms": round(1000 * self.replay_timer.mean, 3),
+            "grad_time_ms": round(1000 * self.grad_timer.mean, 3),
+            "update_time_ms": round(1000 * self.update_weights_timer.mean, 3),
+            "opt_peak_throughput": round(self.grad_timer.mean_throughput, 3),
+            "opt_samples": round(self.grad_timer.mean_units_processed, 3),
+        })

python/ray/rllib/optimizers/multi_gpu.py

@@ -26,9 +26,11 @@ class LocalMultiGPUOptimizer(Optimizer):
     the TFMultiGPUSupport API.
     """
 
-    def _init(self):
+    def _init(self, sgd_batch_size=128, sgd_stepsize=5e-5, num_sgd_iter=10):
         assert isinstance(self.local_evaluator, TFMultiGPUSupport)
-        self.batch_size = self.config.get("sgd_batch_size", 128)
+        self.batch_size = sgd_batch_size
+        self.sgd_stepsize = sgd_stepsize
+        self.num_sgd_iter = num_sgd_iter
         gpu_ids = ray.get_gpu_ids()
         if not gpu_ids:
             self.devices = ["/cpu:0"]
@@ -51,12 +53,12 @@ class LocalMultiGPUOptimizer(Optimizer):
         tf.get_variable_scope().reuse_variables()
 
         self.par_opt = LocalSyncParallelOptimizer(
-            tf.train.AdamOptimizer(self.config.get("sgd_stepsize", 5e-5)),
+            tf.train.AdamOptimizer(self.sgd_stepsize),
             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()))
+            os.getcwd())
 
         self.sess = self.local_evaluator.sess
         self.sess.run(tf.global_variables_initializer())
@@ -83,7 +85,7 @@ class LocalMultiGPUOptimizer(Optimizer):
                 samples.columns([key for key, _ in self.loss_inputs]))
 
         with self.grad_timer:
-            for i in range(self.config.get("num_sgd_iter", 10)):
+            for i in range(self.num_sgd_iter):
                 batch_index = 0
                 num_batches = (
                     int(tuples_per_device) // int(self.per_device_batch_size))
@@ -96,10 +98,13 @@ class LocalMultiGPUOptimizer(Optimizer):
                         permutation[batch_index] * self.per_device_batch_size)
                     batch_index += 1
 
+        self.num_steps_sampled += samples.count
+        self.num_steps_trained += samples.count
+
     def stats(self):
-        return {
+        return dict(Optimizer.stats(), **{
             "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/optimizer.py

@@ -2,6 +2,8 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import ray
+
 
 class Optimizer(object):
     """RLlib optimizers encapsulate distributed RL optimization strategies.
@@ -16,20 +18,45 @@ class Optimizer(object):
     environment and compute model gradient updates.
     """
 
+    @classmethod
+    def make(
+            cls, evaluator_cls, evaluator_args, num_workers, optimizer_config):
+        """Create evaluators and an optimizer instance using those evaluators.
+
+        Args:
+            evaluator_cls (class): Python class of the evaluators to create.
+            evaluator_args (list): List of constructor args for the evaluators.
+            num_workers (int): Number of remote evaluators to create in
+                addition to a local evaluator. This can be zero or greater.
+            optimizer_config (dict): Keyword arguments to pass to the
+                optimizer class constructor.
+        """
+
+        local_evaluator = evaluator_cls(*evaluator_args)
+        remote_cls = ray.remote(num_cpus=1)(evaluator_cls)
+        remote_evaluators = [
+            remote_cls.remote(*evaluator_args)
+            for _ in range(num_workers)]
+        return cls(optimizer_config, local_evaluator, remote_evaluators)
+
     def __init__(self, config, local_evaluator, remote_evaluators):
         """Create an optimizer instance.
 
         Args:
-            config (dict): Optimizer-specific configuration data.
+            config (dict): Optimizer-specific arguments.
             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.
+            remote_evaluators (list): A list of Ray actor handles to remote
+                evaluators instances. If empty, the optimizer should fall back
+                to using only the local evaluator.
         """
         self.config = config
         self.local_evaluator = local_evaluator
         self.remote_evaluators = remote_evaluators
-        self._init()
+        self._init(**config)
+
+        # Counters that should be updated by sub-classes
+        self.num_steps_trained = 0
+        self.num_steps_sampled = 0
 
     def _init(self):
         pass
@@ -42,4 +69,14 @@ class Optimizer(object):
     def stats(self):
         """Returns a dictionary of internal performance statistics."""
 
-        return {}
+        return {
+            "num_steps_trained": self.num_steps_trained,
+            "num_steps_sampled": self.num_steps_sampled,
+        }
+
+    def save(self):
+        return [self.num_steps_trained, self.num_steps_sampled]
+
+    def restore(self, data):
+        self.num_steps_trained = data[0]
+        self.num_steps_sampled = data[1]

python/ray/rllib/optimizers/replay_buffer.py

@@ -0,0 +1,237 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import random
+import sys
+
+from ray.rllib.optimizers.segment_tree import SumSegmentTree, MinSegmentTree
+from ray.rllib.utils.compression import unpack
+from ray.rllib.utils.window_stat import WindowStat
+
+
+class ReplayBuffer(object):
+    def __init__(self, size):
+        """Create Prioritized Replay buffer.
+
+        Parameters
+        ----------
+        size: int
+          Max number of transitions to store in the buffer. When the buffer
+          overflows the old memories are dropped.
+        """
+        self._storage = []
+        self._maxsize = size
+        self._next_idx = 0
+        self._hit_count = np.zeros(size)
+        self._eviction_started = False
+        self._num_added = 0
+        self._num_sampled = 0
+        self._evicted_hit_stats = WindowStat("evicted_hit", 1000)
+        self._est_size_bytes = 0
+
+    def __len__(self):
+        return len(self._storage)
+
+    def add(self, obs_t, action, reward, obs_tp1, done, weight):
+        data = (obs_t, action, reward, obs_tp1, done)
+        self._num_added += 1
+
+        if self._next_idx >= len(self._storage):
+            self._storage.append(data)
+            self._est_size_bytes += sum([sys.getsizeof(d) for d in data])
+        else:
+            self._storage[self._next_idx] = data
+        if self._next_idx + 1 >= self._maxsize:
+            self._eviction_started = True
+        self._next_idx = (self._next_idx + 1) % self._maxsize
+        if self._eviction_started:
+            self._evicted_hit_stats.push(self._hit_count[self._next_idx])
+            self._hit_count[self._next_idx] = 0
+
+    def _encode_sample(self, idxes):
+        obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
+        for i in idxes:
+            data = self._storage[i]
+            obs_t, action, reward, obs_tp1, done = data
+            obses_t.append(np.array(unpack(obs_t), copy=False))
+            actions.append(np.array(action, copy=False))
+            rewards.append(reward)
+            obses_tp1.append(np.array(unpack(obs_tp1), copy=False))
+            dones.append(done)
+            self._hit_count[i] += 1
+        return (np.array(obses_t), np.array(actions), np.array(rewards),
+                np.array(obses_tp1), np.array(dones))
+
+    def sample(self, batch_size):
+        """Sample a batch of experiences.
+
+        Parameters
+        ----------
+        batch_size: int
+            How many transitions to sample.
+
+        Returns
+        -------
+        obs_batch: np.array
+          batch of observations
+        act_batch: np.array
+          batch of actions executed given obs_batch
+        rew_batch: np.array
+          rewards received as results of executing act_batch
+        next_obs_batch: np.array
+          next set of observations seen after executing act_batch
+        done_mask: np.array
+          done_mask[i] = 1 if executing act_batch[i] resulted in
+          the end of an episode and 0 otherwise.
+        """
+        idxes = [random.randint(0, len(self._storage) - 1)
+                 for _ in range(batch_size)]
+        self._num_sampled += batch_size
+        return self._encode_sample(idxes)
+
+    def stats(self):
+        data = {
+            "added_count": self._num_added,
+            "sampled_count": self._num_sampled,
+            "est_size_bytes": self._est_size_bytes,
+            "num_entries": len(self._storage),
+        }
+        data.update(self._evicted_hit_stats.stats())
+        return data
+
+
+class PrioritizedReplayBuffer(ReplayBuffer):
+    def __init__(self, size, alpha):
+        """Create Prioritized Replay buffer.
+
+        Parameters
+        ----------
+        size: int
+          Max number of transitions to store in the buffer. When the buffer
+          overflows the old memories are dropped.
+        alpha: float
+          how much prioritization is used
+          (0 - no prioritization, 1 - full prioritization)
+
+        See Also
+        --------
+        ReplayBuffer.__init__
+        """
+        super(PrioritizedReplayBuffer, self).__init__(size)
+        assert alpha > 0
+        self._alpha = alpha
+
+        it_capacity = 1
+        while it_capacity < size:
+            it_capacity *= 2
+
+        self._it_sum = SumSegmentTree(it_capacity)
+        self._it_min = MinSegmentTree(it_capacity)
+        self._max_priority = 1.0
+        self._prio_change_stats = WindowStat("reprio", 1000)
+
+    def add(self, obs_t, action, reward, obs_tp1, done, weight):
+        """See ReplayBuffer.store_effect"""
+        idx = self._next_idx
+        super(PrioritizedReplayBuffer, self).add(
+            obs_t, action, reward, obs_tp1, done, weight)
+        if weight is None:
+            weight = self._max_priority
+        self._it_sum[idx] = weight ** self._alpha
+        self._it_min[idx] = weight ** self._alpha
+
+    def _sample_proportional(self, batch_size):
+        res = []
+        for _ in range(batch_size):
+            # TODO(szymon): should we ensure no repeats?
+            mass = random.random() * self._it_sum.sum(0,
+                                                      len(self._storage) - 1)
+            idx = self._it_sum.find_prefixsum_idx(mass)
+            res.append(idx)
+        return res
+
+    def sample(self, batch_size, beta):
+        """Sample a batch of experiences.
+
+        compared to ReplayBuffer.sample
+        it also returns importance weights and idxes
+        of sampled experiences.
+
+
+        Parameters
+        ----------
+        batch_size: int
+          How many transitions to sample.
+        beta: float
+          To what degree to use importance weights
+          (0 - no corrections, 1 - full correction)
+
+        Returns
+        -------
+        obs_batch: np.array
+          batch of observations
+        act_batch: np.array
+          batch of actions executed given obs_batch
+        rew_batch: np.array
+          rewards received as results of executing act_batch
+        next_obs_batch: np.array
+          next set of observations seen after executing act_batch
+        done_mask: np.array
+          done_mask[i] = 1 if executing act_batch[i] resulted in
+          the end of an episode and 0 otherwise.
+        weights: np.array
+          Array of shape (batch_size,) and dtype np.float32
+          denoting importance weight of each sampled transition
+        idxes: np.array
+          Array of shape (batch_size,) and dtype np.int32
+          idexes in buffer of sampled experiences
+        """
+        assert beta > 0
+        self._num_sampled += batch_size
+
+        idxes = self._sample_proportional(batch_size)
+
+        weights = []
+        p_min = self._it_min.min() / self._it_sum.sum()
+        max_weight = (p_min * len(self._storage)) ** (-beta)
+
+        for idx in idxes:
+            p_sample = self._it_sum[idx] / self._it_sum.sum()
+            weight = (p_sample * len(self._storage)) ** (-beta)
+            weights.append(weight / max_weight)
+        weights = np.array(weights)
+        encoded_sample = self._encode_sample(idxes)
+        return tuple(list(encoded_sample) + [weights, idxes])
+
+    def update_priorities(self, idxes, priorities):
+        """Update priorities of sampled transitions.
+
+        sets priority of transition at index idxes[i] in buffer
+        to priorities[i].
+
+        Parameters
+        ----------
+        idxes: [int]
+          List of idxes of sampled transitions
+        priorities: [float]
+          List of updated priorities corresponding to
+          transitions at the sampled idxes denoted by
+          variable `idxes`.
+        """
+        assert len(idxes) == len(priorities)
+        for idx, priority in zip(idxes, priorities):
+            assert priority > 0
+            assert 0 <= idx < len(self._storage)
+            delta = priority ** self._alpha - self._it_sum[idx]
+            self._prio_change_stats.push(delta)
+            self._it_sum[idx] = priority ** self._alpha
+            self._it_min[idx] = priority ** self._alpha
+
+            self._max_priority = max(self._max_priority, priority)
+
+    def stats(self):
+        parent = ReplayBuffer.stats(self)
+        parent.update(self._prio_change_stats.stats())
+        return parent

python/ray/rllib/optimizers/sample_batch.py

@@ -37,7 +37,7 @@ class SampleBatch(object):
     def concat_samples(samples):
         out = {}
         for k in samples[0].data.keys():
-            out[k] = np.concatenate([arrayify(s.data[k]) for s in samples])
+            out[k] = np.concatenate([s.data[k] for s in samples])
         return SampleBatch(out)
 
     def concat(self, other):

python/ray/rllib/optimizers/segment_tree.py

@@ -0,0 +1,153 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import operator
+
+
+class SegmentTree(object):
+    def __init__(self, capacity, operation, neutral_element):
+        """Build a Segment Tree data structure.
+
+        https://en.wikipedia.org/wiki/Segment_tree
+
+        Can be used as regular array, but with two
+        important differences:
+
+          a) setting item's value is slightly slower.
+             It is O(lg capacity) instead of O(1).
+          b) user has access to an efficient `reduce`
+             operation which reduces `operation` over
+             a contiguous subsequence of items in the
+             array.
+
+        Paramters
+        ---------
+        capacity: int
+          Total size of the array - must be a power of two.
+        operation: lambda obj, obj -> obj
+          and operation for combining elements (eg. sum, max)
+          must for a mathematical group together with the set of
+          possible values for array elements.
+        neutral_element: obj
+          neutral element for the operation above. eg. float('-inf')
+          for max and 0 for sum.
+        """
+
+        assert capacity > 0 and capacity & (capacity - 1) == 0, \
+            "capacity must be positive and a power of 2."
+        self._capacity = capacity
+        self._value = [neutral_element for _ in range(2 * capacity)]
+        self._operation = operation
+
+    def _reduce_helper(self, start, end, node, node_start, node_end):
+        if start == node_start and end == node_end:
+            return self._value[node]
+        mid = (node_start + node_end) // 2
+        if end <= mid:
+            return self._reduce_helper(start, end, 2 * node, node_start, mid)
+        else:
+            if mid + 1 <= start:
+                return self._reduce_helper(start, end, 2 * node + 1, mid + 1,
+                                           node_end)
+            else:
+                return self._operation(
+                    self._reduce_helper(start, mid, 2 * node, node_start, mid),
+                    self._reduce_helper(mid + 1, end, 2 * node + 1, mid + 1,
+                                        node_end)
+                )
+
+    def reduce(self, start=0, end=None):
+        """Returns result of applying `self.operation`
+        to a contiguous subsequence of the array.
+
+          self.operation(
+              arr[start], operation(arr[start+1], operation(... arr[end])))
+
+        Parameters
+        ----------
+        start: int
+          beginning of the subsequence
+        end: int
+          end of the subsequences
+
+        Returns
+        -------
+        reduced: obj
+          result of reducing self.operation over the specified range of array
+          elements.
+        """
+        if end is None:
+            end = self._capacity
+        if end < 0:
+            end += self._capacity
+        end -= 1
+        return self._reduce_helper(start, end, 1, 0, self._capacity - 1)
+
+    def __setitem__(self, idx, val):
+        # index of the leaf
+        idx += self._capacity
+        self._value[idx] = val
+        idx //= 2
+        while idx >= 1:
+            self._value[idx] = self._operation(
+                self._value[2 * idx],
+                self._value[2 * idx + 1])
+            idx //= 2
+
+    def __getitem__(self, idx):
+        assert 0 <= idx < self._capacity
+        return self._value[self._capacity + idx]
+
+
+class SumSegmentTree(SegmentTree):
+    def __init__(self, capacity):
+        super(SumSegmentTree, self).__init__(
+            capacity=capacity,
+            operation=operator.add,
+            neutral_element=0.0)
+
+    def sum(self, start=0, end=None):
+        """Returns arr[start] + ... + arr[end]"""
+        return super(SumSegmentTree, self).reduce(start, end)
+
+    def find_prefixsum_idx(self, prefixsum):
+        """Find the highest index `i` in the array such that
+          sum(arr[0] + arr[1] + ... + arr[i - i]) <= prefixsum
+
+        if array values are probabilities, this function
+        allows to sample indexes according to the discrete
+        probability efficiently.
+
+        Parameters
+        ----------
+        perfixsum: float
+          upperbound on the sum of array prefix
+
+        Returns
+        -------
+        idx: int
+          highest index satisfying the prefixsum constraint
+        """
+        assert 0 <= prefixsum <= self.sum() + 1e-5
+        idx = 1
+        while idx < self._capacity:  # while non-leaf
+            if self._value[2 * idx] > prefixsum:
+                idx = 2 * idx
+            else:
+                prefixsum -= self._value[2 * idx]
+                idx = 2 * idx + 1
+        return idx - self._capacity
+
+
+class MinSegmentTree(SegmentTree):
+    def __init__(self, capacity):
+        super(MinSegmentTree, self).__init__(
+            capacity=capacity,
+            operation=min,
+            neutral_element=float('inf'))
+
+    def min(self, start=0, end=None):
+        """Returns min(arr[start], ...,  arr[end])"""
+
+        return super(MinSegmentTree, self).reduce(start, end)

python/ray/rllib/optimizers/tests/test_segment_tree.py

@@ -0,0 +1,107 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from ray.rllib.optimizers.segment_tree import SumSegmentTree, MinSegmentTree
+
+
+def test_tree_set():
+    tree = SumSegmentTree(4)
+
+    tree[2] = 1.0
+    tree[3] = 3.0
+
+    assert np.isclose(tree.sum(), 4.0)
+    assert np.isclose(tree.sum(0, 2), 0.0)
+    assert np.isclose(tree.sum(0, 3), 1.0)
+    assert np.isclose(tree.sum(2, 3), 1.0)
+    assert np.isclose(tree.sum(2, -1), 1.0)
+    assert np.isclose(tree.sum(2, 4), 4.0)
+
+
+def test_tree_set_overlap():
+    tree = SumSegmentTree(4)
+
+    tree[2] = 1.0
+    tree[2] = 3.0
+
+    assert np.isclose(tree.sum(), 3.0)
+    assert np.isclose(tree.sum(2, 3), 3.0)
+    assert np.isclose(tree.sum(2, -1), 3.0)
+    assert np.isclose(tree.sum(2, 4), 3.0)
+    assert np.isclose(tree.sum(1, 2), 0.0)
+
+
+def test_prefixsum_idx():
+    tree = SumSegmentTree(4)
+
+    tree[2] = 1.0
+    tree[3] = 3.0
+
+    assert tree.find_prefixsum_idx(0.0) == 2
+    assert tree.find_prefixsum_idx(0.5) == 2
+    assert tree.find_prefixsum_idx(0.99) == 2
+    assert tree.find_prefixsum_idx(1.01) == 3
+    assert tree.find_prefixsum_idx(3.00) == 3
+    assert tree.find_prefixsum_idx(4.00) == 3
+
+
+def test_prefixsum_idx2():
+    tree = SumSegmentTree(4)
+
+    tree[0] = 0.5
+    tree[1] = 1.0
+    tree[2] = 1.0
+    tree[3] = 3.0
+
+    assert tree.find_prefixsum_idx(0.00) == 0
+    assert tree.find_prefixsum_idx(0.55) == 1
+    assert tree.find_prefixsum_idx(0.99) == 1
+    assert tree.find_prefixsum_idx(1.51) == 2
+    assert tree.find_prefixsum_idx(3.00) == 3
+    assert tree.find_prefixsum_idx(5.50) == 3
+
+
+def test_max_interval_tree():
+    tree = MinSegmentTree(4)
+
+    tree[0] = 1.0
+    tree[2] = 0.5
+    tree[3] = 3.0
+
+    assert np.isclose(tree.min(), 0.5)
+    assert np.isclose(tree.min(0, 2), 1.0)
+    assert np.isclose(tree.min(0, 3), 0.5)
+    assert np.isclose(tree.min(0, -1), 0.5)
+    assert np.isclose(tree.min(2, 4), 0.5)
+    assert np.isclose(tree.min(3, 4), 3.0)
+
+    tree[2] = 0.7
+
+    assert np.isclose(tree.min(), 0.7)
+    assert np.isclose(tree.min(0, 2), 1.0)
+    assert np.isclose(tree.min(0, 3), 0.7)
+    assert np.isclose(tree.min(0, -1), 0.7)
+    assert np.isclose(tree.min(2, 4), 0.7)
+    assert np.isclose(tree.min(3, 4), 3.0)
+
+    tree[2] = 4.0
+
+    assert np.isclose(tree.min(), 1.0)
+    assert np.isclose(tree.min(0, 2), 1.0)
+    assert np.isclose(tree.min(0, 3), 1.0)
+    assert np.isclose(tree.min(0, -1), 1.0)
+    assert np.isclose(tree.min(2, 4), 3.0)
+    assert np.isclose(tree.min(2, 3), 4.0)
+    assert np.isclose(tree.min(2, -1), 4.0)
+    assert np.isclose(tree.min(3, 4), 3.0)
+
+
+if __name__ == '__main__':
+    test_tree_set()
+    test_tree_set_overlap()
+    test_prefixsum_idx()
+    test_prefixsum_idx2()
+    test_max_interval_tree()