From 7e3e4cd321ac22906aca3f0e74cf690cfa9ef9f0 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Tue, 11 Aug 2020 23:58:41 -0700 Subject: [PATCH] [rllib] Execution plan API documentation (#10000) * wip * updte * comments --- doc/source/rllib-concepts.rst | 149 ++++++++++++++++++++++++++++++++-- doc/source/rllib-toc.rst | 2 +- 2 files changed, 143 insertions(+), 8 deletions(-) diff --git a/doc/source/rllib-concepts.rst b/doc/source/rllib-concepts.rst index 12e107e28..7f40ac4ef 100644 --- a/doc/source/rllib-concepts.rst +++ b/doc/source/rllib-concepts.rst @@ -1,7 +1,7 @@ RLlib Concepts and Custom Algorithms ==================================== -This page describes the internal concepts used to implement algorithms in RLlib. You might find this useful if modifying or adding new algorithms to RLlib. +This page describes the internal concepts used to implement algorithms in RLlib. You might find this useful if modifying or adding new algorithms to RLlib. The three main concepts covered here are `policies <#policies>`__, `policy evaluation <#policy-evaluation>`__, and `execution plans <#execution-plans>`__. RLlib relies on these concepts to define algorithms in a distributed, framework-agnostic way that is amenable to customization and operation in many types of RL environments. Policies -------- @@ -214,11 +214,11 @@ In the above section you saw how to compose a simple policy gradient algorithm w Besides some boilerplate for defining the PPO configuration and some warnings, the most important argument to take note of is the ``execution_plan``. -The trainer's `execution plan <#execution-plan>`__ defines the distributed training workflow. Depending on the ``simple_optimizer`` trainer config, PPO can switch between a simple synchronous plan, or a multi-GPU plan that implements minibatch SGD (the default): +The trainer's `execution plan <#execution-plans>`__ defines the distributed training workflow. Depending on the ``simple_optimizer`` trainer config, PPO can switch between a simple synchronous plan, or a multi-GPU plan that implements minibatch SGD (the default): .. code-block:: python - def execution_plan(workers, config): + def execution_plan(workers: WorkerSet, config: TrainerConfigDict): rollouts = ParallelRollouts(workers, mode="bulk_sync") # Collect large batches of relevant experiences & standardize. @@ -581,12 +581,147 @@ Here is an example of creating a set of rollout workers and using them gather ex for w in workers.remote_workers(): w.set_weights.remote(weights) -Execution Plan --------------- +Execution Plans +--------------- -.. note:: +Execution plans let you easily express the execution of an RL algorithm as a sequence of steps that +occur either sequentially in the learner, or in parallel across many actors. +Under the hood, RLlib *translates* these plans into ``ray.get()`` and ``ray.wait()`` operations over Ray actors, +so you easily write high-performance algorithms without needing to manage individual low-level Ray actor calls. - Policy optimizers have been replaced by the execution plan API. This documentation will be updated in the future. +Execution plans represent the **dataflow of the RL training job**. For example, the A2C algorithm can be thought +of a sequence of repeating steps, or *dataflow*, of: + + 1. ``ParallelRollouts``: Generate experiences from many envs in parallel using rollout workers. + 2. ``ConcatBatches``: The experiences are concatenated into one batch for training. + 3. ``TrainOneStep``: Take a gradient step with respect to the policy loss, and update the worker weights. + +In code, this dataflow can be expressed as the following execution plan, which is a simple function that can be passed to ``build_trainer`` to define a new algorithm. It takes in a ``WorkerSet`` and config, and returns an iterator over training results: + +.. code-block:: python + + def execution_plan(workers: WorkerSet, config: TrainerConfigDict): + # type: LocalIterator[SampleBatchType] + rollouts = ParallelRollouts(workers, mode="bulk_sync") + + # type: LocalIterator[(SampleBatchType, List[LearnerStatsDict])] + train_op = rollouts \ + .combine(ConcatBatches( + min_batch_size=config["train_batch_size"])) \ + .for_each(TrainOneStep(workers)) + + # type: LocalIterator[ResultDict] + return StandardMetricsReporting(train_op, workers, config) + + +As you can see, each step returns an *iterator* over objects (if you're unfamiliar with distributed iterators, see Ray's `parallel iterators documentation `__). The reason it is a ``LocalIterator`` is that, though it is based on a parallel computation, the iterator has been turned into one that can be consumed locally in sequence by the program. A couple other points to note: + + - The reason the plan returns an iterator over training results, is that ``trainer.train()`` is pulling results from this iterator to return as the result of the train call. + - The rollout workers have been already created ahead of time in the ``WorkerSet``, so the execution plan function is only defining a sequence of operations over the results of the rollouts. + +These iterators represent the infinite stream of data items that can be produced from the dataflow. Each operator (e.g., ``ConcatBatches``, ``TrainOneStep``), executes an operation over each item and returns a transformed item (e.g., concatenated batches, learner stats from training). Finally, some operators such as TrainOneStep have the *side-effect* of updating the rollout worker weights (that's why ``TrainOneStep`` takes the list of worker actors ``workers`` as an argument). + +Understanding and Debugging Execution Plans +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Execution plans are based on Ray `parallel iterators `__ and can be inspected similarly. For example, suppose you wanted to print out the intermediate data items during training. This can be done by inserting a print function into the dataflow, e.g., for A2C: + +.. code-block:: python + + def debug_print(item): + print("I saw", type(item)) + return item + + train_op = rollouts \ + .combine(ConcatBatches( + min_batch_size=config["train_batch_size"])) \ + .for_each(debug_print) \ + .for_each(TrainOneStep(workers)) + +You'll see output like this on the console: + +.. code-block:: bash + + (pid=6555) I saw + (pid=6555) I saw + (pid=6555) I saw + (pid=6555) I saw + +It is important to understand that the iterators of an execution plan are evaluated *lazily*. This means that no computation happens until the `trainer <#trainers>`__ tries to read the next item from the iterator (i.e., get the next training result for a ``Trainer.train()`` call). + +Execution Plan Concepts +~~~~~~~~~~~~~~~~~~~~~~~ + +RLlib provides a library of operators `(GitHub link) `__ that can be used in execution plans. You can of course write your own operators (which are just normal Python functions). As a reminder, operators are simply functions (or stateful function objects) that can be chained on the iterator (e.g., the ``debug_print`` operator above). A few categories of operators are summarized below: + +**Rollout ops** (`rollout_ops.py `__): These are functions for generating and working with experiences, including ``ParallelRollouts`` (for generating experiences synchronously or asynchronously), ``ConcatBatches`` (for combining batches together), ``SelectExperiences`` (for selecting relevant experiences in a multi-agent setting), and ``AsyncGradients`` (for computing gradients over new experiences on the fly, asynchronously, as in A3C). + +**Train ops** (`train_ops.py `__): These are functions that improve the policy and update workers. The most basic operator, ``TrainOneStep``, take in as input a batch of experiences and emit metrics as output. Important operators here include ``TrainOneStep``, ``TrainTFMultiGPU`` (for multi-GPU optimization), ``ComputeGradients`` (to compute gradients without updating the policy), and ``ApplyGradients`` (to apply computed gradients to a policy). + +**Replay ops** (`replay_ops.py `__): The main operator provided here is ``StoreToReplayBuffer``, which can save experiences batches to either a local replay buffer or a set of distributed replay actors. It has a counterpart, ``Replay``, that produces a new stream of experiences replayed from one of the aforementioned replay buffers. Algorithms that use ``StoreToReplayBuffer`` and ``Replay`` are necessarily composed of *multiple sub-dataflows* (different iterators), that are combined with *concurrency ops*. + +**Concurrency ops** (`concurrency_ops.py `__): The main operator provided here is ``Concurrently``, which composes multiple iterators (dataflows) into a single dataflow by executing them in an interleaved fashion. The output can be defined to be the mixture of the two dataflows, or filtered to that of one of the sub-dataflows. It has two modes: + + - ``round_robin``: Alternate taking items from each input dataflow. This ensures a fixed ratio of computations between, e.g., experience generation and experience replay. The ratio can be adjusted by setting ``round_robin_weights``. + - ``async``: Execute each input dataflow as fast as possible without blocking. You might want to use this when, e.g., you want replay to proceed as fast as possible irregardless of how fast experiences are being generated. + +**Metric ops** (`metric_ops.py `__): Finally, we provide a ``StandardMetricsReporting`` operator that collects training metrics from the rollout workers in a unified fashion, and returns a stream of training result dicts. Execution plans should always end with this operator. This metrics op also reports various internal performance metrics stored by other operators in the shared metrics context accessible via ``_get_shared_metrics()``. + +Example: Asynchrony +~~~~~~~~~~~~~~~~~~~~ + +Suppose we wanted to make the above A2C example asynchronous (i.e., A3C). We would switch the synchronous ``ParallelRollouts`` operator with ``AsyncGradients``, and use ``ApplyGradients`` to apply gradient updates as fast as they are collected. The ``AsyncGradients`` operator is going to execute rollouts in parallel, compute the policy gradient over the new batches (of size ``rollout_fragment_length``) on the remote workers, and then return a stream of the computed gradients: + +.. code-block:: python + + def execution_plan(workers: WorkerSet, config: TrainerConfigDict): + # type: LocalIterator[(ModelGradients, int)] + grads = AsyncGradients(workers) + + # type: LocalIterator[_] + train_op = grads.for_each(ApplyGradients(workers, update_all=False)) + + # type: LocalIterator[ResultDict] + return StandardMetricsReporting(train_op, workers, config) + +See also the `actual A3C implementation `__. + +Example: Replay +~~~~~~~~~~~~~~~ + +Let's try adding a replay buffer to A2C. This can be done as follows by inserting store / replay ops and using ``Concurrently`` to compose them together: + +.. code-block:: python + + def execution_plan(workers: WorkerSet, config: TrainerConfigDict): + # Construct a replay buffer. + replay_buffer = LocalReplayBuffer(...) + + # type: LocalIterator[_] + store_op = ParallelRollouts(workers, mode="bulk_sync") \ + .for_each(StoreToReplayBuffer(local_buffer=replay_buffer)) + + # type: LocalIterator[(SampleBatchType, List[LearnerStatsDict])] + replay_op = Replay(local_buffer=replay_buffer) \ + .for_each(TrainOneStep(workers)) + + # type: LocalIterator[(SampleBatchType, List[LearnerStatsDict])] + train_op = Concurrently( + [store_op, replay_op], mode="round_robin", output_indexes=[1]) + + # type: LocalIterator[ResultDict] + return StandardMetricsReporting(train_op, workers, config) + + +Note that here we set ``output_indexes=[1]`` for the ``Concurrently`` operator, which makes it only return results from the replay op. See also the `DQN implementation of replay `__ for a complete example including the implementation of options such as *training intensity*. + + +Example: Multi-agent +~~~~~~~~~~~~~~~~~~~~ + +One of the primary motivations behind execution plans, beyond their conciseness, is to enable complex multi-agent training workflows to be easily composed. For example, suppose one wanted to, in a multi-agent environment, concurrently train one set of agents with ``DQN``, and another set with ``PPO``. This requires stitching together two entirely different distributed dataflows. Fortunately, as we've seen earlier, this is quite simple with the ``Concurrently`` operator. + +Check out the `PPO + DQN multi-agent workflow example `__ for more details. One line to pay particular attention to in this example is the use of ``LocalIterator.duplicate()`` to clone the iterator of experiences into two separate iterators, which are filtered via ``SelectExperiences`` and then consumed by PPO and DQN sub-dataflows respectively. Trainers -------- diff --git a/doc/source/rllib-toc.rst b/doc/source/rllib-toc.rst index 0dcbea76d..e460305f5 100644 --- a/doc/source/rllib-toc.rst +++ b/doc/source/rllib-toc.rst @@ -156,7 +156,7 @@ Concepts and Custom Algorithms - `Extending Existing Policies `__ * `Policy Evaluation `__ -* `Policy Optimization `__ +* `Execution Plans `__ * `Trainers `__ Examples