[rllib] Support multi-agent training in pipeline impls, add easy flag to enable (#7338)

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"]

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)

rllib/agents/a3c/a2c_pipeline.py

@@ -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)

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)

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)
 

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"
 ]

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)

rllib/agents/pg/pg_pipeline.py

@@ -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)

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):

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,
 }
 
 

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

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)

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()

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