[rllib] Add microbatch optimizer with A2C example (#6161)

doc/source/rllib-algorithms.rst

@@ -126,6 +126,8 @@ Advantage Actor-Critic (A2C, A3C)
 `[paper] <https://arxiv.org/abs/1602.01783>`__ `[implementation] <https://github.com/ray-project/ray/blob/master/rllib/agents/a3c/a3c.py>`__
 RLlib implements A2C and A3C using SyncSamplesOptimizer and AsyncGradientsOptimizer respectively for policy optimization. These algorithms scale to up to 16-32 worker processes depending on the environment. Both a TensorFlow (LSTM), and PyTorch version are available.
 
+A2C also supports microbatching (i.e., gradient accumulation), which can be enabled by setting the ``microbatch_size`` config. Microbatching allows for training with a ``train_batch_size`` much larger than GPU memory. See also the `microbatch optimizer implementation <https://github.com/ray-project/ray/blob/master/rllib/optimizers/microbatch_optimizer.py>`__.
+
 .. figure:: a2c-arch.svg
 
     A2C architecture

rllib/agents/a3c/a2c.py

@@ -4,6 +4,7 @@ from __future__ import print_function
 
 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
@@ -14,12 +15,31 @@ A2C_DEFAULT_CONFIG = merge_dicts(
         "sample_batch_size": 20,
         "min_iter_time_s": 10,
         "sample_async": False,
+
+        # A2C supports microbatching, in which we accumulate gradients over
+        # batch of this size until the train batch size is reached. This allows
+        # training with batch sizes much larger than can fit in GPU memory.
+        # To enable, set this to a value less than the train batch size.
+        "microbatch_size": None,
     },
 )
 
+
+def choose_policy_optimizer(workers, config):
+    if config["microbatch_size"]:
+        return MicrobatchOptimizer(
+            workers,
+            train_batch_size=config["train_batch_size"],
+            microbatch_size=config["microbatch_size"])
+    else:
+        return SyncSamplesOptimizer(
+            workers, train_batch_size=config["train_batch_size"])
+
+
 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)

rllib/optimizers/__init__.py

@@ -7,6 +7,7 @@ from ray.rllib.optimizers.sync_samples_optimizer import SyncSamplesOptimizer
 from ray.rllib.optimizers.sync_replay_optimizer import SyncReplayOptimizer
 from ray.rllib.optimizers.sync_batch_replay_optimizer import \
     SyncBatchReplayOptimizer
+from ray.rllib.optimizers.microbatch_optimizer import MicrobatchOptimizer
 from ray.rllib.optimizers.multi_gpu_optimizer import LocalMultiGPUOptimizer
 
 __all__ = [
@@ -14,6 +15,7 @@ __all__ = [
     "AsyncReplayOptimizer",
     "AsyncSamplesOptimizer",
     "AsyncGradientsOptimizer",
+    "MicrobatchOptimizer",
     "SyncSamplesOptimizer",
     "SyncReplayOptimizer",
     "LocalMultiGPUOptimizer",

rllib/optimizers/microbatch_optimizer.py

@@ -0,0 +1,143 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+
+import ray
+from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer
+from ray.rllib.policy.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \
+    MultiAgentBatch
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.filter import RunningStat
+from ray.rllib.utils.timer import TimerStat
+from ray.rllib.utils.memory import ray_get_and_free
+
+logger = logging.getLogger(__name__)
+
+
+class MicrobatchOptimizer(PolicyOptimizer):
+    """A microbatching synchronous RL optimizer.
+
+    This optimizer pulls sample batches from workers until the target
+    microbatch size is reached. Then, it computes and accumulates the policy
+    gradient in a local buffer. This process is repeated until the number of
+    samples collected equals the train batch size. Then, an accumulated
+    gradient update is made.
+
+    This allows for training with effective batch sizes much larger than can
+    fit in GPU or host memory.
+    """
+
+    def __init__(self, workers, train_batch_size=10000, microbatch_size=1000):
+        PolicyOptimizer.__init__(self, workers)
+
+        if train_batch_size <= microbatch_size:
+            raise ValueError(
+                "The microbatch size must be smaller than the train batch "
+                "size, got {} vs {}".format(microbatch_size, train_batch_size))
+
+        self.update_weights_timer = TimerStat()
+        self.sample_timer = TimerStat()
+        self.grad_timer = TimerStat()
+        self.throughput = RunningStat()
+        self.train_batch_size = train_batch_size
+        self.microbatch_size = microbatch_size
+        self.learner_stats = {}
+        self.policies = dict(self.workers.local_worker()
+                             .foreach_trainable_policy(lambda p, i: (i, p)))
+        logger.debug("Policies to train: {}".format(self.policies))
+
+    @override(PolicyOptimizer)
+    def step(self):
+        with self.update_weights_timer:
+            if self.workers.remote_workers():
+                weights = ray.put(self.workers.local_worker().get_weights())
+                for e in self.workers.remote_workers():
+                    e.set_weights.remote(weights)
+
+        fetches = {}
+        accumulated_gradients = {}
+        samples_so_far = 0
+
+        # Accumulate minibatches.
+        i = 0
+        while samples_so_far < self.train_batch_size:
+            i += 1
+            with self.sample_timer:
+                samples = []
+                while sum(s.count for s in samples) < self.microbatch_size:
+                    if self.workers.remote_workers():
+                        samples.extend(
+                            ray_get_and_free([
+                                e.sample.remote()
+                                for e in self.workers.remote_workers()
+                            ]))
+                    else:
+                        samples.append(self.workers.local_worker().sample())
+                samples = SampleBatch.concat_samples(samples)
+                self.sample_timer.push_units_processed(samples.count)
+                samples_so_far += samples.count
+
+            logger.info(
+                "Computing gradients for microbatch {} ({}/{} samples)".format(
+                    i, samples_so_far, self.train_batch_size))
+
+            # Handle everything as if multiagent
+            if isinstance(samples, SampleBatch):
+                samples = MultiAgentBatch({
+                    DEFAULT_POLICY_ID: samples
+                }, samples.count)
+
+            with self.grad_timer:
+                for policy_id, policy in self.policies.items():
+                    if policy_id not in samples.policy_batches:
+                        continue
+                    batch = samples.policy_batches[policy_id]
+                    grad_out, info_out = (
+                        self.workers.local_worker().compute_gradients(
+                            MultiAgentBatch({
+                                policy_id: batch
+                            }, batch.count)))
+                    grad = grad_out[policy_id]
+                    fetches.update(info_out)
+                    if policy_id not in accumulated_gradients:
+                        accumulated_gradients[policy_id] = grad
+                    else:
+                        grad_size = len(accumulated_gradients[policy_id])
+                        assert grad_size == len(grad), (grad_size, len(grad))
+                        c = []
+                        for a, b in zip(accumulated_gradients[policy_id],
+                                        grad):
+                            c.append(a + b)
+                        accumulated_gradients[policy_id] = c
+            self.grad_timer.push_units_processed(samples.count)
+
+        # Apply the accumulated gradient
+        logger.info("Applying accumulated gradients ({} samples)".format(
+            samples_so_far))
+        self.workers.local_worker().apply_gradients(accumulated_gradients)
+
+        if len(fetches) == 1 and DEFAULT_POLICY_ID in fetches:
+            self.learner_stats = fetches[DEFAULT_POLICY_ID]
+        else:
+            self.learner_stats = fetches
+        self.num_steps_sampled += samples_so_far
+        self.num_steps_trained += samples_so_far
+        return self.learner_stats
+
+    @override(PolicyOptimizer)
+    def stats(self):
+        return dict(
+            PolicyOptimizer.stats(self), **{
+                "sample_time_ms": round(1000 * self.sample_timer.mean, 3),
+                "grad_time_ms": round(1000 * self.grad_timer.mean, 3),
+                "update_time_ms": round(1000 * self.update_weights_timer.mean,
+                                        3),
+                "opt_peak_throughput": round(self.grad_timer.mean_throughput,
+                                             3),
+                "sample_peak_throughput": round(
+                    self.sample_timer.mean_throughput, 3),
+                "opt_samples": round(self.grad_timer.mean_units_processed, 3),
+                "learner": self.learner_stats,
+            })

rllib/tuned_examples/regression_tests/cartpole-a2c-microbatch.yaml

@@ -0,0 +1,11 @@
+cartpole-a2c-microbatch:
+    env: CartPole-v0
+    run: A2C
+    stop:
+        episode_reward_mean: 100
+        timesteps_total: 100000
+    config:
+        num_workers: 1
+        gamma: 0.95
+        microbatch_size: 50
+        train_batch_size: 100