[rllib] Parallel-data loading and multi-gpu support for IMPALA (#2766)

python/ray/rllib/agents/impala/impala.py

@@ -11,8 +11,16 @@ from ray.rllib.optimizers import AsyncSamplesOptimizer
 from ray.tune.trial import Resources
 
 OPTIMIZER_SHARED_CONFIGS = [
+    "lr",
+    "num_envs_per_worker",
+    "num_gpus",
     "sample_batch_size",
     "train_batch_size",
+    "replay_buffer_num_slots",
+    "replay_proportion",
+    "num_parallel_data_loaders",
+    "grad_clip",
+    "max_sample_requests_in_flight_per_worker",
 ]
 
 DEFAULT_CONFIG = with_common_config({
@@ -25,10 +33,22 @@ DEFAULT_CONFIG = with_common_config({
     "sample_batch_size": 50,
     "train_batch_size": 500,
     "min_iter_time_s": 10,
-    "gpu": True,
     "num_workers": 2,
     "num_cpus_per_worker": 1,
     "num_gpus_per_worker": 0,
+    # number of GPUs the learner should use.
+    "num_gpus": 1,
+    # set >1 to load data into GPUs in parallel. Increases GPU memory usage
+    # proportionally with the number of loaders.
+    "num_parallel_data_loaders": 1,
+    # level of queuing for sampling.
+    "max_sample_requests_in_flight_per_worker": 2,
+    # set >0 to enable experience replay. Saved samples will be replayed with
+    # a p:1 proportion to new data samples.
+    "replay_proportion": 0.0,
+    # number of sample batches to store for replay. The number of transitions
+    # saved total will be (replay_buffer_num_slots * sample_batch_size).
+    "replay_buffer_num_slots": 100,
 
     # Learning params.
     "grad_clip": 40.0,
@@ -65,7 +85,7 @@ class ImpalaAgent(Agent):
         cf = dict(cls._default_config, **config)
         return Resources(
             cpu=1,
-            gpu=cf["gpu"] and cf["gpu_fraction"] or 0,
+            gpu=cf["num_gpus"] and cf["num_gpus"] * cf["gpu_fraction"] or 0,
             extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
             extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
 

python/ray/rllib/agents/impala/vtrace_policy_graph.py

@@ -31,6 +31,7 @@ class VTraceLoss(object):
                  rewards,
                  values,
                  bootstrap_value,
+                 valid_mask,
                  vf_loss_coeff=0.5,
                  entropy_coeff=-0.01,
                  clip_rho_threshold=1.0,
@@ -52,6 +53,7 @@ class VTraceLoss(object):
             rewards: A float32 tensor of shape [T, B].
             values: A float32 tensor of shape [T, B].
             bootstrap_value: A float32 tensor of shape [B].
+            valid_mask: A bool tensor of valid RNN input elements (#2992).
         """
 
         # Compute vtrace on the CPU for better perf.
@@ -70,14 +72,16 @@ class VTraceLoss(object):
 
         # The policy gradients loss
         self.pi_loss = -tf.reduce_sum(
-            actions_logp * self.vtrace_returns.pg_advantages)
+            tf.boolean_mask(actions_logp * self.vtrace_returns.pg_advantages,
+                            valid_mask))
 
         # The baseline loss
-        delta = values - self.vtrace_returns.vs
+        delta = tf.boolean_mask(values - self.vtrace_returns.vs, valid_mask)
         self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta))
 
         # The entropy loss
-        self.entropy = tf.reduce_sum(actions_entropy)
+        self.entropy = tf.reduce_sum(
+            tf.boolean_mask(actions_entropy, valid_mask))
 
         # The summed weighted loss
         self.total_loss = (self.pi_loss + self.vf_loss * vf_loss_coeff +
@@ -85,20 +89,49 @@ class VTraceLoss(object):
 
 
 class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
-    def __init__(self, observation_space, action_space, config):
+    def __init__(self,
+                 observation_space,
+                 action_space,
+                 config,
+                 existing_inputs=None):
         config = dict(ray.rllib.agents.impala.impala.DEFAULT_CONFIG, **config)
         assert config["batch_mode"] == "truncate_episodes", \
             "Must use `truncate_episodes` batch mode with V-trace."
         self.config = config
         self.sess = tf.get_default_session()
 
+        # Create input placeholders
+        if existing_inputs:
+            actions, dones, behaviour_logits, rewards, observations = \
+                existing_inputs[:5]
+            existing_state_in = existing_inputs[5:-1]
+            existing_seq_lens = existing_inputs[-1]
+        else:
+            if isinstance(action_space, gym.spaces.Discrete):
+                ac_size = action_space.n
+                actions = tf.placeholder(tf.int64, [None], name="ac")
+            else:
+                raise UnsupportedSpaceException(
+                    "Action space {} is not supported for IMPALA.".format(
+                        action_space))
+            dones = tf.placeholder(tf.bool, [None], name="dones")
+            rewards = tf.placeholder(tf.float32, [None], name="rewards")
+            behaviour_logits = tf.placeholder(
+                tf.float32, [None, ac_size], name="behaviour_logits")
+            observations = tf.placeholder(
+                tf.float32, [None] + list(observation_space.shape))
+            existing_state_in = None
+            existing_seq_lens = None
+
         # Setup the policy
-        self.observations = tf.placeholder(
-            tf.float32, [None] + list(observation_space.shape))
         dist_class, logit_dim = ModelCatalog.get_action_dist(
             action_space, self.config["model"])
-        self.model = ModelCatalog.get_model(self.observations, logit_dim,
-                                            self.config["model"])
+        self.model = ModelCatalog.get_model(
+            observations,
+            logit_dim,
+            self.config["model"],
+            state_in=existing_state_in,
+            seq_lens=existing_seq_lens)
         action_dist = dist_class(self.model.outputs)
         values = tf.reshape(
             linear(self.model.last_layer, 1, "value", normc_initializer(1.0)),
@@ -106,19 +139,6 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
         self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                           tf.get_variable_scope().name)
 
-        # Setup the policy loss
-        if isinstance(action_space, gym.spaces.Discrete):
-            ac_size = action_space.n
-            actions = tf.placeholder(tf.int64, [None], name="ac")
-        else:
-            raise UnsupportedSpaceException(
-                "Action space {} is not supported for IMPALA.".format(
-                    action_space))
-        dones = tf.placeholder(tf.bool, [None], name="dones")
-        rewards = tf.placeholder(tf.float32, [None], name="rewards")
-        behaviour_logits = tf.placeholder(
-            tf.float32, [None, ac_size], name="behaviour_logits")
-
         def to_batches(tensor):
             if self.config["model"]["use_lstm"]:
                 B = tf.shape(self.model.seq_lens)[0]
@@ -135,6 +155,13 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
                 rs,
                 [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0]))))
 
+        if self.model.state_in:
+            max_seq_len = tf.reduce_max(self.model.seq_lens) - 1
+            mask = tf.sequence_mask(self.model.seq_lens, max_seq_len)
+            mask = tf.reshape(mask, [-1])
+        else:
+            mask = tf.ones_like(rewards)
+
         # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc.
         self.loss = VTraceLoss(
             actions=to_batches(actions)[:-1],
@@ -147,6 +174,7 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
             rewards=to_batches(rewards)[:-1],
             values=to_batches(values)[:-1],
             bootstrap_value=to_batches(values)[-1],
+            valid_mask=to_batches(mask)[:-1],
             vf_loss_coeff=self.config["vf_loss_coeff"],
             entropy_coeff=self.config["entropy_coeff"],
             clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
@@ -158,7 +186,7 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
             ("dones", dones),
             ("behaviour_logits", behaviour_logits),
             ("rewards", rewards),
-            ("obs", self.observations),
+            ("obs", observations),
         ]
         LearningRateSchedule.__init__(self, self.config["lr"],
                                       self.config["lr_schedule"])
@@ -167,7 +195,7 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
             observation_space,
             action_space,
             self.sess,
-            obs_input=self.observations,
+            obs_input=observations,
             action_sampler=action_dist.sample(),
             loss=self.loss.total_loss,
             loss_inputs=loss_in,
@@ -218,3 +246,10 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
 
     def get_initial_state(self):
         return self.model.state_init
+
+    def copy(self, existing_inputs):
+        return VTracePolicyGraph(
+            self.observation_space,
+            self.action_space,
+            self.config,
+            existing_inputs=existing_inputs)

python/ray/rllib/agents/ppo/ppo_policy_graph.py

@@ -24,6 +24,7 @@ class PPOLoss(object):
                  curr_action_dist,
                  value_fn,
                  cur_kl_coeff,
+                 valid_mask,
                  entropy_coeff=0,
                  clip_param=0.1,
                  vf_clip_param=0.1,
@@ -48,28 +49,33 @@ class PPOLoss(object):
             value_fn (Tensor): Current value function output Tensor.
             cur_kl_coeff (Variable): Variable holding the current PPO KL
                 coefficient.
+            valid_mask (Tensor): A bool mask of valid input elements (#2992).
             entropy_coeff (float): Coefficient of the entropy regularizer.
             clip_param (float): Clip parameter
             vf_clip_param (float): Clip parameter for the value function
             vf_loss_coeff (float): Coefficient of the value function loss
             use_gae (bool): If true, use the Generalized Advantage Estimator.
         """
+
+        def reduce_mean_valid(t):
+            return tf.reduce_mean(tf.boolean_mask(t, valid_mask))
+
         dist_cls, _ = ModelCatalog.get_action_dist(action_space, {})
         prev_dist = dist_cls(logits)
         # Make loss functions.
         logp_ratio = tf.exp(
             curr_action_dist.logp(actions) - prev_dist.logp(actions))
         action_kl = prev_dist.kl(curr_action_dist)
-        self.mean_kl = tf.reduce_mean(action_kl)
+        self.mean_kl = reduce_mean_valid(action_kl)
 
         curr_entropy = curr_action_dist.entropy()
-        self.mean_entropy = tf.reduce_mean(curr_entropy)
+        self.mean_entropy = reduce_mean_valid(curr_entropy)
 
         surrogate_loss = tf.minimum(
             advantages * logp_ratio,
             advantages * tf.clip_by_value(logp_ratio, 1 - clip_param,
                                           1 + clip_param))
-        self.mean_policy_loss = tf.reduce_mean(-surrogate_loss)
+        self.mean_policy_loss = reduce_mean_valid(-surrogate_loss)
 
         if use_gae:
             vf_loss1 = tf.square(value_fn - value_targets)
@@ -77,14 +83,15 @@ class PPOLoss(object):
                 value_fn - vf_preds, -vf_clip_param, vf_clip_param)
             vf_loss2 = tf.square(vf_clipped - value_targets)
             vf_loss = tf.maximum(vf_loss1, vf_loss2)
-            self.mean_vf_loss = tf.reduce_mean(vf_loss)
-            loss = tf.reduce_mean(-surrogate_loss + cur_kl_coeff * action_kl +
-                                  vf_loss_coeff * vf_loss -
-                                  entropy_coeff * curr_entropy)
+            self.mean_vf_loss = reduce_mean_valid(vf_loss)
+            loss = reduce_mean_valid(
+                -surrogate_loss + cur_kl_coeff * action_kl +
+                vf_loss_coeff * vf_loss - entropy_coeff * curr_entropy)
         else:
             self.mean_vf_loss = tf.constant(0.0)
-            loss = tf.reduce_mean(-surrogate_loss + cur_kl_coeff * action_kl -
-                                  entropy_coeff * curr_entropy)
+            loss = reduce_mean_valid(-surrogate_loss +
+                                     cur_kl_coeff * action_kl -
+                                     entropy_coeff * curr_entropy)
         self.loss = loss
 
 
@@ -179,6 +186,13 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
         else:
             self.value_function = tf.zeros(shape=tf.shape(obs_ph)[:1])
 
+        if self.model.state_in:
+            max_seq_len = tf.reduce_max(self.model.seq_lens)
+            mask = tf.sequence_mask(self.model.seq_lens, max_seq_len)
+            mask = tf.reshape(mask, [-1])
+        else:
+            mask = tf.ones_like(adv_ph)
+
         self.loss_obj = PPOLoss(
             action_space,
             value_targets_ph,
@@ -189,6 +203,7 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             curr_action_dist,
             self.value_function,
             self.kl_coeff,
+            mask,
             entropy_coeff=self.config["entropy_coeff"],
             clip_param=self.config["clip_param"],
             vf_clip_param=self.config["vf_clip_param"],
@@ -227,7 +242,7 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
     def copy(self, existing_inputs):
         """Creates a copy of self using existing input placeholders."""
         return PPOPolicyGraph(
-            None,
+            self.observation_space,
             self.action_space,
             self.config,
             existing_inputs=existing_inputs)

python/ray/rllib/examples/cartpole_lstm.py

@@ -14,6 +14,7 @@ import numpy as np
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--stop", type=int, default=200)
+parser.add_argument("--run", type=str, default="PPO")
 
 
 class CartPoleStatelessEnv(gym.Env):
@@ -163,18 +164,29 @@ if __name__ == "__main__":
     tune.register_env("cartpole_stateless", lambda _: CartPoleStatelessEnv())
 
     ray.init()
+
+    configs = {
+        "PPO": {
+            "num_sgd_iter": 5,
+        },
+        "IMPALA": {
+            "num_workers": 2,
+            "num_gpus": 0,
+            "vf_loss_coeff": 0.01,
+        },
+    }
+
     tune.run_experiments({
         "test": {
             "env": "cartpole_stateless",
-            "run": "PPO",
+            "run": args.run,
             "stop": {
                 "episode_reward_mean": args.stop
             },
-            "config": {
-                "num_sgd_iter": 5,
+            "config": dict(configs[args.run], **{
                 "model": {
                     "use_lstm": True,
                 },
-            },
+            }),
         }
     })

python/ray/rllib/optimizers/async_replay_optimizer.py

@@ -87,14 +87,14 @@ class ReplayActor(object):
             new_priorities = (np.abs(td_errors) + self.prioritized_replay_eps)
             self.replay_buffer.update_priorities(batch_indexes, new_priorities)
 
-    def stats(self):
+    def stats(self, debug=False):
         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())
+        stat.update(self.replay_buffer.stats(debug=debug))
         return stat
 
 
@@ -274,7 +274,7 @@ class AsyncReplayOptimizer(PolicyOptimizer):
         return sample_timesteps, train_timesteps
 
     def stats(self):
-        replay_stats = ray.get(self.replay_actors[0].stats.remote())
+        replay_stats = ray.get(self.replay_actors[0].stats.remote(self.debug))
         timing = {
             "{}_time_ms".format(k): round(1000 * self.timers[k].mean, 3)
             for k in self.timers
@@ -288,13 +288,13 @@ class AsyncReplayOptimizer(PolicyOptimizer):
                                        3),
             "train_throughput": round(self.timers["train"].mean_throughput, 3),
             "num_weight_syncs": self.num_weight_syncs,
+            "learner_queue": self.learner.learner_queue_size.stats(),
+            "replay_shard_0": replay_stats,
         }
         debug_stats = {
-            "replay_shard_0": replay_stats,
             "timing_breakdown": timing,
             "pending_sample_tasks": self.sample_tasks.count,
             "pending_replay_tasks": self.replay_tasks.count,
-            "learner_queue": self.learner.learner_queue_size.stats(),
         }
         if self.debug:
             stats.update(debug_stats)

python/ray/rllib/optimizers/async_samples_optimizer.py

@@ -6,19 +6,22 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+import random
 import time
 import threading
 
 from six.moves import queue
 
 import ray
+from ray.rllib.optimizers.multi_gpu_impl import LocalSyncParallelOptimizer
 from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer
 from ray.rllib.utils.actors import TaskPool
 from ray.rllib.utils.timer import TimerStat
 from ray.rllib.utils.window_stat import WindowStat
 
-SAMPLE_QUEUE_DEPTH = 2
 LEARNER_QUEUE_MAX_SIZE = 16
+NUM_DATA_LOAD_THREADS = 16
 
 
 class LearnerThread(threading.Thread):
@@ -38,8 +41,10 @@ class LearnerThread(threading.Thread):
         self.outqueue = queue.Queue()
         self.queue_timer = TimerStat()
         self.grad_timer = TimerStat()
+        self.load_timer = TimerStat()
+        self.load_wait_timer = TimerStat()
         self.daemon = True
-        self.weights_updated = 0
+        self.weights_updated = False
         self.stats = {}
 
     def run(self):
@@ -48,18 +53,129 @@ class LearnerThread(threading.Thread):
 
     def step(self):
         with self.queue_timer:
-            ra, batch = self.inqueue.get()
+            batch = self.inqueue.get()
 
-        if batch is not None:
-            with self.grad_timer:
-                fetches = self.local_evaluator.compute_apply(batch)
-                self.weights_updated += 1
-                if "stats" in fetches:
-                    self.stats = fetches["stats"]
-            self.outqueue.put(batch.count)
+        with self.grad_timer:
+            fetches = self.local_evaluator.compute_apply(batch)
+            self.weights_updated = True
+            self.stats = fetches.get("stats", {})
+
+        self.outqueue.put(batch.count)
         self.learner_queue_size.push(self.inqueue.qsize())
 
 
+class TFMultiGPULearner(LearnerThread):
+    """Learner that can use multiple GPUs and parallel loading."""
+
+    def __init__(self,
+                 local_evaluator,
+                 num_gpus=1,
+                 lr=0.0005,
+                 train_batch_size=500,
+                 grad_clip=40,
+                 num_parallel_data_loaders=1):
+        # Multi-GPU requires TensorFlow to function.
+        import tensorflow as tf
+
+        LearnerThread.__init__(self, local_evaluator)
+        self.lr = lr
+        self.train_batch_size = train_batch_size
+        if not num_gpus:
+            self.devices = ["/cpu:0"]
+        else:
+            self.devices = ["/gpu:{}".format(i) for i in range(num_gpus)]
+            print("TFMultiGPULearner devices", self.devices)
+        assert self.train_batch_size % len(self.devices) == 0
+        assert self.train_batch_size >= len(self.devices), "batch too small"
+        self.policy = self.local_evaluator.policy_map["default"]
+
+        # per-GPU graph copies created below must share vars with the policy
+        # reuse is set to AUTO_REUSE because Adam nodes are created after
+        # all of the device copies are created.
+        self.par_opt = []
+        with self.local_evaluator.tf_sess.graph.as_default():
+            with self.local_evaluator.tf_sess.as_default():
+                with tf.variable_scope("default", reuse=tf.AUTO_REUSE):
+                    if self.policy._state_inputs:
+                        rnn_inputs = self.policy._state_inputs + [
+                            self.policy._seq_lens
+                        ]
+                    else:
+                        rnn_inputs = []
+                    adam = tf.train.AdamOptimizer(self.lr)
+                    for _ in range(num_parallel_data_loaders):
+                        self.par_opt.append(
+                            LocalSyncParallelOptimizer(
+                                adam,
+                                self.devices,
+                                [v for _, v in self.policy.loss_inputs()],
+                                rnn_inputs,
+                                999999,  # it will get rounded down
+                                self.policy.copy,
+                                grad_norm_clipping=grad_clip))
+
+                self.sess = self.local_evaluator.tf_sess
+                self.sess.run(tf.global_variables_initializer())
+
+        self.idle_optimizers = queue.Queue()
+        self.ready_optimizers = queue.Queue()
+        for opt in self.par_opt:
+            self.idle_optimizers.put(opt)
+        for i in range(NUM_DATA_LOAD_THREADS):
+            self.loader_thread = _LoaderThread(self, share_stats=(i == 0))
+            self.loader_thread.start()
+
+    def step(self):
+        assert self.loader_thread.is_alive()
+        with self.load_wait_timer:
+            opt = self.ready_optimizers.get()
+
+        with self.grad_timer:
+            fetches = opt.optimize(self.sess, 0)
+            self.weights_updated = True
+            self.stats = fetches.get("stats", {})
+
+        self.idle_optimizers.put(opt)
+        self.outqueue.put(self.train_batch_size)
+        self.learner_queue_size.push(self.inqueue.qsize())
+
+
+class _LoaderThread(threading.Thread):
+    def __init__(self, learner, share_stats):
+        threading.Thread.__init__(self)
+        self.learner = learner
+        self.daemon = True
+        if share_stats:
+            self.queue_timer = learner.queue_timer
+            self.load_timer = learner.load_timer
+        else:
+            self.queue_timer = TimerStat()
+            self.load_timer = TimerStat()
+
+    def run(self):
+        while True:
+            self.step()
+
+    def step(self):
+        s = self.learner
+        with self.queue_timer:
+            batch = s.inqueue.get()
+
+        opt = s.idle_optimizers.get()
+
+        with self.load_timer:
+            tuples = s.policy._get_loss_inputs_dict(batch)
+            data_keys = [ph for _, ph in s.policy.loss_inputs()]
+            if s.policy._state_inputs:
+                state_keys = s.policy._state_inputs + [s.policy._seq_lens]
+            else:
+                state_keys = []
+            opt.load_data(s.sess, [tuples[k] for k in data_keys],
+                          [tuples[k] for k in state_keys])
+
+        s.ready_optimizers.put(opt)
+
+
 class AsyncSamplesOptimizer(PolicyOptimizer):
     """Main event loop of the IMPALA architecture.
 
@@ -67,13 +183,38 @@ class AsyncSamplesOptimizer(PolicyOptimizer):
     and remote evaluators (IMPALA actors).
     """
 
-    def _init(self, train_batch_size=512, sample_batch_size=50, debug=False):
-
-        self.debug = debug
+    def _init(self,
+              train_batch_size=500,
+              sample_batch_size=50,
+              num_envs_per_worker=1,
+              num_gpus=0,
+              lr=0.0005,
+              grad_clip=40,
+              replay_buffer_num_slots=0,
+              replay_proportion=0.0,
+              num_parallel_data_loaders=1,
+              max_sample_requests_in_flight_per_worker=2):
         self.learning_started = False
         self.train_batch_size = train_batch_size
+        self.sample_batch_size = sample_batch_size
 
-        self.learner = LearnerThread(self.local_evaluator)
+        if num_gpus > 1 or num_parallel_data_loaders > 1:
+            print(
+                "Enabling multi-GPU mode, {} GPUs, {} parallel loaders".format(
+                    num_gpus, num_parallel_data_loaders))
+            if train_batch_size // max(1, num_gpus) % (
+                    sample_batch_size // num_envs_per_worker) != 0:
+                raise ValueError(
+                    "Sample batches must evenly divide across GPUs.")
+            self.learner = TFMultiGPULearner(
+                self.local_evaluator,
+                lr=lr,
+                num_gpus=num_gpus,
+                train_batch_size=train_batch_size,
+                grad_clip=grad_clip,
+                num_parallel_data_loaders=num_parallel_data_loaders)
+        else:
+            self.learner = LearnerThread(self.local_evaluator)
         self.learner.start()
 
         assert len(self.remote_evaluators) > 0
@@ -85,6 +226,7 @@ class AsyncSamplesOptimizer(PolicyOptimizer):
             ["put_weights", "enqueue", "sample_processing", "train", "sample"]
         }
         self.num_weight_syncs = 0
+        self.num_replayed = 0
         self.learning_started = False
 
         # Kick off async background sampling
@@ -92,11 +234,19 @@ class AsyncSamplesOptimizer(PolicyOptimizer):
         weights = self.local_evaluator.get_weights()
         for ev in self.remote_evaluators:
             ev.set_weights.remote(weights)
-            for _ in range(SAMPLE_QUEUE_DEPTH):
+            for _ in range(max_sample_requests_in_flight_per_worker):
                 self.sample_tasks.add(ev, ev.sample.remote())
 
         self.batch_buffer = []
 
+        if replay_proportion:
+            assert replay_buffer_num_slots > 0
+            assert (replay_buffer_num_slots * sample_batch_size >
+                    train_batch_size)
+        self.replay_proportion = replay_proportion
+        self.replay_buffer_num_slots = replay_buffer_num_slots
+        self.replay_batches = []
+
     def step(self):
         assert self.learner.is_alive()
         start = time.time()
@@ -112,23 +262,52 @@ class AsyncSamplesOptimizer(PolicyOptimizer):
         self.num_steps_sampled += sample_timesteps
         self.num_steps_trained += train_timesteps
 
+    def _augment_with_replay(self, sample_futures):
+        def can_replay():
+            num_needed = int(
+                np.ceil(self.train_batch_size / self.sample_batch_size))
+            return len(self.replay_batches) > num_needed
+
+        for ev, sample_batch in sample_futures:
+            sample_batch = ray.get(sample_batch)
+            yield ev, sample_batch
+
+            if can_replay():
+                f = self.replay_proportion
+                while random.random() < f:
+                    f -= 1
+                    replay_batch = random.choice(self.replay_batches)
+                    self.num_replayed += replay_batch.count
+                    yield None, replay_batch
+
     def _step(self):
         sample_timesteps, train_timesteps = 0, 0
         weights = None
 
         with self.timers["sample_processing"]:
-            for ev, sample_batch in self.sample_tasks.completed_prefetch():
-                sample_batch = ray.get(sample_batch)
-                sample_timesteps += sample_batch.count
+            for ev, sample_batch in self._augment_with_replay(
+                    self.sample_tasks.completed_prefetch()):
                 self.batch_buffer.append(sample_batch)
                 if sum(b.count
                        for b in self.batch_buffer) >= self.train_batch_size:
                     train_batch = self.batch_buffer[0].concat_samples(
                         self.batch_buffer)
                     with self.timers["enqueue"]:
-                        self.learner.inqueue.put((ev, train_batch))
+                        self.learner.inqueue.put(train_batch)
                     self.batch_buffer = []
 
+                # If the batch was replayed, skip the update below.
+                if ev is None:
+                    continue
+
+                sample_timesteps += sample_batch.count
+
+                # Put in replay buffer if enabled
+                if self.replay_buffer_num_slots > 0:
+                    self.replay_batches.append(sample_batch)
+                    if len(self.replay_batches) > self.replay_buffer_num_slots:
+                        self.replay_batches.pop(0)
+
                 # Note that it's important to pull new weights once
                 # updated to avoid excessive correlation between actors
                 if weights is None or self.learner.weights_updated:
@@ -154,6 +333,10 @@ class AsyncSamplesOptimizer(PolicyOptimizer):
         }
         timing["learner_grad_time_ms"] = round(
             1000 * self.learner.grad_timer.mean, 3)
+        timing["learner_load_time_ms"] = round(
+            1000 * self.learner.load_timer.mean, 3)
+        timing["learner_load_wait_time_ms"] = round(
+            1000 * self.learner.load_wait_timer.mean, 3)
         timing["learner_dequeue_time_ms"] = round(
             1000 * self.learner.queue_timer.mean, 3)
         stats = {
@@ -161,14 +344,10 @@ class AsyncSamplesOptimizer(PolicyOptimizer):
                                        3),
             "train_throughput": round(self.timers["train"].mean_throughput, 3),
             "num_weight_syncs": self.num_weight_syncs,
-        }
-        debug_stats = {
+            "num_steps_replayed": self.num_replayed,
             "timing_breakdown": timing,
-            "pending_sample_tasks": self.sample_tasks.count,
             "learner_queue": self.learner.learner_queue_size.stats(),
         }
-        if self.debug:
-            stats.update(debug_stats)
         if self.learner.stats:
             stats["learner"] = self.learner.stats
         return dict(PolicyOptimizer.stats(self), **stats)

python/ray/rllib/optimizers/multi_gpu_impl.py

@@ -36,13 +36,13 @@ class LocalSyncParallelOptimizer(object):
             to define the per-device loss ops.
         rnn_inputs: Extra input placeholders for RNN inputs. These will have
             shape [BATCH_SIZE // MAX_SEQ_LEN, ...].
-        per_device_batch_size: Number of tuples to optimize over at a time per
-            device. In each call to `optimize()`,
+        max_per_device_batch_size: Number of tuples to optimize over at a time
+            per device. In each call to `optimize()`,
             `len(devices) * per_device_batch_size` tuples of data will be
-            processed.
+            processed. If this is larger than the total data size, it will be
+            clipped.
         build_graph: Function that takes the specified inputs and returns a
             TF Policy Graph instance.
-        logdir: Directory to place debugging output in.
         grad_norm_clipping: None or int stdev to clip grad norms by
     """
 
@@ -51,18 +51,14 @@ class LocalSyncParallelOptimizer(object):
                  devices,
                  input_placeholders,
                  rnn_inputs,
-                 per_device_batch_size,
+                 max_per_device_batch_size,
                  build_graph,
-                 logdir,
                  grad_norm_clipping=None):
-        # TODO(rliaw): remove logdir
         self.optimizer = optimizer
         self.devices = devices
-        self.batch_size = per_device_batch_size * len(devices)
-        self.per_device_batch_size = per_device_batch_size
+        self.max_per_device_batch_size = max_per_device_batch_size
         self.loss_inputs = input_placeholders + rnn_inputs
         self.build_graph = build_graph
-        self.logdir = logdir
 
         # First initialize the shared loss network
         with tf.name_scope(TOWER_SCOPE_NAME):
@@ -71,6 +67,11 @@ class LocalSyncParallelOptimizer(object):
         # Then setup the per-device loss graphs that use the shared weights
         self._batch_index = tf.placeholder(tf.int32, name="batch_index")
 
+        # Dynamic batch size, which may be shrunk if there isn't enough data
+        self._per_device_batch_size = tf.placeholder(
+            tf.int32, name="per_device_batch_size")
+        self._loaded_per_device_batch_size = max_per_device_batch_size
+
         # When loading RNN input, we dynamically determine the max seq len
         self._max_seq_len = tf.placeholder(tf.int32, name="max_seq_len")
         self._loaded_max_seq_len = 1
@@ -88,9 +89,12 @@ class LocalSyncParallelOptimizer(object):
 
         avg = average_gradients([t.grads for t in self._towers])
         if grad_norm_clipping:
+            clipped = []
+            for grad, _ in avg:
+                clipped.append(grad)
+            clipped, _ = tf.clip_by_global_norm(clipped, grad_norm_clipping)
             for i, (grad, var) in enumerate(avg):
-                if grad is not None:
-                    avg[i] = (tf.clip_by_norm(grad, grad_norm_clipping), var)
+                avg[i] = (clipped[i], var)
         self._train_op = self.optimizer.apply_gradients(avg)
 
     def load_data(self, sess, inputs, state_inputs):
@@ -117,44 +121,64 @@ class LocalSyncParallelOptimizer(object):
         assert len(self.loss_inputs) == len(inputs + state_inputs), \
             (self.loss_inputs, inputs, state_inputs)
 
-        # The RNN truncation case is more complicated
+        # Let's suppose we have the following input data, and 2 devices:
+        # 1 2 3 4 5 6 7                              <- state inputs shape
+        # A A A B B B C C C D D D E E E F F F G G G  <- inputs shape
+        # The data is truncated and split across devices as follows:
+        # |---| seq len = 3
+        # |---------------------------------| seq batch size = 6 seqs
+        # |----------------| per device batch size = 9 tuples
+
         if len(state_inputs) > 0:
+            smallest_array = state_inputs[0]
             seq_len = len(inputs[0]) // len(state_inputs[0])
             self._loaded_max_seq_len = seq_len
-            assert len(state_inputs[0]) * seq_len == len(inputs[0])
-            # Make sure the shorter state inputs arrays are evenly divisible
+        else:
+            smallest_array = inputs[0]
+            self._loaded_max_seq_len = 1
+
+        seq_batch_size = (self.max_per_device_batch_size //
+                          self._loaded_max_seq_len * len(self.devices))
+        if len(smallest_array) < seq_batch_size:
+            # Dynamically shrink the batch size if insufficient data
+            seq_batch_size = make_divisible_by(
+                len(smallest_array), len(self.devices))
+        if seq_batch_size < len(self.devices):
+            raise ValueError("Must load at least 1 tuple sequence per device, "
+                             "got only {} total.".format(len(smallest_array)))
+        self._loaded_per_device_batch_size = (
+            seq_batch_size // len(self.devices) * self._loaded_max_seq_len)
+
+        if len(state_inputs) > 0:
+            # First truncate the RNN state arrays to the seq_batch_size
             state_inputs = [
-                make_divisible_by(arr, self.batch_size) for arr in state_inputs
+                make_divisible_by(arr, seq_batch_size) for arr in state_inputs
             ]
             # Then truncate the data inputs to match
             inputs = [arr[:len(state_inputs[0]) * seq_len] for arr in inputs]
-            assert len(state_inputs[0]) * seq_len == len(inputs[0])
-            assert len(state_inputs[0]) % self.batch_size == 0
+            assert len(state_inputs[0]) * seq_len == len(inputs[0]), \
+                (len(state_inputs[0]), seq_batch_size, seq_len, len(inputs[0]))
             for ph, arr in zip(self.loss_inputs, inputs + state_inputs):
                 feed_dict[ph] = arr
             truncated_len = len(inputs[0])
         else:
             for ph, arr in zip(self.loss_inputs, inputs + state_inputs):
-                truncated_arr = make_divisible_by(arr, self.batch_size)
+                truncated_arr = make_divisible_by(arr, seq_batch_size)
                 feed_dict[ph] = truncated_arr
                 truncated_len = len(truncated_arr)
 
         sess.run([t.init_op for t in self._towers], feed_dict=feed_dict)
 
         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
+        assert tuples_per_device > 0, "No data loaded?"
+        assert tuples_per_device % self._loaded_per_device_batch_size == 0
         return tuples_per_device
 
     def optimize(self, sess, batch_index):
         """Run a single step of SGD.
 
         Runs a SGD step over a slice of the preloaded batch with size given by
-        self.per_device_batch_size and offset given by the batch_index
+        self._loaded_per_device_batch_size and offset given by the batch_index
         argument.
 
         Updates shared model weights based on the averaged per-device
@@ -164,13 +188,14 @@ class LocalSyncParallelOptimizer(object):
             sess: TensorFlow session.
             batch_index: Offset into the preloaded data. This value must be
                 between `0` and `tuples_per_device`. The amount of data to
-                process is always fixed to `per_device_batch_size`.
+                process is at most `max_per_device_batch_size`.
 
         Returns:
             The outputs of extra_ops evaluated over the batch.
         """
         feed_dict = {
             self._batch_index: batch_index,
+            self._per_device_batch_size: self._loaded_per_device_batch_size,
             self._max_seq_len: self._loaded_max_seq_len,
         }
         for tower in self._towers:
@@ -213,7 +238,7 @@ class LocalSyncParallelOptimizer(object):
                         current_batch,
                         ([self._batch_index // scale * granularity] +
                          [0] * len(ph.shape[1:])),
-                        ([self.per_device_batch_size // scale * granularity] +
+                        ([self._per_device_batch_size // scale * granularity] +
                          [-1] * len(ph.shape[1:])))
                     current_slice.set_shape(ph.shape)
                     device_input_slices.append(current_slice)
@@ -229,8 +254,10 @@ class LocalSyncParallelOptimizer(object):
 Tower = namedtuple("Tower", ["init_op", "grads", "loss_graph"])
 
 
-def make_divisible_by(array, n):
-    return array[0:array.shape[0] - array.shape[0] % n]
+def make_divisible_by(a, n):
+    if type(a) is int:
+        return a - a % n
+    return a[0:a.shape[0] - a.shape[0] % n]
 
 
 def average_gradients(tower_grads):

python/ray/rllib/optimizers/multi_gpu_optimizer.py

@@ -4,7 +4,6 @@ from __future__ import print_function
 
 import numpy as np
 from collections import defaultdict
-import os
 import tensorflow as tf
 
 import ray
@@ -81,8 +80,7 @@ class LocalMultiGPUOptimizer(PolicyOptimizer):
                     self.par_opt = LocalSyncParallelOptimizer(
                         self.policy.optimizer(), self.devices,
                         [v for _, v in self.policy.loss_inputs()], rnn_inputs,
-                        self.per_device_batch_size, self.policy.copy,
-                        os.getcwd())
+                        self.per_device_batch_size, self.policy.copy)
 
                 self.sess = self.local_evaluator.tf_sess
                 self.sess.run(tf.global_variables_initializer())

python/ray/rllib/optimizers/replay_buffer.py

@@ -93,14 +93,15 @@ class ReplayBuffer(object):
         self._num_sampled += batch_size
         return self._encode_sample(idxes)
 
-    def stats(self):
+    def stats(self, debug=False):
         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())
+        if debug:
+            data.update(self._evicted_hit_stats.stats())
         return data
 
 
@@ -233,7 +234,8 @@ class PrioritizedReplayBuffer(ReplayBuffer):
 
             self._max_priority = max(self._max_priority, priority)
 
-    def stats(self):
-        parent = ReplayBuffer.stats(self)
-        parent.update(self._prio_change_stats.stats())
+    def stats(self, debug=False):
+        parent = ReplayBuffer.stats(self, debug)
+        if debug:
+            parent.update(self._prio_change_stats.stats())
         return parent

python/ray/rllib/test/test_supported_spaces.py

@@ -94,7 +94,7 @@ class ModelSupportedSpaces(unittest.TestCase):
     def testAll(self):
         ray.init()
         stats = {}
-        check_support("IMPALA", {"gpu": False}, stats)
+        check_support("IMPALA", {"num_gpus": 0}, stats)
         check_support("DDPG", {"timesteps_per_iteration": 1}, stats)
         check_support("DQN", {"timesteps_per_iteration": 1}, stats)
         check_support("A3C", {