diff --git a/doc/source/rllib-algorithms.rst b/doc/source/rllib-algorithms.rst index d764fc7ad..be39b61bb 100644 --- a/doc/source/rllib-algorithms.rst +++ b/doc/source/rllib-algorithms.rst @@ -43,7 +43,7 @@ Importance Weighted Actor-Learner Architecture (IMPALA) `[paper] `__ `[implementation] `__ -In IMPALA, a central learner runs SGD in a tight loop while asynchronously pulling sample batches from many actor processes. RLlib's IMPALA implementation uses DeepMind's reference `V-trace code `__. Note that we do not provide a deep residual network out of the box, but one can be plugged in as a `custom model `__. +In IMPALA, a central learner runs SGD in a tight loop while asynchronously pulling sample batches from many actor processes. RLlib's IMPALA implementation uses DeepMind's reference `V-trace code `__. Note that we do not provide a deep residual network out of the box, but one can be plugged in as a `custom model `__. Multiple learner GPUs and experience replay are also supported. Tuned examples: `PongNoFrameskip-v4 `__, `vectorized configuration `__, `{BeamRider,Breakout,Qbert,SpaceInvaders}NoFrameskip-v4 `__ diff --git a/python/ray/rllib/agents/impala/impala.py b/python/ray/rllib/agents/impala/impala.py index cfa55bd73..1ad2b673f 100644 --- a/python/ray/rllib/agents/impala/impala.py +++ b/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"]) diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index f69846871..2ba018b91 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/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) diff --git a/python/ray/rllib/agents/ppo/ppo_policy_graph.py b/python/ray/rllib/agents/ppo/ppo_policy_graph.py index 9456ebe94..638ae0eb8 100644 --- a/python/ray/rllib/agents/ppo/ppo_policy_graph.py +++ b/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) diff --git a/python/ray/rllib/examples/cartpole_lstm.py b/python/ray/rllib/examples/cartpole_lstm.py index e3d0ddc4c..67fd35d28 100644 --- a/python/ray/rllib/examples/cartpole_lstm.py +++ b/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, }, - }, + }), } }) diff --git a/python/ray/rllib/optimizers/async_replay_optimizer.py b/python/ray/rllib/optimizers/async_replay_optimizer.py index 3ed5f37d3..c48fd1860 100644 --- a/python/ray/rllib/optimizers/async_replay_optimizer.py +++ b/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) diff --git a/python/ray/rllib/optimizers/async_samples_optimizer.py b/python/ray/rllib/optimizers/async_samples_optimizer.py index 3b6bb861b..69f5e849b 100644 --- a/python/ray/rllib/optimizers/async_samples_optimizer.py +++ b/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) diff --git a/python/ray/rllib/optimizers/multi_gpu_impl.py b/python/ray/rllib/optimizers/multi_gpu_impl.py index 7233e37e9..1affe8df3 100644 --- a/python/ray/rllib/optimizers/multi_gpu_impl.py +++ b/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): diff --git a/python/ray/rllib/optimizers/multi_gpu_optimizer.py b/python/ray/rllib/optimizers/multi_gpu_optimizer.py index e47457036..4595415a1 100644 --- a/python/ray/rllib/optimizers/multi_gpu_optimizer.py +++ b/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()) diff --git a/python/ray/rllib/optimizers/replay_buffer.py b/python/ray/rllib/optimizers/replay_buffer.py index 77d954345..cd5ec7328 100644 --- a/python/ray/rllib/optimizers/replay_buffer.py +++ b/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 diff --git a/python/ray/rllib/test/test_supported_spaces.py b/python/ray/rllib/test/test_supported_spaces.py index 2ced3402a..16bdd485f 100644 --- a/python/ray/rllib/test/test_supported_spaces.py +++ b/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", { diff --git a/test/jenkins_tests/run_multi_node_tests.sh b/test/jenkins_tests/run_multi_node_tests.sh index 43815f470..0a6db03b9 100755 --- a/test/jenkins_tests/run_multi_node_tests.sh +++ b/test/jenkins_tests/run_multi_node_tests.sh @@ -189,14 +189,28 @@ docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ --env CartPole-v0 \ --run IMPALA \ --stop '{"training_iteration": 2}' \ - --config '{"gpu": false, "num_workers": 2, "min_iter_time_s": 1}' + --config '{"num_gpus": 0, "num_workers": 2, "min_iter_time_s": 1}' docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/rllib/train.py \ --env CartPole-v0 \ --run IMPALA \ --stop '{"training_iteration": 2}' \ - --config '{"gpu": false, "num_workers": 2, "min_iter_time_s": 1, "model": {"use_lstm": true}}' + --config '{"num_gpus": 0, "num_workers": 2, "min_iter_time_s": 1, "model": {"use_lstm": true}}' + +docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ + python /ray/python/ray/rllib/train.py \ + --env CartPole-v0 \ + --run IMPALA \ + --stop '{"training_iteration": 2}' \ + --config '{"num_gpus": 0, "num_workers": 2, "min_iter_time_s": 1, "num_parallel_data_loaders": 2, "replay_proportion": 1.0}' + +docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ + python /ray/python/ray/rllib/train.py \ + --env CartPole-v0 \ + --run IMPALA \ + --stop '{"training_iteration": 2}' \ + --config '{"num_gpus": 0, "num_workers": 2, "min_iter_time_s": 1, "num_parallel_data_loaders": 2, "replay_proportion": 1.0, "model": {"use_lstm": true}}' docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/rllib/train.py \ @@ -295,7 +309,10 @@ docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/rllib/examples/multiagent_two_trainers.py --num-iters=2 docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ - python /ray/python/ray/rllib/examples/cartpole_lstm.py --stop=200 + python /ray/python/ray/rllib/examples/cartpole_lstm.py --run=PPO --stop=200 + +docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ + python /ray/python/ray/rllib/examples/cartpole_lstm.py --run=IMPALA --stop=100 docker run -e "RAY_USE_XRAY=1" --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/experimental/sgd/test_sgd.py --num-iters=2