From 0f88444686442fee367d814fa4af6c776b9a7148 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Mon, 2 Mar 2020 15:16:37 -0800 Subject: [PATCH] [rllib] Support multi-agent training in pipeline impls, add easy flag to enable (#7338) --- rllib/agents/a3c/__init__.py | 6 +---- rllib/agents/a3c/a2c.py | 36 +++++++++++++++++++++++++++- rllib/agents/a3c/a2c_pipeline.py | 38 ------------------------------ rllib/agents/a3c/a3c.py | 17 ++++++++++++- rllib/agents/a3c/tests/test_a2c.py | 23 +++++++++++------- rllib/agents/pg/__init__.py | 3 +-- rllib/agents/pg/pg.py | 22 ++++++++++++++++- rllib/agents/pg/pg_pipeline.py | 24 ------------------- rllib/agents/pg/tests/test_pg.py | 9 +++++-- rllib/agents/registry.py | 20 ---------------- rllib/agents/trainer.py | 3 +++ rllib/agents/trainer_template.py | 8 ++++++- rllib/tests/test_pipeline.py | 18 ++++++++++---- rllib/utils/experimental_dsl.py | 33 ++++++++++++++++---------- 14 files changed, 138 insertions(+), 122 deletions(-) delete mode 100644 rllib/agents/a3c/a2c_pipeline.py delete mode 100644 rllib/agents/pg/pg_pipeline.py diff --git a/rllib/agents/a3c/__init__.py b/rllib/agents/a3c/__init__.py index 9467b770c..96a1498df 100644 --- a/rllib/agents/a3c/__init__.py +++ b/rllib/agents/a3c/__init__.py @@ -1,8 +1,4 @@ from ray.rllib.agents.a3c.a3c import A3CTrainer, DEFAULT_CONFIG from ray.rllib.agents.a3c.a2c import A2CTrainer -from ray.rllib.agents.a3c.a2c_pipeline import A2CPipeline -from ray.rllib.agents.a3c.a3c_pipeline import A3CPipeline -__all__ = [ - "A2CTrainer", "A3CTrainer", "DEFAULT_CONFIG", "A2CPipeline", "A3CPipeline" -] +__all__ = ["A2CTrainer", "A3CTrainer", "DEFAULT_CONFIG"] diff --git a/rllib/agents/a3c/a2c.py b/rllib/agents/a3c/a2c.py index f8bb4c372..eb002d92c 100644 --- a/rllib/agents/a3c/a2c.py +++ b/rllib/agents/a3c/a2c.py @@ -1,9 +1,14 @@ +import math + from ray.rllib.agents.a3c.a3c import DEFAULT_CONFIG as A3C_CONFIG, \ validate_config, get_policy_class from ray.rllib.optimizers import SyncSamplesOptimizer, MicrobatchOptimizer from ray.rllib.agents.a3c.a3c_tf_policy import A3CTFPolicy from ray.rllib.agents.trainer_template import build_trainer from ray.rllib.utils import merge_dicts +from ray.rllib.utils.experimental_dsl import ( + ParallelRollouts, ConcatBatches, ComputeGradients, AverageGradients, + ApplyGradients, TrainOneStep, StandardMetricsReporting) A2C_DEFAULT_CONFIG = merge_dicts( A3C_CONFIG, @@ -32,10 +37,39 @@ def choose_policy_optimizer(workers, config): workers, train_batch_size=config["train_batch_size"]) +# Experimental pipeline-based impl; enable with "use_pipeline_impl": True. +def training_pipeline(workers, config): + rollouts = ParallelRollouts(workers, mode="bulk_sync") + + if config["microbatch_size"]: + num_microbatches = math.ceil( + config["train_batch_size"] / config["microbatch_size"]) + # In microbatch mode, we want to compute gradients on experience + # microbatches, average a number of these microbatches, and then apply + # the averaged gradient in one SGD step. This conserves GPU memory, + # allowing for extremely large experience batches to be used. + train_op = ( + rollouts.combine( + ConcatBatches(min_batch_size=config["microbatch_size"])) + .for_each(ComputeGradients(workers)) # (grads, info) + .batch(num_microbatches) # List[(grads, info)] + .for_each(AverageGradients()) # (avg_grads, info) + .for_each(ApplyGradients(workers))) + else: + # In normal mode, we execute one SGD step per each train batch. + train_op = rollouts \ + .combine(ConcatBatches( + min_batch_size=config["train_batch_size"])) \ + .for_each(TrainOneStep(workers)) + + return StandardMetricsReporting(train_op, workers, config) + + A2CTrainer = build_trainer( name="A2C", default_config=A2C_DEFAULT_CONFIG, default_policy=A3CTFPolicy, get_policy_class=get_policy_class, make_policy_optimizer=choose_policy_optimizer, - validate_config=validate_config) + validate_config=validate_config, + training_pipeline=training_pipeline) diff --git a/rllib/agents/a3c/a2c_pipeline.py b/rllib/agents/a3c/a2c_pipeline.py deleted file mode 100644 index aa12cfe7f..000000000 --- a/rllib/agents/a3c/a2c_pipeline.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Experimental pipeline-based impl; run this with --run='A2C_pl'""" - -import math - -from ray.rllib.agents.a3c.a2c import A2CTrainer -from ray.rllib.utils.experimental_dsl import ( - ParallelRollouts, ConcatBatches, ComputeGradients, AverageGradients, - ApplyGradients, TrainOneStep, StandardMetricsReporting) - - -def training_pipeline(workers, config): - rollouts = ParallelRollouts(workers, mode="bulk_sync") - - if config["microbatch_size"]: - num_microbatches = math.ceil( - config["train_batch_size"] / config["microbatch_size"]) - # In microbatch mode, we want to compute gradients on experience - # microbatches, average a number of these microbatches, and then apply - # the averaged gradient in one SGD step. This conserves GPU memory, - # allowing for extremely large experience batches to be used. - train_op = ( - rollouts.combine( - ConcatBatches(min_batch_size=config["microbatch_size"])) - .for_each(ComputeGradients(workers)) # (grads, info) - .batch(num_microbatches) # List[(grads, info)] - .for_each(AverageGradients()) # (avg_grads, info) - .for_each(ApplyGradients(workers))) - else: - # In normal mode, we execute one SGD step per each train batch. - train_op = rollouts \ - .combine(ConcatBatches( - min_batch_size=config["train_batch_size"])) \ - .for_each(TrainOneStep(workers)) - - return StandardMetricsReporting(train_op, workers, config) - - -A2CPipeline = A2CTrainer.with_updates(training_pipeline=training_pipeline) diff --git a/rllib/agents/a3c/a3c.py b/rllib/agents/a3c/a3c.py index 7a1980c82..36339f11e 100644 --- a/rllib/agents/a3c/a3c.py +++ b/rllib/agents/a3c/a3c.py @@ -4,6 +4,8 @@ from ray.rllib.agents.a3c.a3c_tf_policy import A3CTFPolicy from ray.rllib.agents.trainer import with_common_config from ray.rllib.agents.trainer_template import build_trainer from ray.rllib.optimizers import AsyncGradientsOptimizer +from ray.rllib.utils.experimental_dsl import (AsyncGradients, ApplyGradients, + StandardMetricsReporting) logger = logging.getLogger(__name__) @@ -63,10 +65,23 @@ def make_async_optimizer(workers, config): return AsyncGradientsOptimizer(workers, **config["optimizer"]) +# Experimental pipeline-based impl; enable with "use_pipeline_impl": True. +def training_pipeline(workers, config): + # For A3C, compute policy gradients remotely on the rollout workers. + grads = AsyncGradients(workers) + + # Apply the gradients as they arrive. We set update_all to False so that + # only the worker sending the gradient is updated with new weights. + train_op = grads.for_each(ApplyGradients(workers, update_all=False)) + + return StandardMetricsReporting(train_op, workers, config) + + A3CTrainer = build_trainer( name="A3C", default_config=DEFAULT_CONFIG, default_policy=A3CTFPolicy, get_policy_class=get_policy_class, validate_config=validate_config, - make_policy_optimizer=make_async_optimizer) + make_policy_optimizer=make_async_optimizer, + training_pipeline=training_pipeline) diff --git a/rllib/agents/a3c/tests/test_a2c.py b/rllib/agents/a3c/tests/test_a2c.py index 935444b49..19ff7bf89 100644 --- a/rllib/agents/a3c/tests/test_a2c.py +++ b/rllib/agents/a3c/tests/test_a2c.py @@ -1,11 +1,11 @@ import unittest import ray -from ray.rllib.agents.a3c import a2c_pipeline +from ray.rllib.agents.a3c import A2CTrainer class TestA2C(unittest.TestCase): - """Sanity tests for A2CPipeline.""" + """Sanity tests for A2C pipeline.""" def setUp(self): ray.init() @@ -14,16 +14,21 @@ class TestA2C(unittest.TestCase): ray.shutdown() def test_a2c_pipeline(ray_start_regular): - trainer = a2c_pipeline.A2CPipeline( - env="CartPole-v0", config={"min_iter_time_s": 0}) - assert isinstance(trainer.train(), dict) - - def test_a2c_pipeline_microbatch(ray_start_regular): - trainer = a2c_pipeline.A2CPipeline( + trainer = A2CTrainer( env="CartPole-v0", config={ "min_iter_time_s": 0, - "microbatch_size": 10 + "use_pipeline_impl": True + }) + assert isinstance(trainer.train(), dict) + + def test_a2c_pipeline_microbatch(ray_start_regular): + trainer = A2CTrainer( + env="CartPole-v0", + config={ + "min_iter_time_s": 0, + "microbatch_size": 10, + "use_pipeline_impl": True, }) assert isinstance(trainer.train(), dict) diff --git a/rllib/agents/pg/__init__.py b/rllib/agents/pg/__init__.py index 6aa398214..b8d289645 100644 --- a/rllib/agents/pg/__init__.py +++ b/rllib/agents/pg/__init__.py @@ -1,10 +1,9 @@ from ray.rllib.agents.pg.pg import PGTrainer, DEFAULT_CONFIG -from ray.rllib.agents.pg.pg_pipeline import PGPipeline from ray.rllib.agents.pg.pg_tf_policy import pg_tf_loss, \ post_process_advantages from ray.rllib.agents.pg.pg_torch_policy import pg_torch_loss __all__ = [ "PGTrainer", "pg_tf_loss", "pg_torch_loss", "post_process_advantages", - "DEFAULT_CONFIG", "PGPipeline" + "DEFAULT_CONFIG" ] diff --git a/rllib/agents/pg/pg.py b/rllib/agents/pg/pg.py index 6c86e751e..6a495f024 100644 --- a/rllib/agents/pg/pg.py +++ b/rllib/agents/pg/pg.py @@ -1,6 +1,8 @@ from ray.rllib.agents.trainer import with_common_config from ray.rllib.agents.trainer_template import build_trainer from ray.rllib.agents.pg.pg_tf_policy import PGTFPolicy +from ray.rllib.utils.experimental_dsl import ( + ParallelRollouts, ConcatBatches, TrainOneStep, StandardMetricsReporting) # yapf: disable # __sphinx_doc_begin__ @@ -22,8 +24,26 @@ def get_policy_class(config): return PGTFPolicy +# Experimental pipeline-based impl; enable with "use_pipeline_impl": True. +def training_pipeline(workers, config): + # Collects experiences in parallel from multiple RolloutWorker actors. + rollouts = ParallelRollouts(workers, mode="bulk_sync") + + # Combine experiences batches until we hit `train_batch_size` in size. + # Then, train the policy on those experiences and update the workers. + train_op = rollouts \ + .combine(ConcatBatches( + min_batch_size=config["train_batch_size"])) \ + .for_each(TrainOneStep(workers)) + + # Add on the standard episode reward, etc. metrics reporting. This returns + # a LocalIterator[metrics_dict] representing metrics for each train step. + return StandardMetricsReporting(train_op, workers, config) + + PGTrainer = build_trainer( name="PG", default_config=DEFAULT_CONFIG, default_policy=PGTFPolicy, - get_policy_class=get_policy_class) + get_policy_class=get_policy_class, + training_pipeline=training_pipeline) diff --git a/rllib/agents/pg/pg_pipeline.py b/rllib/agents/pg/pg_pipeline.py deleted file mode 100644 index 23ca07ae7..000000000 --- a/rllib/agents/pg/pg_pipeline.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Experimental pipeline-based impl; run this with --run='PG_pl'""" - -from ray.rllib.agents.pg.pg import PGTrainer -from ray.rllib.utils.experimental_dsl import ( - ParallelRollouts, ConcatBatches, TrainOneStep, StandardMetricsReporting) - - -def training_pipeline(workers, config): - # Collects experiences in parallel from multiple RolloutWorker actors. - rollouts = ParallelRollouts(workers, mode="bulk_sync") - - # Combine experiences batches until we hit `train_batch_size` in size. - # Then, train the policy on those experiences and update the workers. - train_op = rollouts \ - .combine(ConcatBatches( - min_batch_size=config["train_batch_size"])) \ - .for_each(TrainOneStep(workers)) - - # Add on the standard episode reward, etc. metrics reporting. This returns - # a LocalIterator[metrics_dict] representing metrics for each train step. - return StandardMetricsReporting(train_op, workers, config) - - -PGPipeline = PGTrainer.with_updates(training_pipeline=training_pipeline) diff --git a/rllib/agents/pg/tests/test_pg.py b/rllib/agents/pg/tests/test_pg.py index 62ff3ee0d..5ada2e144 100644 --- a/rllib/agents/pg/tests/test_pg.py +++ b/rllib/agents/pg/tests/test_pg.py @@ -3,7 +3,7 @@ import unittest import ray import ray.rllib.agents.pg as pg -from ray.rllib.agents.pg import PGPipeline +from ray.rllib.agents.pg import PGTrainer from ray.rllib.evaluation.postprocessing import Postprocessing from ray.rllib.models.tf.tf_action_dist import Categorical from ray.rllib.models.torch.torch_action_dist import TorchCategorical @@ -19,7 +19,12 @@ class TestPG(unittest.TestCase): ray.shutdown() def test_pg_pipeline(ray_start_regular): - trainer = PGPipeline(env="CartPole-v0", config={"min_iter_time_s": 0}) + trainer = PGTrainer( + env="CartPole-v0", + config={ + "min_iter_time_s": 0, + "use_pipeline_impl": True + }) assert isinstance(trainer.train(), dict) def test_pg_compilation(self): diff --git a/rllib/agents/registry.py b/rllib/agents/registry.py index b8b7de285..be6e0920a 100644 --- a/rllib/agents/registry.py +++ b/rllib/agents/registry.py @@ -100,21 +100,6 @@ def _import_marwil(): return marwil.MARWILTrainer -def _import_a2c_pipeline(): - from ray.rllib.agents import a3c - return a3c.A2CPipeline - - -def _import_a3c_pipeline(): - from ray.rllib.agents import a3c - return a3c.A3CPipeline - - -def _import_pg_pipeline(): - from ray.rllib.agents import pg - return pg.PGPipeline - - ALGORITHMS = { "SAC": _import_sac, "DDPG": _import_ddpg, @@ -135,11 +120,6 @@ ALGORITHMS = { "APPO": _import_appo, "DDPPO": _import_ddppo, "MARWIL": _import_marwil, - - # Experimental pipeline-based impls. - "A2C_pl": _import_a2c_pipeline, - "A3C_pl": _import_a3c_pipeline, - "PG_pl": _import_pg_pipeline, } diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 6cb4c471a..150cb31fb 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -213,6 +213,9 @@ COMMON_CONFIG = { # trainer guarantees all eval workers have the latest policy state before # this function is called. "custom_eval_function": None, + # EXPERIMENTAL: use the pipeline based implementation of the algo. Can also + # be enabled by setting RLLIB_USE_PIPELINE_IMPL=1. + "use_pipeline_impl": False, # === Advanced Rollout Settings === # Use a background thread for sampling (slightly off-policy, usually not diff --git a/rllib/agents/trainer_template.py b/rllib/agents/trainer_template.py index d5734cbc8..09bb33838 100644 --- a/rllib/agents/trainer_template.py +++ b/rllib/agents/trainer_template.py @@ -1,3 +1,5 @@ +import logging +import os import time from ray.rllib.agents.trainer import Trainer, COMMON_CONFIG @@ -5,6 +7,8 @@ from ray.rllib.optimizers import SyncSamplesOptimizer from ray.rllib.utils import add_mixins from ray.rllib.utils.annotations import override, DeveloperAPI +logger = logging.getLogger(__name__) + @DeveloperAPI def build_trainer(name, @@ -106,7 +110,9 @@ def build_trainer(name, self.train_pipeline = None self.optimizer = None - if training_pipeline: + if training_pipeline and (self.config["use_pipeline_impl"] or + "RLLIB_USE_PIPELINE_IMPL" in os.environ): + logger.warning("Using experimental pipeline based impl.") self.train_pipeline = training_pipeline(self.workers, config) elif make_policy_optimizer: self.optimizer = make_policy_optimizer(self.workers, config) diff --git a/rllib/tests/test_pipeline.py b/rllib/tests/test_pipeline.py index 03c2db218..abd80c910 100644 --- a/rllib/tests/test_pipeline.py +++ b/rllib/tests/test_pipeline.py @@ -1,7 +1,7 @@ import unittest import ray -from ray.rllib.agents.a3c import a2c_pipeline +from ray.rllib.agents.a3c import A2CTrainer class TestPipeline(unittest.TestCase): @@ -14,8 +14,12 @@ class TestPipeline(unittest.TestCase): ray.shutdown() def test_pipeline_stats(ray_start_regular): - trainer = a2c_pipeline.A2CPipeline( - env="CartPole-v0", config={"min_iter_time_s": 0}) + trainer = A2CTrainer( + env="CartPole-v0", + config={ + "min_iter_time_s": 0, + "use_pipeline_impl": True + }) result = trainer.train() assert isinstance(result, dict) assert "info" in result @@ -30,8 +34,12 @@ class TestPipeline(unittest.TestCase): assert "update_time_ms" in result["timers"] def test_pipeline_save_restore(ray_start_regular): - trainer = a2c_pipeline.A2CPipeline( - env="CartPole-v0", config={"min_iter_time_s": 0}) + trainer = A2CTrainer( + env="CartPole-v0", + config={ + "min_iter_time_s": 0, + "use_pipeline_impl": True + }) res1 = trainer.train() checkpoint = trainer.save() res2 = trainer.train() diff --git a/rllib/utils/experimental_dsl.py b/rllib/utils/experimental_dsl.py index e324c2b26..637fba842 100644 --- a/rllib/utils/experimental_dsl.py +++ b/rllib/utils/experimental_dsl.py @@ -3,17 +3,17 @@ TODO(ekl): describe the concepts.""" import logging -from typing import List, Any, Tuple +from typing import List, Any, Tuple, Union import time import ray from ray.util.iter import from_actors, LocalIterator from ray.util.iter_metrics import MetricsContext -from ray.rllib.evaluation.metrics import collect_episodes, summarize_episodes +from ray.rllib.evaluation.metrics import collect_episodes, \ + summarize_episodes, get_learner_stats from ray.rllib.evaluation.rollout_worker import get_global_worker from ray.rllib.evaluation.worker_set import WorkerSet -from ray.rllib.policy.sample_batch import SampleBatch -from ray.rllib.policy.policy import LEARNER_STATS_KEY +from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch logger = logging.getLogger(__name__) @@ -30,6 +30,13 @@ LEARNER_INFO = "learner" # Type aliases. GradientType = dict +SampleBatchType = Union[SampleBatch, MultiAgentBatch] + + +def _check_sample_batch_type(batch): + if not isinstance(batch, SampleBatchType.__args__): + raise ValueError("Expected either SampleBatch or MultiAgentBatch, " + "got {}: {}".format(type(batch), batch)) def ParallelRollouts(workers: WorkerSet, @@ -125,7 +132,7 @@ def AsyncGradients( (grads, info), count = item metrics = LocalIterator.get_metrics() metrics.counters[STEPS_SAMPLED_COUNTER] += count - metrics.info[LEARNER_INFO] = info[LEARNER_STATS_KEY] + metrics.info[LEARNER_INFO] = get_learner_stats(info) metrics.timers[GRAD_WAIT_TIMER].push(time.perf_counter() - self.fetch_start_time) return grads, count @@ -186,10 +193,8 @@ class ConcatBatches: if self.batch_start_time is None: self.batch_start_time = time.perf_counter() - def __call__(self, batch: SampleBatch) -> List[SampleBatch]: - if not isinstance(batch, SampleBatch): - raise ValueError("Expected type SampleBatch, got {}: {}".format( - type(batch), batch)) + def __call__(self, batch: SampleBatchType) -> List[SampleBatchType]: + _check_sample_batch_type(batch) self.buffer.append(batch) self.count += batch.count if self.count >= self.min_batch_size: @@ -222,14 +227,15 @@ class TrainOneStep: def __init__(self, workers: WorkerSet): self.workers = workers - def __call__(self, batch: SampleBatch) -> List[dict]: + def __call__(self, batch: SampleBatchType) -> List[dict]: + _check_sample_batch_type(batch) metrics = LocalIterator.get_metrics() learn_timer = metrics.timers[LEARN_ON_BATCH_TIMER] with learn_timer: info = self.workers.local_worker().learn_on_batch(batch) learn_timer.push_units_processed(batch.count) metrics.counters[STEPS_TRAINED_COUNTER] += batch.count - metrics.info[LEARNER_INFO] = info[LEARNER_STATS_KEY] + metrics.info[LEARNER_INFO] = get_learner_stats(info) if self.workers.remote_workers(): with metrics.timers[WORKER_UPDATE_TIMER]: weights = ray.put(self.workers.local_worker().get_weights()) @@ -338,11 +344,12 @@ class ComputeGradients: def __init__(self, workers): self.workers = workers - def __call__(self, samples): + def __call__(self, samples: SampleBatchType): + _check_sample_batch_type(samples) metrics = LocalIterator.get_metrics() with metrics.timers[COMPUTE_GRADS_TIMER]: grad, info = self.workers.local_worker().compute_gradients(samples) - metrics.info[LEARNER_INFO] = info[LEARNER_STATS_KEY] + metrics.info[LEARNER_INFO] = get_learner_stats(info) return grad, samples.count