From 243b1b728122a4c5839ecabcb1596ff0e43bf34e Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Thu, 14 Nov 2019 12:14:00 -0800 Subject: [PATCH] [rllib] Add microbatch optimizer with A2C example (#6161) --- doc/source/rllib-algorithms.rst | 2 + rllib/agents/a3c/a2c.py | 20 +++ rllib/optimizers/__init__.py | 2 + rllib/optimizers/microbatch_optimizer.py | 143 ++++++++++++++++++ .../cartpole-a2c-microbatch.yaml | 11 ++ 5 files changed, 178 insertions(+) create mode 100644 rllib/optimizers/microbatch_optimizer.py create mode 100644 rllib/tuned_examples/regression_tests/cartpole-a2c-microbatch.yaml diff --git a/doc/source/rllib-algorithms.rst b/doc/source/rllib-algorithms.rst index 04426d5e8..90826674e 100644 --- a/doc/source/rllib-algorithms.rst +++ b/doc/source/rllib-algorithms.rst @@ -126,6 +126,8 @@ Advantage Actor-Critic (A2C, A3C) `[paper] `__ `[implementation] `__ 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 `__. + .. figure:: a2c-arch.svg A2C architecture diff --git a/rllib/agents/a3c/a2c.py b/rllib/agents/a3c/a2c.py index 0b6592e74..0a1282f08 100644 --- a/rllib/agents/a3c/a2c.py +++ b/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) diff --git a/rllib/optimizers/__init__.py b/rllib/optimizers/__init__.py index 41ea54466..48a07efe8 100644 --- a/rllib/optimizers/__init__.py +++ b/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", diff --git a/rllib/optimizers/microbatch_optimizer.py b/rllib/optimizers/microbatch_optimizer.py new file mode 100644 index 000000000..c8bbdc04e --- /dev/null +++ b/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, + }) diff --git a/rllib/tuned_examples/regression_tests/cartpole-a2c-microbatch.yaml b/rllib/tuned_examples/regression_tests/cartpole-a2c-microbatch.yaml new file mode 100644 index 000000000..ac1bd74ed --- /dev/null +++ b/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