diff --git a/doc/source/ddppo-arch.svg b/doc/source/ddppo-arch.svg new file mode 100644 index 000000000..e46e9a556 --- /dev/null +++ b/doc/source/ddppo-arch.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/doc/source/rllib-algorithms.rst b/doc/source/rllib-algorithms.rst index f36f0cbc7..0aa1005ea 100644 --- a/doc/source/rllib-algorithms.rst +++ b/doc/source/rllib-algorithms.rst @@ -104,7 +104,9 @@ Asynchronous Proximal Policy Optimization (APPO) `[implementation] `__ We include an asynchronous variant of Proximal Policy Optimization (PPO) based on the IMPALA architecture. This is similar to IMPALA but using a surrogate policy loss with clipping. Compared to synchronous PPO, APPO is more efficient in wall-clock time due to its use of asynchronous sampling. Using a clipped loss also allows for multiple SGD passes, and therefore the potential for better sample efficiency compared to IMPALA. V-trace can also be enabled to correct for off-policy samples. -APPO is not always more efficient; it is often better to simply use `PPO `__ or `IMPALA `__. +.. tip:: + + APPO is not always more efficient; it is often better to use `standard PPO `__ or `IMPALA `__. .. figure:: impala-arch.svg @@ -119,6 +121,30 @@ Tuned examples: `PongNoFrameskip-v4 `__ +`[implementation] `__ +Unlike APPO or PPO, with DD-PPO policy improvement is no longer done centralized in the trainer process. Instead, gradients are computed remotely on each rollout worker and all-reduced at each mini-batch using `torch distributed `__. This allows each worker's GPU to be used both for sampling and for training. + +.. tip:: + + DD-PPO is best for envs that require GPUs to function, or if you need to scale out SGD to multiple nodes. If you don't meet these requirements, `standard PPO <#proximal-policy-optimization-ppo>`__ will be more efficient. + +.. figure:: ddppo-arch.svg + + DD-PPO architecture (both sampling and learning are done on worker GPUs) + +Tuned examples: `CartPole-v0 `__, `BreakoutNoFrameskip-v4 `__ + +**DDPPO-specific configs** (see also `common configs `__): + +.. literalinclude:: ../../rllib/agents/ppo/ddppo.py + :language: python + :start-after: __sphinx_doc_begin__ + :end-before: __sphinx_doc_end__ + Gradient-based ~~~~~~~~~~~~~~ @@ -241,7 +267,11 @@ Proximal Policy Optimization (PPO) ---------------------------------- |pytorch| |tensorflow| `[paper] `__ `[implementation] `__ -PPO's clipped objective supports multiple SGD passes over the same batch of experiences. RLlib's multi-GPU optimizer pins that data in GPU memory to avoid unnecessary transfers from host memory, substantially improving performance over a naive implementation. RLlib's PPO scales out using multiple workers for experience collection, and also with multiple GPUs for SGD. +PPO's clipped objective supports multiple SGD passes over the same batch of experiences. RLlib's multi-GPU optimizer pins that data in GPU memory to avoid unnecessary transfers from host memory, substantially improving performance over a naive implementation. PPO scales out using multiple workers for experience collection, and also to multiple GPUs for SGD. + +.. tip:: + + If you need to scale out with GPUs on multiple nodes, consider using `decentralized PPO <#decentralized-distributed-proximal-policy-optimization-dd-ppo>`__. .. figure:: ppo-arch.svg diff --git a/doc/source/rllib-toc.rst b/doc/source/rllib-toc.rst index 43fc2ba74..1b6335027 100644 --- a/doc/source/rllib-toc.rst +++ b/doc/source/rllib-toc.rst @@ -86,6 +86,8 @@ Algorithms - |tensorflow| `Asynchronous Proximal Policy Optimization (APPO) `__ + - |pytorch| `Decentralized Distributed Proximal Policy Optimization (DD-PPO) `__ + - |pytorch| `Single-Player AlphaZero (contrib/AlphaZero) `__ * Gradient-based diff --git a/doc/source/rllib.rst b/doc/source/rllib.rst index fa9c0cec3..15d9a2df4 100644 --- a/doc/source/rllib.rst +++ b/doc/source/rllib.rst @@ -10,7 +10,7 @@ To get started, take a look over the `custom env example `__ and `RLlib blog posts `__. You may also want to skim the `list of built-in algorithms `__. Look out for the |tensorflow| and |pytorch| icons to see which algorithms are available for each framework. +The following is a whirlwind overview of RLlib. For a more in-depth guide, see also the `full table of contents `__ and `RLlib blog posts `__. You may also want to skim the `list of built-in algorithms `__. Look out for the |tensorflow| and |pytorch| icons to see which algorithms are `available `__ for each framework. Running RLlib ~~~~~~~~~~~~~ diff --git a/rllib/agents/ppo/__init__.py b/rllib/agents/ppo/__init__.py index 239008743..741ed15c7 100644 --- a/rllib/agents/ppo/__init__.py +++ b/rllib/agents/ppo/__init__.py @@ -1,4 +1,5 @@ from ray.rllib.agents.ppo.ppo import PPOTrainer, DEFAULT_CONFIG from ray.rllib.agents.ppo.appo import APPOTrainer +from ray.rllib.agents.ppo.ddppo import DDPPOTrainer -__all__ = ["APPOTrainer", "PPOTrainer", "DEFAULT_CONFIG"] +__all__ = ["APPOTrainer", "DDPPOTrainer", "PPOTrainer", "DEFAULT_CONFIG"] diff --git a/rllib/agents/ppo/ddppo.py b/rllib/agents/ppo/ddppo.py new file mode 100644 index 000000000..cc50d13d8 --- /dev/null +++ b/rllib/agents/ppo/ddppo.py @@ -0,0 +1,87 @@ +from ray.rllib.agents.ppo import ppo +from ray.rllib.agents.trainer import with_base_config +from ray.rllib.optimizers import TorchDistributedDataParallelOptimizer +"""Decentralized Distributed PPO implementation. + +Unlike APPO or PPO, learning is no longer done centralized in the trainer +process. Instead, gradients are computed remotely on each rollout worker and +all-reduced to sync them at each mini-batch. This allows each worker's GPU +to be used both for sampling and for training. + +DD-PPO should be used if you have envs that require GPUs to function, or have +a very large model that cannot be effectively optimized with the GPUs available +on a single machine (DD-PPO allows scaling to arbitrary numbers of GPUs across +multiple nodes, unlike PPO/APPO which is limited to GPUs on a single node). + +Paper reference: https://arxiv.org/abs/1911.00357 +Note that unlike the paper, we currently do not implement straggler mitigation. +""" + +# yapf: disable +# __sphinx_doc_begin__ +DEFAULT_CONFIG = with_base_config(ppo.DEFAULT_CONFIG, { + # During the sampling phase, each rollout worker will collect a batch + # `sample_batch_size * num_envs_per_worker` steps in size. + "sample_batch_size": 100, + # Vectorize the env (should enable by default since each worker has a GPU). + "num_envs_per_worker": 5, + # During the SGD phase, workers iterate over minibatches of this size. + # The effective minibatch size will be `sgd_minibatch_size * num_workers`. + "sgd_minibatch_size": 50, + # Number of SGD epochs per optimization round. + "num_sgd_iter": 10, + + # *** WARNING: configs below are DDPPO overrides over PPO; you + # shouldn't need to adjust them. *** + "use_pytorch": True, # DDPPO requires PyTorch distributed. + "num_gpus": 0, # Learning is no longer done on the driver process, so + # giving GPUs to the driver does not make sense! + "num_gpus_per_worker": 1, # Each rollout worker gets a GPU. + "truncate_episodes": True, # Require evenly sized batches. Otherwise, + # collective allreduce could fail. + "train_batch_size": -1, # This is auto set based on sample batch size. +}) +# __sphinx_doc_end__ +# yapf: enable + + +def validate_config(config): + if config["train_batch_size"] == -1: + # Auto set. + config["train_batch_size"] = ( + config["sample_batch_size"] * config["num_envs_per_worker"]) + else: + raise ValueError( + "Set sample_batch_size instead of train_batch_size for DDPPO.") + ppo.validate_config(config) + + +def make_distributed_allreduce_optimizer(workers, config): + if not config["use_pytorch"]: + raise ValueError( + "Distributed data parallel is only supported for PyTorch") + if config["num_gpus"]: + raise ValueError( + "When using distributed data parallel, you should set " + "num_gpus=0 since all optimization " + "is happening on workers. Enable GPUs for workers by setting " + "num_gpus_per_worker=1.") + if config["batch_mode"] != "truncate_episodes": + raise ValueError( + "Distributed data parallel requires truncate_episodes " + "batch mode.") + + return TorchDistributedDataParallelOptimizer( + workers, + expected_batch_size=config["sample_batch_size"] * + config["num_envs_per_worker"], + num_sgd_iter=config["num_sgd_iter"], + sgd_minibatch_size=config["sgd_minibatch_size"], + standardize_fields=["advantages"]) + + +DDPPOTrainer = ppo.PPOTrainer.with_updates( + name="DDPPO", + default_config=DEFAULT_CONFIG, + make_policy_optimizer=make_distributed_allreduce_optimizer, + validate_config=validate_config) diff --git a/rllib/agents/ppo/ppo.py b/rllib/agents/ppo/ppo.py index 8c1919c22..25f990656 100644 --- a/rllib/agents/ppo/ppo.py +++ b/rllib/agents/ppo/ppo.py @@ -3,8 +3,7 @@ import logging from ray.rllib.agents import with_common_config from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicy from ray.rllib.agents.trainer_template import build_trainer -from ray.rllib.optimizers import SyncSamplesOptimizer, \ - LocalMultiGPUOptimizer, TorchDistributedDataParallelOptimizer +from ray.rllib.optimizers import SyncSamplesOptimizer, LocalMultiGPUOptimizer from ray.rllib.utils import try_import_tf tf = try_import_tf() @@ -69,8 +68,6 @@ DEFAULT_CONFIG = with_common_config({ # usually slower, but you might want to try it if you run into issues with # the default optimizer. "simple_optimizer": False, - # Use the experimental torch multi-node SGD optimizer. - "distributed_data_parallel_optimizer": False, # Use PyTorch as framework? "use_pytorch": False }) @@ -79,33 +76,6 @@ DEFAULT_CONFIG = with_common_config({ def choose_policy_optimizer(workers, config): - if config["distributed_data_parallel_optimizer"]: - if not config["use_pytorch"]: - raise ValueError( - "Distributed data parallel is only supported for PyTorch") - if config["num_gpus"]: - raise ValueError( - "When using distributed data parallel, you should set " - "num_gpus=0 since all optimization " - "is happening on workers. Enable GPUs for workers by setting " - "num_gpus_per_worker=1.") - if config["batch_mode"] != "truncate_episodes": - raise ValueError( - "Distributed data parallel requires truncate_episodes " - "batch mode.") - if config["sample_batch_size"] != config["train_batch_size"]: - raise ValueError( - "Distributed data parallel requires sample_batch_size to be " - "equal to train_batch_size. Each worker will sample and learn " - "on train_batch_size samples per iteration.") - - return TorchDistributedDataParallelOptimizer( - workers, - num_sgd_iter=config["num_sgd_iter"], - train_batch_size=config["train_batch_size"], - sgd_minibatch_size=config["sgd_minibatch_size"], - standardize_fields=["advantages"]) - if config["simple_optimizer"]: return SyncSamplesOptimizer( workers, diff --git a/rllib/agents/ppo/ppo_torch_policy.py b/rllib/agents/ppo/ppo_torch_policy.py index 92bd3c8fb..92bfb3e6a 100644 --- a/rllib/agents/ppo/ppo_torch_policy.py +++ b/rllib/agents/ppo/ppo_torch_policy.py @@ -161,9 +161,10 @@ def vf_preds_and_logits_fetches(policy, input_dict, state_batches, model, action_dist): """Adds value function and logits outputs to experience train_batches.""" return { - SampleBatch.VF_PREDS: policy.model.value_function(), - BEHAVIOUR_LOGITS: policy.model.last_output().numpy(), - ACTION_LOGP: action_dist.logp(input_dict[SampleBatch.ACTIONS]) + SampleBatch.VF_PREDS: policy.model.value_function().cpu().numpy(), + BEHAVIOUR_LOGITS: policy.model.last_output().cpu().numpy(), + ACTION_LOGP: action_dist.logp( + input_dict[SampleBatch.ACTIONS]).cpu().numpy(), } @@ -187,11 +188,14 @@ class ValueNetworkMixin: def value(ob, prev_action, prev_reward, *state): model_out, _ = self.model({ - SampleBatch.CUR_OBS: torch.Tensor([ob]), - SampleBatch.PREV_ACTIONS: torch.Tensor([prev_action]), - SampleBatch.PREV_REWARDS: torch.Tensor([prev_reward]), + SampleBatch.CUR_OBS: torch.Tensor([ob]).to(self.device), + SampleBatch.PREV_ACTIONS: torch.Tensor([prev_action]).to( + self.device), + SampleBatch.PREV_REWARDS: torch.Tensor([prev_reward]).to( + self.device), "is_training": False, - }, [torch.Tensor([s]) for s in state], torch.Tensor([1])) + }, [torch.Tensor([s]).to(self.device) for s in state], + torch.Tensor([1]).to(self.device)) return self.model.value_function()[0] else: diff --git a/rllib/agents/registry.py b/rllib/agents/registry.py index 35bd160ba..be6e0920a 100644 --- a/rllib/agents/registry.py +++ b/rllib/agents/registry.py @@ -15,6 +15,11 @@ def _import_appo(): return ppo.APPOTrainer +def _import_ddppo(): + from ray.rllib.agents import ppo + return ppo.DDPPOTrainer + + def _import_qmix(): from ray.rllib.agents import qmix return qmix.QMixTrainer @@ -113,6 +118,7 @@ ALGORITHMS = { "QMIX": _import_qmix, "APEX_QMIX": _import_apex_qmix, "APPO": _import_appo, + "DDPPO": _import_ddppo, "MARWIL": _import_marwil, } diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index 6ec4bf6c6..1bc7c604e 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/rllib/evaluation/rollout_worker.py @@ -625,14 +625,14 @@ class RolloutWorker(EvaluatorInterface): logger.debug("Training out:\n\n{}\n".format(summarize(info_out))) return info_out - def sample_and_learn(self, train_batch_size, num_sgd_iter, + def sample_and_learn(self, expected_batch_size, num_sgd_iter, sgd_minibatch_size, standardize_fields): """Sample and batch and learn on it. This is typically used in combination with distributed allreduce. Arguments: - train_batch_size (int): Number of samples to learn on. + expected_batch_size (int): Expected number of samples to learn on. num_sgd_iter (int): Number of SGD iterations. sgd_minibatch_size (int): SGD minibatch size. standardize_fields (list): List of sample fields to normalize. @@ -642,10 +642,12 @@ class RolloutWorker(EvaluatorInterface): count: number of samples learned on. """ batch = self.sample() - assert batch.count == train_batch_size, \ - (batch.count, "Batch size possibly out of sync between workers") + assert batch.count == expected_batch_size, \ + ("Batch size possibly out of sync between workers, expected:", + expected_batch_size, "got:", batch.count) logger.info("Executing distributed minibatch SGD " - "on batch of size {}".format(batch.count)) + "with epoch size {}, minibatch size {}".format( + batch.count, sgd_minibatch_size)) info = do_minibatch_sgd(batch, self.policy_map, self, num_sgd_iter, sgd_minibatch_size, standardize_fields) return info, batch.count diff --git a/rllib/optimizers/torch_distributed_data_parallel_optimizer.py b/rllib/optimizers/torch_distributed_data_parallel_optimizer.py index 73402b70c..33b25d2b4 100644 --- a/rllib/optimizers/torch_distributed_data_parallel_optimizer.py +++ b/rllib/optimizers/torch_distributed_data_parallel_optimizer.py @@ -13,8 +13,8 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer): def __init__(self, workers, + expected_batch_size, num_sgd_iter=1, - train_batch_size=1, sgd_minibatch_size=0, standardize_fields=frozenset([]), keep_local_weights_in_sync=True, @@ -22,11 +22,12 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer): PolicyOptimizer.__init__(self, workers) self.learner_stats = {} self.num_sgd_iter = num_sgd_iter - self.train_batch_size = train_batch_size + self.expected_batch_size = expected_batch_size self.sgd_minibatch_size = sgd_minibatch_size self.standardize_fields = standardize_fields self.keep_local_weights_in_sync = keep_local_weights_in_sync - self.update_weights_timer = TimerStat() + self.sync_down_timer = TimerStat() + self.sync_up_timer = TimerStat() self.learn_timer = TimerStat() # Setup the distributed processes. @@ -52,7 +53,7 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer): # add too much overhead and handles the case where the user manually # updates the local weights. if self.keep_local_weights_in_sync: - with self.update_weights_timer: + with self.sync_up_timer: weights = ray.put(self.workers.local_worker().get_weights()) for e in self.workers.remote_workers(): e.set_weights.remote(weights) @@ -60,7 +61,7 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer): with self.learn_timer: results = ray.get([ w.sample_and_learn.remote( - self.train_batch_size, self.num_sgd_iter, + self.expected_batch_size, self.num_sgd_iter, self.sgd_minibatch_size, self.standardize_fields) for w in self.workers.remote_workers() ]) @@ -87,8 +88,10 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer): # Sync down the weights. As with the sync up, this is not really # needed unless the user is reading the local weights. if self.keep_local_weights_in_sync: - self.workers.local_worker().set_weights( - ray.get(self.workers.remote_workers()[0].get_weights.remote())) + with self.sync_down_timer: + self.workers.local_worker().set_weights( + ray.get( + self.workers.remote_workers()[0].get_weights.remote())) return self.learner_stats @@ -96,8 +99,10 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer): def stats(self): return dict( PolicyOptimizer.stats(self), **{ - "update_weights_time_ms": round( - 1000 * self.update_weights_timer.mean, 3), + "sync_weights_up_time": round(1000 * self.sync_up_timer.mean, + 3), + "sync_weights_down_time": round( + 1000 * self.sync_down_timer.mean, 3), "learn_time_ms": round(1000 * self.learn_timer.mean, 3), "learner": self.learner_stats, }) diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index a68bc133d..4d63af818 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -108,8 +108,14 @@ class TorchPolicy(Policy): if p.grad is not None: grads.append(p.grad) start = time.time() - torch.distributed.all_reduce_coalesced( - grads, op=torch.distributed.ReduceOp.SUM) + if torch.cuda.is_available(): + # Sadly, allreduce_coalesced does not work with CUDA yet. + for g in grads: + torch.distributed.all_reduce( + g, op=torch.distributed.ReduceOp.SUM) + else: + torch.distributed.all_reduce_coalesced( + grads, op=torch.distributed.ReduceOp.SUM) for p in self.model.parameters(): if p.grad is not None: p.grad /= self.distributed_world_size @@ -208,6 +214,8 @@ class TorchPolicy(Policy): train_batch = UsageTrackingDict(postprocessed_batch) def convert(arr): + if torch.is_tensor(arr): + return arr.to(self.device) tensor = torch.from_numpy(np.asarray(arr)) if tensor.dtype == torch.double: tensor = tensor.float() diff --git a/rllib/tuned_examples/atari-ddppo.yaml b/rllib/tuned_examples/atari-ddppo.yaml new file mode 100644 index 000000000..221a1f5be --- /dev/null +++ b/rllib/tuned_examples/atari-ddppo.yaml @@ -0,0 +1,30 @@ +# Basically the same as atari-ppo, but adapted for DDPPO. Note that DDPPO +# isn't actually any more efficient on Atari, since the network size is +# relatively small and the env doesn't require a GPU. +atari-ddppo: + env: + grid_search: + - BreakoutNoFrameskip-v4 + run: DDPPO + config: + # Worker config: 10 workers, each of which requires a GPU. + num_workers: 10 + num_gpus_per_worker: 1 + # Each worker will sample 100 * 5 envs per worker steps = 500 steps + # per optimization round. This is 5000 steps summed across workers. + sample_batch_size: 100 + num_envs_per_worker: 5 + # Each worker will take a minibatch of 50. There are 10 workers total, + # so the effective minibatch size will be 500. + sgd_minibatch_size: 50 + num_sgd_iter: 10 + # Params from standard PPO Atari config: + lambda: 0.95 + kl_coeff: 0.5 + clip_rewards: True + clip_param: 0.1 + vf_clip_param: 10.0 + entropy_coeff: 0.01 + batch_mode: truncate_episodes + observation_filter: NoFilter + vf_share_layers: true diff --git a/rllib/tuned_examples/regression_tests/cartpole-ddppo.yaml b/rllib/tuned_examples/regression_tests/cartpole-ddppo.yaml new file mode 100644 index 000000000..0d9516cbb --- /dev/null +++ b/rllib/tuned_examples/regression_tests/cartpole-ddppo.yaml @@ -0,0 +1,8 @@ +cartpole-ddppo: + env: CartPole-v0 + run: DDPPO + stop: + episode_reward_mean: 100 + timesteps_total: 100000 + config: + num_gpus_per_worker: 0 diff --git a/rllib/tuned_examples/regression_tests/cartpole-torch-dist.yaml b/rllib/tuned_examples/regression_tests/cartpole-torch-dist.yaml deleted file mode 100644 index c346d9671..000000000 --- a/rllib/tuned_examples/regression_tests/cartpole-torch-dist.yaml +++ /dev/null @@ -1,14 +0,0 @@ -cartpole-torch-dist: - env: CartPole-v0 - run: PPO - stop: - episode_reward_mean: 150 - timesteps_total: 100000 - config: - num_workers: 2 - sample_batch_size: 4000 - train_batch_size: 4000 - batch_mode: truncate_episodes - observation_filter: MeanStdFilter - use_pytorch: true - distributed_data_parallel_optimizer: true