From 31b40b00f67163e9723648fac8a4de5cb99063dc Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 10 Apr 2020 00:56:08 -0700 Subject: [PATCH] [rllib] Pull out experimental dsl into rllib.execution module, add initial unit tests (#7958) --- rllib/BUILD | 7 + rllib/agents/a3c/a2c.py | 7 +- rllib/agents/a3c/a3c.py | 5 +- rllib/agents/dqn/apex.py | 12 +- rllib/agents/dqn/dqn.py | 8 +- rllib/agents/pg/pg.py | 5 +- rllib/execution/common.py | 40 ++ rllib/execution/concurrency_ops.py | 94 ++++ rllib/execution/metric_ops.py | 123 +++++ rllib/execution/replay_ops.py | 151 ++++++ rllib/execution/rollout_ops.py | 164 +++++++ rllib/execution/train_ops.py | 201 ++++++++ rllib/tests/conftest.py | 1 + rllib/tests/test_execution.py | 161 +++++++ rllib/tests/test_io.py | 4 +- rllib/utils/experimental_dsl.py | 742 ----------------------------- 16 files changed, 967 insertions(+), 758 deletions(-) create mode 100644 rllib/execution/common.py create mode 100644 rllib/execution/concurrency_ops.py create mode 100644 rllib/execution/metric_ops.py create mode 100644 rllib/execution/replay_ops.py create mode 100644 rllib/execution/rollout_ops.py create mode 100644 rllib/execution/train_ops.py create mode 100644 rllib/tests/conftest.py create mode 100644 rllib/tests/test_execution.py delete mode 100644 rllib/utils/experimental_dsl.py diff --git a/rllib/BUILD b/rllib/BUILD index 132fc0d0c..e90f77649 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1076,6 +1076,13 @@ py_test( srcs = ["tests/test_io.py"] ) +py_test( + name = "tests/test_execution", + tags = ["tests_dir", "tests_dir_E"], + size = "medium", + srcs = ["tests/test_execution.py"] +) + py_test( name = "tests/test_local", tags = ["tests_dir", "tests_dir_L"], diff --git a/rllib/agents/a3c/a2c.py b/rllib/agents/a3c/a2c.py index 030188462..8ddef62ac 100644 --- a/rllib/agents/a3c/a2c.py +++ b/rllib/agents/a3c/a2c.py @@ -5,10 +5,11 @@ from ray.rllib.agents.a3c.a3c import DEFAULT_CONFIG as A3C_CONFIG, \ 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.execution.rollout_ops import ParallelRollouts, ConcatBatches +from ray.rllib.execution.train_ops import ComputeGradients, AverageGradients, \ + ApplyGradients, TrainOneStep +from ray.rllib.execution.metric_ops import StandardMetricsReporting 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, diff --git a/rllib/agents/a3c/a3c.py b/rllib/agents/a3c/a3c.py index 30d49712f..6583f121e 100644 --- a/rllib/agents/a3c/a3c.py +++ b/rllib/agents/a3c/a3c.py @@ -3,9 +3,10 @@ import logging 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.execution.rollout_ops import AsyncGradients +from ray.rllib.execution.train_ops import ApplyGradients +from ray.rllib.execution.metric_ops import StandardMetricsReporting from ray.rllib.optimizers import AsyncGradientsOptimizer -from ray.rllib.utils.experimental_dsl import (AsyncGradients, ApplyGradients, - StandardMetricsReporting) logger = logging.getLogger(__name__) diff --git a/rllib/agents/dqn/apex.py b/rllib/agents/dqn/apex.py index 0604343eb..7dba27371 100644 --- a/rllib/agents/dqn/apex.py +++ b/rllib/agents/dqn/apex.py @@ -2,14 +2,16 @@ import collections import ray from ray.rllib.agents.dqn.dqn import DQNTrainer, DEFAULT_CONFIG as DQN_CONFIG +from ray.rllib.execution.common import STEPS_TRAINED_COUNTER +from ray.rllib.execution.rollout_ops import ParallelRollouts +from ray.rllib.execution.concurrency_ops import Concurrently, Enqueue, Dequeue +from ray.rllib.execution.replay_ops import StoreToReplayActors, ParallelReplay +from ray.rllib.execution.train_ops import UpdateTargetNetwork +from ray.rllib.execution.metric_ops import StandardMetricsReporting from ray.rllib.optimizers import AsyncReplayOptimizer from ray.rllib.optimizers.async_replay_optimizer import ReplayActor from ray.rllib.utils import merge_dicts from ray.rllib.utils.actors import create_colocated -from ray.rllib.utils.experimental_dsl import ( - ParallelRollouts, Concurrently, ParallelReplay, StandardMetricsReporting, - StoreToReplayActors, UpdateTargetNetwork, Enqueue, Dequeue, - STEPS_TRAINED_COUNTER) from ray.rllib.optimizers.async_replay_optimizer import LearnerThread from ray.util.iter import LocalIterator @@ -101,6 +103,8 @@ def execution_plan(workers, config): actor, prio_dict, count = item actor.update_priorities.remote(prio_dict) metrics = LocalIterator.get_metrics() + # Manually update the steps trained counter since the learner thread + # is executing outside the pipeline. metrics.counters[STEPS_TRAINED_COUNTER] += count metrics.timers["learner_dequeue"] = learner_thread.queue_timer metrics.timers["learner_grad"] = learner_thread.grad_timer diff --git a/rllib/agents/dqn/dqn.py b/rllib/agents/dqn/dqn.py index 7777b6d9c..00a940a6b 100644 --- a/rllib/agents/dqn/dqn.py +++ b/rllib/agents/dqn/dqn.py @@ -8,9 +8,11 @@ from ray.rllib.optimizers import SyncReplayOptimizer from ray.rllib.optimizers.replay_buffer import ReplayBuffer from ray.rllib.utils.deprecation import deprecation_warning, DEPRECATED_VALUE from ray.rllib.utils.exploration import PerWorkerEpsilonGreedy -from ray.rllib.utils.experimental_dsl import ( - ParallelRollouts, Concurrently, StoreToReplayBuffer, LocalReplay, - TrainOneStep, StandardMetricsReporting, UpdateTargetNetwork) +from ray.rllib.execution.rollout_ops import ParallelRollouts +from ray.rllib.execution.concurrency_ops import Concurrently +from ray.rllib.execution.replay_ops import StoreToReplayBuffer, LocalReplay +from ray.rllib.execution.train_ops import TrainOneStep, UpdateTargetNetwork +from ray.rllib.execution.metric_ops import StandardMetricsReporting logger = logging.getLogger(__name__) diff --git a/rllib/agents/pg/pg.py b/rllib/agents/pg/pg.py index a23c42fd7..fd31eb52c 100644 --- a/rllib/agents/pg/pg.py +++ b/rllib/agents/pg/pg.py @@ -1,8 +1,9 @@ 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) +from ray.rllib.execution.rollout_ops import ParallelRollouts, ConcatBatches +from ray.rllib.execution.train_ops import TrainOneStep +from ray.rllib.execution.metric_ops import StandardMetricsReporting # yapf: disable # __sphinx_doc_begin__ diff --git a/rllib/execution/common.py b/rllib/execution/common.py new file mode 100644 index 000000000..7f865b5be --- /dev/null +++ b/rllib/execution/common.py @@ -0,0 +1,40 @@ +from typing import Union + +from ray.util.iter import LocalIterator +from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch + +# Counters for training progress (keys for metrics.counters). +STEPS_SAMPLED_COUNTER = "num_steps_sampled" +STEPS_TRAINED_COUNTER = "num_steps_trained" + +# Counters to track target network updates. +LAST_TARGET_UPDATE_TS = "last_target_update_ts" +NUM_TARGET_UPDATES = "num_target_updates" + +# Performance timers (keys for metrics.timers). +APPLY_GRADS_TIMER = "apply_grad" +COMPUTE_GRADS_TIMER = "compute_grads" +WORKER_UPDATE_TIMER = "update" +GRAD_WAIT_TIMER = "grad_wait" +SAMPLE_TIMER = "sample" +LEARN_ON_BATCH_TIMER = "learn" + +# Instant metrics (keys for metrics.info). +LEARNER_INFO = "learner" + +# Type aliases. +GradientType = dict +SampleBatchType = Union[SampleBatch, MultiAgentBatch] + + +# Asserts that an object is a type of SampleBatch. +def _check_sample_batch_type(batch): + if not isinstance(batch, SampleBatchType.__args__): + raise ValueError("Expected either SampleBatch or MultiAgentBatch, " + "got {}: {}".format(type(batch), batch)) + + +# Returns pipeline global vars that should be periodically sent to each worker. +def _get_global_vars(): + metrics = LocalIterator.get_metrics() + return {"timestep": metrics.counters[STEPS_SAMPLED_COUNTER]} diff --git a/rllib/execution/concurrency_ops.py b/rllib/execution/concurrency_ops.py new file mode 100644 index 000000000..73865252d --- /dev/null +++ b/rllib/execution/concurrency_ops.py @@ -0,0 +1,94 @@ +from typing import List +import queue + +from ray.util.iter import LocalIterator, _NextValueNotReady +from ray.util.iter_metrics import SharedMetrics + + +def Concurrently(ops: List[LocalIterator], *, mode="round_robin"): + """Operator that runs the given parent iterators concurrently. + + Arguments: + mode (str): One of {'round_robin', 'async'}. + - In 'round_robin' mode, we alternate between pulling items from + each parent iterator in order deterministically. + - In 'async' mode, we pull from each parent iterator as fast as + they are produced. This is non-deterministic. + + >>> sim_op = ParallelRollouts(...).for_each(...) + >>> replay_op = LocalReplay(...).for_each(...) + >>> combined_op = Concurrently([sim_op, replay_op], mode="async") + """ + + if len(ops) < 2: + raise ValueError("Should specify at least 2 ops.") + if mode == "round_robin": + deterministic = True + elif mode == "async": + deterministic = False + else: + raise ValueError("Unknown mode {}".format(mode)) + return ops[0].union(*ops[1:], deterministic=deterministic) + + +class Enqueue: + """Enqueue data items into a queue.Queue instance. + + The enqueue is non-blocking, so Enqueue operations can executed with + Dequeue via the Concurrently() operator. + + Examples: + >>> queue = queue.Queue(100) + >>> write_op = ParallelRollouts(...).for_each(Enqueue(queue)) + >>> read_op = Dequeue(queue) + >>> combined_op = Concurrently([write_op, read_op], mode="async") + >>> next(combined_op) + SampleBatch(...) + """ + + def __init__(self, output_queue: queue.Queue): + if not isinstance(output_queue, queue.Queue): + raise ValueError("Expected queue.Queue, got {}".format( + type(output_queue))) + self.queue = output_queue + + def __call__(self, x): + try: + self.queue.put_nowait(x) + except queue.Full: + return _NextValueNotReady() + + +def Dequeue(input_queue: queue.Queue, check=lambda: True): + """Dequeue data items from a queue.Queue instance. + + The dequeue is non-blocking, so Dequeue operations can executed with + Enqueue via the Concurrently() operator. + + Arguments: + input_queue (Queue): queue to pull items from. + check (fn): liveness check. When this function returns false, + Dequeue() will raise an error to halt execution. + + Examples: + >>> queue = queue.Queue(100) + >>> write_op = ParallelRollouts(...).for_each(Enqueue(queue)) + >>> read_op = Dequeue(queue) + >>> combined_op = Concurrently([write_op, read_op], mode="async") + >>> next(combined_op) + SampleBatch(...) + """ + if not isinstance(input_queue, queue.Queue): + raise ValueError("Expected queue.Queue, got {}".format( + type(input_queue))) + + def base_iterator(timeout=None): + while check(): + try: + item = input_queue.get_nowait() + yield item + except queue.Empty: + yield _NextValueNotReady() + raise RuntimeError("Error raised reading from queue") + + return LocalIterator(base_iterator, SharedMetrics()) diff --git a/rllib/execution/metric_ops.py b/rllib/execution/metric_ops.py new file mode 100644 index 000000000..357538296 --- /dev/null +++ b/rllib/execution/metric_ops.py @@ -0,0 +1,123 @@ +from typing import Any +import time + +from ray.util.iter import LocalIterator +from ray.rllib.evaluation.metrics import collect_episodes, summarize_episodes +from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER +from ray.rllib.evaluation.worker_set import WorkerSet + + +def StandardMetricsReporting(train_op: LocalIterator[Any], workers: WorkerSet, + config: dict) -> LocalIterator[dict]: + """Operator to periodically collect and report metrics. + + Arguments: + train_op (LocalIterator): Operator for executing training steps. + We ignore the output values. + workers (WorkerSet): Rollout workers to collect metrics from. + config (dict): Trainer configuration, used to determine the frequency + of stats reporting. + + Returns: + A local iterator over training results. + + Examples: + >>> train_op = ParallelRollouts(...).for_each(TrainOneStep(...)) + >>> metrics_op = StandardMetricsReporting(train_op, workers, config) + >>> next(metrics_op) + {"episode_reward_max": ..., "episode_reward_mean": ..., ...} + """ + + output_op = train_op \ + .filter(OncePerTimeInterval(max(2, config["min_iter_time_s"]))) \ + .for_each(CollectMetrics( + workers, min_history=config["metrics_smoothing_episodes"], + timeout_seconds=config["collect_metrics_timeout"])) + return output_op + + +class CollectMetrics: + """Callable that collects metrics from workers. + + The metrics are smoothed over a given history window. + + This should be used with the .for_each() operator. For a higher level + API, consider using StandardMetricsReporting instead. + + Examples: + >>> output_op = train_op.for_each(CollectMetrics(workers)) + >>> print(next(output_op)) + {"episode_reward_max": ..., "episode_reward_mean": ..., ...} + """ + + def __init__(self, workers, min_history=100, timeout_seconds=180): + self.workers = workers + self.episode_history = [] + self.to_be_collected = [] + self.min_history = min_history + self.timeout_seconds = timeout_seconds + + def __call__(self, _): + # Collect worker metrics. + episodes, self.to_be_collected = collect_episodes( + self.workers.local_worker(), + self.workers.remote_workers(), + self.to_be_collected, + timeout_seconds=self.timeout_seconds) + orig_episodes = list(episodes) + missing = self.min_history - len(episodes) + if missing > 0: + episodes.extend(self.episode_history[-missing:]) + assert len(episodes) <= self.min_history + self.episode_history.extend(orig_episodes) + self.episode_history = self.episode_history[-self.min_history:] + res = summarize_episodes(episodes, orig_episodes) + + # Add in iterator metrics. + metrics = LocalIterator.get_metrics() + timers = {} + counters = {} + info = {} + info.update(metrics.info) + for k, counter in metrics.counters.items(): + counters[k] = counter + for k, timer in metrics.timers.items(): + timers["{}_time_ms".format(k)] = round(timer.mean * 1000, 3) + if timer.has_units_processed(): + timers["{}_throughput".format(k)] = round( + timer.mean_throughput, 3) + res.update({ + "num_healthy_workers": len(self.workers.remote_workers()), + "timesteps_total": metrics.counters[STEPS_SAMPLED_COUNTER], + }) + res["timers"] = timers + res["info"] = info + res["info"].update(counters) + return res + + +class OncePerTimeInterval: + """Callable that returns True once per given interval. + + This should be used with the .filter() operator to throttle / rate-limit + metrics reporting. For a higher-level API, consider using + StandardMetricsReporting instead. + + Examples: + >>> throttled_op = train_op.filter(OncePerTimeInterval(5)) + >>> start = time.time() + >>> next(throttled_op) + >>> print(time.time() - start) + 5.00001 # will be greater than 5 seconds + """ + + def __init__(self, delay): + self.delay = delay + self.last_called = 0 + + def __call__(self, item): + now = time.time() + if now - self.last_called > self.delay: + self.last_called = now + return True + return False diff --git a/rllib/execution/replay_ops.py b/rllib/execution/replay_ops.py new file mode 100644 index 000000000..6ad326e32 --- /dev/null +++ b/rllib/execution/replay_ops.py @@ -0,0 +1,151 @@ +from typing import List +import numpy as np +import random + +from ray.util.iter import from_actors, LocalIterator +from ray.util.iter_metrics import SharedMetrics +from ray.rllib.optimizers.replay_buffer import PrioritizedReplayBuffer, \ + ReplayBuffer +from ray.rllib.execution.common import SampleBatchType, STEPS_TRAINED_COUNTER +from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch, \ + DEFAULT_POLICY_ID +from ray.rllib.utils.compression import pack_if_needed + + +class StoreToReplayBuffer: + """Callable that stores data into a local replay buffer. + + This should be used with the .for_each() operator on a rollouts iterator. + The batch that was stored is returned. + + Examples: + >>> buf = ReplayBuffer(1000) + >>> rollouts = ParallelRollouts(...) + >>> store_op = rollouts.for_each(StoreToReplayBuffer(buf)) + >>> next(store_op) + SampleBatch(...) + """ + + def __init__(self, replay_buffer: ReplayBuffer): + assert isinstance(replay_buffer, ReplayBuffer) + self.replay_buffers = {DEFAULT_POLICY_ID: replay_buffer} + + def __call__(self, batch: SampleBatchType): + # Handle everything as if multiagent + if isinstance(batch, SampleBatch): + batch = MultiAgentBatch({DEFAULT_POLICY_ID: batch}, batch.count) + + for policy_id, s in batch.policy_batches.items(): + for row in s.rows(): + self.replay_buffers[policy_id].add( + pack_if_needed(row["obs"]), + row["actions"], + row["rewards"], + pack_if_needed(row["new_obs"]), + row["dones"], + weight=None) + return batch + + +class StoreToReplayActors: + """Callable that stores data into a replay buffer actors. + + This should be used with the .for_each() operator on a rollouts iterator. + The batch that was stored is returned. + + Examples: + >>> actors = [ReplayActor.remote() for _ in range(4)] + >>> rollouts = ParallelRollouts(...) + >>> store_op = rollouts.for_each(StoreToReplayActors(actors)) + >>> next(store_op) + SampleBatch(...) + """ + + def __init__(self, replay_actors: List["ActorHandle"]): + self.replay_actors = replay_actors + + def __call__(self, batch: SampleBatchType): + actor = random.choice(self.replay_actors) + actor.add_batch.remote(batch) + return batch + + +def ParallelReplay(replay_actors: List["ActorHandle"], async_queue_depth=4): + """Replay experiences in parallel from the given actors. + + This should be combined with the StoreToReplayActors operation using the + Concurrently() operator. + + Arguments: + replay_actors (list): List of replay actors. + async_queue_depth (int): In async mode, the max number of async + requests in flight per actor. + + Examples: + >>> actors = [ReplayActor.remote() for _ in range(4)] + >>> replay_op = ParallelReplay(actors) + >>> next(replay_op) + SampleBatch(...) + """ + replay = from_actors(replay_actors) + return replay.gather_async( + async_queue_depth=async_queue_depth).filter(lambda x: x is not None) + + +def LocalReplay(replay_buffer: ReplayBuffer, train_batch_size: int): + """Replay experiences from a local buffer instance. + + This should be combined with the StoreToReplayBuffer operation using the + Concurrently() operator. + + Arguments: + replay_buffer (ReplayBuffer): Buffer to replay experiences from. + train_batch_size (int): Batch size of fetches from the buffer. + + Examples: + >>> actors = [ReplayActor.remote() for _ in range(4)] + >>> replay_op = ParallelReplay(actors) + >>> next(replay_op) + SampleBatch(...) + """ + assert isinstance(replay_buffer, ReplayBuffer) + replay_buffers = {DEFAULT_POLICY_ID: replay_buffer} + # TODO(ekl) support more options, or combine with ParallelReplay (?) + synchronize_sampling = False + prioritized_replay_beta = None + + def gen_replay(timeout): + while True: + samples = {} + idxes = None + for policy_id, replay_buffer in replay_buffers.items(): + if synchronize_sampling: + if idxes is None: + idxes = replay_buffer.sample_idxes(train_batch_size) + else: + idxes = replay_buffer.sample_idxes(train_batch_size) + + if isinstance(replay_buffer, PrioritizedReplayBuffer): + metrics = LocalIterator.get_metrics() + num_steps_trained = metrics.counters[STEPS_TRAINED_COUNTER] + (obses_t, actions, rewards, obses_tp1, dones, weights, + batch_indexes) = replay_buffer.sample_with_idxes( + idxes, + beta=prioritized_replay_beta.value(num_steps_trained)) + else: + (obses_t, actions, rewards, obses_tp1, + dones) = replay_buffer.sample_with_idxes(idxes) + weights = np.ones_like(rewards) + batch_indexes = -np.ones_like(rewards) + samples[policy_id] = SampleBatch({ + "obs": obses_t, + "actions": actions, + "rewards": rewards, + "new_obs": obses_tp1, + "dones": dones, + "weights": weights, + "batch_indexes": batch_indexes + }) + yield MultiAgentBatch(samples, train_batch_size) + + return LocalIterator(gen_replay, SharedMetrics()) diff --git a/rllib/execution/rollout_ops.py b/rllib/execution/rollout_ops.py new file mode 100644 index 000000000..9b9d7f09f --- /dev/null +++ b/rllib/execution/rollout_ops.py @@ -0,0 +1,164 @@ +from typing import List, Tuple +import time + +from ray.util.iter import from_actors, LocalIterator +from ray.util.iter_metrics import SharedMetrics +from ray.rllib.evaluation.metrics import get_learner_stats +from ray.rllib.evaluation.rollout_worker import get_global_worker +from ray.rllib.evaluation.worker_set import WorkerSet +from ray.rllib.execution.common import GradientType, SampleBatchType, \ + STEPS_SAMPLED_COUNTER, LEARNER_INFO, SAMPLE_TIMER, \ + GRAD_WAIT_TIMER, _check_sample_batch_type +from ray.rllib.policy.sample_batch import SampleBatch + + +def ParallelRollouts(workers: WorkerSet, + *, + mode="bulk_sync", + async_queue_depth=1) -> LocalIterator[SampleBatch]: + """Operator to collect experiences in parallel from rollout workers. + + If there are no remote workers, experiences will be collected serially from + the local worker instance instead. + + Arguments: + workers (WorkerSet): set of rollout workers to use. + mode (str): One of {'async', 'bulk_sync'}. + - In 'async' mode, batches are returned as soon as they are + computed by rollout workers with no order guarantees. + - In 'bulk_sync' mode, we collect one batch from each worker + and concatenate them together into a large batch to return. + async_queue_depth (int): In async mode, the max number of async + requests in flight per actor. + + Returns: + A local iterator over experiences collected in parallel. + + Examples: + >>> rollouts = ParallelRollouts(workers, mode="async") + >>> batch = next(rollouts) + >>> print(batch.count) + 50 # config.rollout_fragment_length + + >>> rollouts = ParallelRollouts(workers, mode="bulk_sync") + >>> batch = next(rollouts) + >>> print(batch.count) + 200 # config.rollout_fragment_length * config.num_workers + + Updates the STEPS_SAMPLED_COUNTER counter in the local iterator context. + """ + + # Ensure workers are initially in sync. + workers.sync_weights() + + def report_timesteps(batch): + metrics = LocalIterator.get_metrics() + metrics.counters[STEPS_SAMPLED_COUNTER] += batch.count + return batch + + if not workers.remote_workers(): + # Handle the serial sampling case. + def sampler(_): + while True: + yield workers.local_worker().sample() + + return (LocalIterator(sampler, SharedMetrics()) + .for_each(report_timesteps)) + + # Create a parallel iterator over generated experiences. + rollouts = from_actors(workers.remote_workers()) + + if mode == "bulk_sync": + return rollouts \ + .batch_across_shards() \ + .for_each(lambda batches: SampleBatch.concat_samples(batches)) \ + .for_each(report_timesteps) + elif mode == "async": + return rollouts.gather_async( + async_queue_depth=async_queue_depth).for_each(report_timesteps) + else: + raise ValueError( + "mode must be one of 'bulk_sync', 'async', got '{}'".format(mode)) + + +def AsyncGradients( + workers: WorkerSet) -> LocalIterator[Tuple[GradientType, int]]: + """Operator to compute gradients in parallel from rollout workers. + + Arguments: + workers (WorkerSet): set of rollout workers to use. + + Returns: + A local iterator over policy gradients computed on rollout workers. + + Examples: + >>> grads_op = AsyncGradients(workers) + >>> print(next(grads_op)) + {"var_0": ..., ...}, 50 # grads, batch count + + Updates the STEPS_SAMPLED_COUNTER counter and LEARNER_INFO field in the + local iterator context. + """ + + # Ensure workers are initially in sync. + workers.sync_weights() + + # This function will be applied remotely on the workers. + def samples_to_grads(samples): + return get_global_worker().compute_gradients(samples), samples.count + + # Record learner metrics and pass through (grads, count). + class record_metrics: + def _on_fetch_start(self): + self.fetch_start_time = time.perf_counter() + + def __call__(self, item): + (grads, info), count = item + metrics = LocalIterator.get_metrics() + metrics.counters[STEPS_SAMPLED_COUNTER] += count + metrics.info[LEARNER_INFO] = get_learner_stats(info) + metrics.timers[GRAD_WAIT_TIMER].push(time.perf_counter() - + self.fetch_start_time) + return grads, count + + rollouts = from_actors(workers.remote_workers()) + grads = rollouts.for_each(samples_to_grads) + return grads.gather_async().for_each(record_metrics()) + + +class ConcatBatches: + """Callable used to merge batches into larger batches for training. + + This should be used with the .combine() operator. + + Examples: + >>> rollouts = ParallelRollouts(...) + >>> rollouts = rollouts.combine(ConcatBatches(min_batch_size=10000)) + >>> print(next(rollouts).count) + 10000 + """ + + def __init__(self, min_batch_size: int): + self.min_batch_size = min_batch_size + self.buffer = [] + self.count = 0 + self.batch_start_time = None + + def _on_fetch_start(self): + if self.batch_start_time is None: + self.batch_start_time = time.perf_counter() + + 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: + out = SampleBatch.concat_samples(self.buffer) + timer = LocalIterator.get_metrics().timers[SAMPLE_TIMER] + timer.push(time.perf_counter() - self.batch_start_time) + timer.push_units_processed(self.count) + self.batch_start_time = None + self.buffer = [] + self.count = 0 + return [out] + return [] diff --git a/rllib/execution/train_ops.py b/rllib/execution/train_ops.py new file mode 100644 index 000000000..14c7fedaf --- /dev/null +++ b/rllib/execution/train_ops.py @@ -0,0 +1,201 @@ +import logging +from typing import List + +import ray +from ray.util.iter import LocalIterator +from ray.rllib.evaluation.metrics import get_learner_stats +from ray.rllib.evaluation.worker_set import WorkerSet +from ray.rllib.execution.common import SampleBatchType, \ + STEPS_SAMPLED_COUNTER, STEPS_TRAINED_COUNTER, LEARNER_INFO, \ + APPLY_GRADS_TIMER, COMPUTE_GRADS_TIMER, WORKER_UPDATE_TIMER, \ + LEARN_ON_BATCH_TIMER, LAST_TARGET_UPDATE_TS, NUM_TARGET_UPDATES, \ + _get_global_vars, _check_sample_batch_type + +logger = logging.getLogger(__name__) + + +class TrainOneStep: + """Callable that improves the policy and updates workers. + + This should be used with the .for_each() operator. + + Examples: + >>> rollouts = ParallelRollouts(...) + >>> train_op = rollouts.for_each(TrainOneStep(workers)) + >>> print(next(train_op)) # This trains the policy on one batch. + {"learner_stats": ...} + + Updates the STEPS_TRAINED_COUNTER counter and LEARNER_INFO field in the + local iterator context. + """ + + def __init__(self, workers: WorkerSet): + self.workers = workers + + 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] = get_learner_stats(info) + if self.workers.remote_workers(): + with metrics.timers[WORKER_UPDATE_TIMER]: + weights = ray.put(self.workers.local_worker().get_weights()) + for e in self.workers.remote_workers(): + e.set_weights.remote(weights, _get_global_vars()) + # Also update global vars of the local worker. + self.workers.local_worker().set_global_vars(_get_global_vars()) + return info + + +class ComputeGradients: + """Callable that computes gradients with respect to the policy loss. + + This should be used with the .for_each() operator. + + Examples: + >>> grads_op = rollouts.for_each(ComputeGradients(workers)) + >>> print(next(grads_op)) + {"var_0": ..., ...}, 50 # grads, batch count + + Updates the LEARNER_INFO info field in the local iterator context. + """ + + def __init__(self, workers): + self.workers = workers + + 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] = get_learner_stats(info) + return grad, samples.count + + +class ApplyGradients: + """Callable that applies gradients and updates workers. + + This should be used with the .for_each() operator. + + Examples: + >>> apply_op = grads_op.for_each(ApplyGradients(workers)) + >>> print(next(apply_op)) + None + + Updates the STEPS_TRAINED_COUNTER counter in the local iterator context. + """ + + def __init__(self, workers, update_all=True): + """Creates an ApplyGradients instance. + + Arguments: + workers (WorkerSet): workers to apply gradients to. + update_all (bool): If true, updates all workers. Otherwise, only + update the worker that produced the sample batch we are + currently processing (i.e., A3C style). + """ + self.workers = workers + self.update_all = update_all + + def __call__(self, item): + if not isinstance(item, tuple) or len(item) != 2: + raise ValueError( + "Input must be a tuple of (grad_dict, count), got {}".format( + item)) + gradients, count = item + metrics = LocalIterator.get_metrics() + metrics.counters[STEPS_TRAINED_COUNTER] += count + + apply_timer = metrics.timers[APPLY_GRADS_TIMER] + with apply_timer: + self.workers.local_worker().apply_gradients(gradients) + apply_timer.push_units_processed(count) + + # Also update global vars of the local worker. + self.workers.local_worker().set_global_vars(_get_global_vars()) + + if self.update_all: + if self.workers.remote_workers(): + with metrics.timers[WORKER_UPDATE_TIMER]: + weights = ray.put( + self.workers.local_worker().get_weights()) + for e in self.workers.remote_workers(): + e.set_weights.remote(weights, _get_global_vars()) + else: + if metrics.current_actor is None: + raise ValueError( + "Could not find actor to update. When " + "update_all=False, `current_actor` must be set " + "in the iterator context.") + with metrics.timers[WORKER_UPDATE_TIMER]: + weights = self.workers.local_worker().get_weights() + metrics.current_actor.set_weights.remote( + weights, _get_global_vars()) + + +class AverageGradients: + """Callable that averages the gradients in a batch. + + This should be used with the .for_each() operator after a set of gradients + have been batched with .batch(). + + Examples: + >>> batched_grads = grads_op.batch(32) + >>> avg_grads = batched_grads.for_each(AverageGradients()) + >>> print(next(avg_grads)) + {"var_0": ..., ...}, 1600 # averaged grads, summed batch count + """ + + def __call__(self, gradients): + acc = None + sum_count = 0 + for grad, count in gradients: + if acc is None: + acc = grad + else: + acc = [a + b for a, b in zip(acc, grad)] + sum_count += count + logger.info("Computing average of {} microbatch gradients " + "({} samples total)".format(len(gradients), sum_count)) + return acc, sum_count + + +class UpdateTargetNetwork: + """Periodically call policy.update_target() on all trainable policies. + + This should be used with the .for_each() operator after training step + has been taken. + + Examples: + >>> train_op = ParallelRollouts(...).for_each(TrainOneStep(...)) + >>> update_op = train_op.for_each( + ... UpdateTargetIfNeeded(workers, target_update_freq=500)) + >>> print(next(update_op)) + None + + Updates the LAST_TARGET_UPDATE_TS and NUM_TARGET_UPDATES counters in the + local iterator context. The value of the last update counter is used to + track when we should update the target next. + """ + + def __init__(self, workers, target_update_freq, by_steps_trained=False): + self.workers = workers + self.target_update_freq = target_update_freq + if by_steps_trained: + self.metric = STEPS_TRAINED_COUNTER + else: + self.metric = STEPS_SAMPLED_COUNTER + + def __call__(self, _): + metrics = LocalIterator.get_metrics() + cur_ts = metrics.counters[self.metric] + last_update = metrics.counters[LAST_TARGET_UPDATE_TS] + if cur_ts - last_update > self.target_update_freq: + self.workers.local_worker().foreach_trainable_policy( + lambda p, _: p.update_target()) + metrics.counters[NUM_TARGET_UPDATES] += 1 + metrics.counters[LAST_TARGET_UPDATE_TS] = cur_ts diff --git a/rllib/tests/conftest.py b/rllib/tests/conftest.py new file mode 100644 index 000000000..2b71d4c1b --- /dev/null +++ b/rllib/tests/conftest.py @@ -0,0 +1 @@ +from ray.tests.conftest import ray_start_regular_shared # noqa: F401 diff --git a/rllib/tests/test_execution.py b/rllib/tests/test_execution.py new file mode 100644 index 000000000..dba288b93 --- /dev/null +++ b/rllib/tests/test_execution.py @@ -0,0 +1,161 @@ +import pytest +import time +import gym +import queue + +from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicy +from ray.rllib.evaluation.worker_set import WorkerSet +from ray.rllib.evaluation.rollout_worker import RolloutWorker +from ray.rllib.execution.concurrency_ops import Concurrently, Enqueue, Dequeue +from ray.rllib.execution.metric_ops import StandardMetricsReporting +from ray.rllib.execution.rollout_ops import ParallelRollouts, AsyncGradients, \ + ConcatBatches +from ray.rllib.execution.train_ops import TrainOneStep, ComputeGradients, \ + AverageGradients +from ray.util.iter import LocalIterator, from_range +from ray.util.iter_metrics import SharedMetrics + + +def iter_list(values): + return LocalIterator(lambda _: values, SharedMetrics()) + + +def make_workers(n): + local = RolloutWorker( + env_creator=lambda _: gym.make("CartPole-v0"), + policy=PPOTFPolicy, + rollout_fragment_length=100) + remotes = [ + RolloutWorker.as_remote().remote( + env_creator=lambda _: gym.make("CartPole-v0"), + policy=PPOTFPolicy, + rollout_fragment_length=100) for _ in range(n) + ] + workers = WorkerSet._from_existing(local, remotes) + return workers + + +def test_concurrently(ray_start_regular_shared): + a = iter_list([1, 2, 3]) + b = iter_list([4, 5, 6]) + c = Concurrently([a, b], mode="round_robin") + assert c.take(6) == [1, 4, 2, 5, 3, 6] + + a = iter_list([1, 2, 3]) + b = iter_list([4, 5, 6]) + c = Concurrently([a, b], mode="async") + assert c.take(6) == [1, 2, 3, 4, 5, 6] + + +def test_enqueue_dequeue(ray_start_regular_shared): + a = iter_list([1, 2, 3]) + q = queue.Queue(100) + a.for_each(Enqueue(q)).take(3) + assert q.qsize() == 3 + assert q.get_nowait() == 1 + assert q.get_nowait() == 2 + assert q.get_nowait() == 3 + + q.put("a") + q.put("b") + q.put("c") + a = Dequeue(q) + assert a.take(3) == ["a", "b", "c"] + + +def test_metrics(ray_start_regular_shared): + workers = make_workers(1) + workers.foreach_worker(lambda w: w.sample()) + a = from_range(10, repeat=True).gather_sync() + b = StandardMetricsReporting( + a, workers, { + "min_iter_time_s": 2.5, + "metrics_smoothing_episodes": 10, + "collect_metrics_timeout": 10, + }) + + start = time.time() + res1 = next(b) + assert res1["episode_reward_mean"] > 0, res1 + res2 = next(b) + assert res2["episode_reward_mean"] > 0, res2 + assert time.time() - start > 2.4 + workers.stop() + + +def test_rollouts(ray_start_regular_shared): + workers = make_workers(2) + a = ParallelRollouts(workers, mode="bulk_sync") + assert next(a).count == 200 + counters = a.shared_metrics.get().counters + assert counters["num_steps_sampled"] == 200, counters + a = ParallelRollouts(workers, mode="async") + assert next(a).count == 100 + counters = a.shared_metrics.get().counters + assert counters["num_steps_sampled"] == 100, counters + workers.stop() + + +def test_rollouts_local(ray_start_regular_shared): + workers = make_workers(0) + a = ParallelRollouts(workers, mode="bulk_sync") + assert next(a).count == 100 + counters = a.shared_metrics.get().counters + assert counters["num_steps_sampled"] == 100, counters + workers.stop() + + +def test_concat_batches(ray_start_regular_shared): + workers = make_workers(0) + a = ParallelRollouts(workers, mode="async") + b = a.combine(ConcatBatches(1000)) + assert next(b).count == 1000 + timers = b.shared_metrics.get().timers + assert "sample" in timers + + +def test_async_grads(ray_start_regular_shared): + workers = make_workers(2) + a = AsyncGradients(workers) + res1 = next(a) + assert isinstance(res1, tuple) and len(res1) == 2, res1 + counters = a.shared_metrics.get().counters + assert counters["num_steps_sampled"] == 100, counters + workers.stop() + + +def test_train_one_step(ray_start_regular_shared): + workers = make_workers(0) + a = ParallelRollouts(workers, mode="bulk_sync") + b = a.for_each(TrainOneStep(workers)) + assert "learner_stats" in next(b) + counters = a.shared_metrics.get().counters + assert counters["num_steps_sampled"] == 100, counters + assert counters["num_steps_trained"] == 100, counters + timers = a.shared_metrics.get().timers + assert "learn" in timers + workers.stop() + + +def test_compute_gradients(ray_start_regular_shared): + workers = make_workers(0) + a = ParallelRollouts(workers, mode="bulk_sync") + b = a.for_each(ComputeGradients(workers)) + grads, counts = next(b) + assert counts == 100, counts + timers = a.shared_metrics.get().timers + assert "compute_grads" in timers + + +def test_avg_gradients(ray_start_regular_shared): + workers = make_workers(0) + a = ParallelRollouts(workers, mode="bulk_sync") + b = a.for_each(ComputeGradients(workers)).batch(4) + c = b.for_each(AverageGradients()) + grads, counts = next(c) + assert counts == 400, counts + + +if __name__ == "__main__": + import sys + sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/tests/test_io.py b/rllib/tests/test_io.py index a31e83629..cf0ce481e 100644 --- a/rllib/tests/test_io.py +++ b/rllib/tests/test_io.py @@ -222,7 +222,7 @@ class JsonIOTest(unittest.TestCase): def test_write_file_uri(self): ioctx = IOContext(self.test_dir, {}, 0, None) writer = JsonWriter( - "file:" + self.test_dir, + "file://" + self.test_dir, ioctx, max_file_size=1000, compress_columns=["obs"]) @@ -278,7 +278,7 @@ class JsonIOTest(unittest.TestCase): reader = JsonReader([ self.test_dir + "/empty", self.test_dir + "/f1", - "file:" + self.test_dir + "/f2", + "file://" + self.test_dir + "/f2", ]) seen_a = set() for i in range(100): diff --git a/rllib/utils/experimental_dsl.py b/rllib/utils/experimental_dsl.py deleted file mode 100644 index 28658b44f..000000000 --- a/rllib/utils/experimental_dsl.py +++ /dev/null @@ -1,742 +0,0 @@ -"""Experimental distributed execution API. - -TODO(ekl): describe the concepts.""" - -import logging -from typing import List, Any, Tuple, Union -import numpy as np -import queue -import random -import time - -import ray -from ray.util.iter import from_actors, LocalIterator, _NextValueNotReady -from ray.util.iter_metrics import SharedMetrics -from ray.rllib.optimizers.replay_buffer import PrioritizedReplayBuffer, \ - ReplayBuffer -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, MultiAgentBatch, \ - DEFAULT_POLICY_ID -from ray.rllib.utils.compression import pack_if_needed - -logger = logging.getLogger(__name__) - -# Counters for training progress (keys for metrics.counters). -STEPS_SAMPLED_COUNTER = "num_steps_sampled" -STEPS_TRAINED_COUNTER = "num_steps_trained" - -# Counters to track target network updates. -LAST_TARGET_UPDATE_TS = "last_target_update_ts" -NUM_TARGET_UPDATES = "num_target_updates" - -# Performance timers (keys for metrics.timers). -APPLY_GRADS_TIMER = "apply_grad" -COMPUTE_GRADS_TIMER = "compute_grads" -WORKER_UPDATE_TIMER = "update" -GRAD_WAIT_TIMER = "grad_wait" -SAMPLE_TIMER = "sample" -LEARN_ON_BATCH_TIMER = "learn" - -# Instant metrics (keys for metrics.info). -LEARNER_INFO = "learner" - -# Type aliases. -GradientType = dict -SampleBatchType = Union[SampleBatch, MultiAgentBatch] - - -# Asserts that an object is a type of SampleBatch. -def _check_sample_batch_type(batch): - if not isinstance(batch, SampleBatchType.__args__): - raise ValueError("Expected either SampleBatch or MultiAgentBatch, " - "got {}: {}".format(type(batch), batch)) - - -# Returns pipeline global vars that should be periodically sent to each worker. -def _get_global_vars(): - metrics = LocalIterator.get_metrics() - return {"timestep": metrics.counters[STEPS_SAMPLED_COUNTER]} - - -def ParallelRollouts(workers: WorkerSet, mode="bulk_sync", - async_queue_depth=1) -> LocalIterator[SampleBatch]: - """Operator to collect experiences in parallel from rollout workers. - - If there are no remote workers, experiences will be collected serially from - the local worker instance instead. - - Arguments: - workers (WorkerSet): set of rollout workers to use. - mode (str): One of {'async', 'bulk_sync'}. - - In 'async' mode, batches are returned as soon as they are - computed by rollout workers with no order guarantees. - - In 'bulk_sync' mode, we collect one batch from each worker - and concatenate them together into a large batch to return. - async_queue_depth (int): In async mode, the max number of async - requests in flight per actor. - - Returns: - A local iterator over experiences collected in parallel. - - Examples: - >>> rollouts = ParallelRollouts(workers, mode="async") - >>> batch = next(rollouts) - >>> print(batch.count) - 50 # config.rollout_fragment_length - - >>> rollouts = ParallelRollouts(workers, mode="bulk_sync") - >>> batch = next(rollouts) - >>> print(batch.count) - 200 # config.rollout_fragment_length * config.num_workers - - Updates the STEPS_SAMPLED_COUNTER counter in the local iterator context. - """ - - # Ensure workers are initially in sync. - workers.sync_weights() - - def report_timesteps(batch): - metrics = LocalIterator.get_metrics() - metrics.counters[STEPS_SAMPLED_COUNTER] += batch.count - return batch - - if not workers.remote_workers(): - # Handle the serial sampling case. - def sampler(_): - while True: - yield workers.local_worker().sample() - - return (LocalIterator(sampler, SharedMetrics()) - .for_each(report_timesteps)) - - # Create a parallel iterator over generated experiences. - rollouts = from_actors(workers.remote_workers()) - - if mode == "bulk_sync": - return rollouts \ - .batch_across_shards() \ - .for_each(lambda batches: SampleBatch.concat_samples(batches)) \ - .for_each(report_timesteps) - elif mode == "async": - return rollouts.gather_async( - async_queue_depth=async_queue_depth).for_each(report_timesteps) - else: - raise ValueError( - "mode must be one of 'bulk_sync', 'async', got '{}'".format(mode)) - - -def AsyncGradients( - workers: WorkerSet) -> LocalIterator[Tuple[GradientType, int]]: - """Operator to compute gradients in parallel from rollout workers. - - Arguments: - workers (WorkerSet): set of rollout workers to use. - - Returns: - A local iterator over policy gradients computed on rollout workers. - - Examples: - >>> grads_op = AsyncGradients(workers) - >>> print(next(grads_op)) - {"var_0": ..., ...}, 50 # grads, batch count - - Updates the STEPS_SAMPLED_COUNTER counter and LEARNER_INFO field in the - local iterator context. - """ - - # Ensure workers are initially in sync. - workers.sync_weights() - - # This function will be applied remotely on the workers. - def samples_to_grads(samples): - return get_global_worker().compute_gradients(samples), samples.count - - # Record learner metrics and pass through (grads, count). - class record_metrics: - def _on_fetch_start(self): - self.fetch_start_time = time.perf_counter() - - def __call__(self, item): - (grads, info), count = item - metrics = LocalIterator.get_metrics() - metrics.counters[STEPS_SAMPLED_COUNTER] += count - metrics.info[LEARNER_INFO] = get_learner_stats(info) - metrics.timers[GRAD_WAIT_TIMER].push(time.perf_counter() - - self.fetch_start_time) - return grads, count - - rollouts = from_actors(workers.remote_workers()) - grads = rollouts.for_each(samples_to_grads) - return grads.gather_async().for_each(record_metrics()) - - -def StandardMetricsReporting(train_op: LocalIterator[Any], workers: WorkerSet, - config: dict) -> LocalIterator[dict]: - """Operator to periodically collect and report metrics. - - Arguments: - train_op (LocalIterator): Operator for executing training steps. - We ignore the output values. - workers (WorkerSet): Rollout workers to collect metrics from. - config (dict): Trainer configuration, used to determine the frequency - of stats reporting. - - Returns: - A local iterator over training results. - - Examples: - >>> train_op = ParallelRollouts(...).for_each(TrainOneStep(...)) - >>> metrics_op = StandardMetricsReporting(train_op, workers, config) - >>> next(metrics_op) - {"episode_reward_max": ..., "episode_reward_mean": ..., ...} - """ - - output_op = train_op \ - .filter(OncePerTimeInterval(max(2, config["min_iter_time_s"]))) \ - .for_each(CollectMetrics( - workers, min_history=config["metrics_smoothing_episodes"], - timeout_seconds=config["collect_metrics_timeout"])) - return output_op - - -class ConcatBatches: - """Callable used to merge batches into larger batches for training. - - This should be used with the .combine() operator. - - Examples: - >>> rollouts = ParallelRollouts(...) - >>> rollouts = rollouts.combine(ConcatBatches(min_batch_size=10000)) - >>> print(next(rollouts).count) - 10000 - """ - - def __init__(self, min_batch_size: int): - self.min_batch_size = min_batch_size - self.buffer = [] - self.count = 0 - self.batch_start_time = None - - def _on_fetch_start(self): - if self.batch_start_time is None: - self.batch_start_time = time.perf_counter() - - 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: - out = SampleBatch.concat_samples(self.buffer) - timer = LocalIterator.get_metrics().timers[SAMPLE_TIMER] - timer.push(time.perf_counter() - self.batch_start_time) - timer.push_units_processed(self.count) - self.batch_start_time = None - self.buffer = [] - self.count = 0 - return [out] - return [] - - -class TrainOneStep: - """Callable that improves the policy and updates workers. - - This should be used with the .for_each() operator. - - Examples: - >>> rollouts = ParallelRollouts(...) - >>> train_op = rollouts.for_each(TrainOneStep(workers)) - >>> print(next(train_op)) # This trains the policy on one batch. - None - - Updates the STEPS_TRAINED_COUNTER counter and LEARNER_INFO field in the - local iterator context. - """ - - def __init__(self, workers: WorkerSet): - self.workers = workers - - 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] = get_learner_stats(info) - if self.workers.remote_workers(): - with metrics.timers[WORKER_UPDATE_TIMER]: - weights = ray.put(self.workers.local_worker().get_weights()) - for e in self.workers.remote_workers(): - e.set_weights.remote(weights, _get_global_vars()) - # Also update global vars of the local worker. - self.workers.local_worker().set_global_vars(_get_global_vars()) - return info - - -class CollectMetrics: - """Callable that collects metrics from workers. - - The metrics are smoothed over a given history window. - - This should be used with the .for_each() operator. For a higher level - API, consider using StandardMetricsReporting instead. - - Examples: - >>> output_op = train_op.for_each(CollectMetrics(workers)) - >>> print(next(output_op)) - {"episode_reward_max": ..., "episode_reward_mean": ..., ...} - """ - - def __init__(self, workers, min_history=100, timeout_seconds=180): - self.workers = workers - self.episode_history = [] - self.to_be_collected = [] - self.min_history = min_history - self.timeout_seconds = timeout_seconds - - def __call__(self, _): - # Collect worker metrics. - episodes, self.to_be_collected = collect_episodes( - self.workers.local_worker(), - self.workers.remote_workers(), - self.to_be_collected, - timeout_seconds=self.timeout_seconds) - orig_episodes = list(episodes) - missing = self.min_history - len(episodes) - if missing > 0: - episodes.extend(self.episode_history[-missing:]) - assert len(episodes) <= self.min_history - self.episode_history.extend(orig_episodes) - self.episode_history = self.episode_history[-self.min_history:] - res = summarize_episodes(episodes, orig_episodes) - - # Add in iterator metrics. - metrics = LocalIterator.get_metrics() - timers = {} - counters = {} - info = {} - info.update(metrics.info) - for k, counter in metrics.counters.items(): - counters[k] = counter - for k, timer in metrics.timers.items(): - timers["{}_time_ms".format(k)] = round(timer.mean * 1000, 3) - if timer.has_units_processed(): - timers["{}_throughput".format(k)] = round( - timer.mean_throughput, 3) - res.update({ - "num_healthy_workers": len(self.workers.remote_workers()), - "timesteps_total": metrics.counters[STEPS_SAMPLED_COUNTER], - }) - res["timers"] = timers - res["info"] = info - res["info"].update(counters) - return res - - -class OncePerTimeInterval: - """Callable that returns True once per given interval. - - This should be used with the .filter() operator to throttle / rate-limit - metrics reporting. For a higher-level API, consider using - StandardMetricsReporting instead. - - Examples: - >>> throttled_op = train_op.filter(OncePerTimeInterval(5)) - >>> start = time.time() - >>> next(throttled_op) - >>> print(time.time() - start) - 5.00001 # will be greater than 5 seconds - """ - - def __init__(self, delay): - self.delay = delay - self.last_called = 0 - - def __call__(self, item): - now = time.time() - if now - self.last_called > self.delay: - self.last_called = now - return True - return False - - -class ComputeGradients: - """Callable that computes gradients with respect to the policy loss. - - This should be used with the .for_each() operator. - - Examples: - >>> grads_op = rollouts.for_each(ComputeGradients(workers)) - >>> print(next(grads_op)) - {"var_0": ..., ...}, 50 # grads, batch count - - Updates the LEARNER_INFO info field in the local iterator context. - """ - - def __init__(self, workers): - self.workers = workers - - 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] = get_learner_stats(info) - return grad, samples.count - - -class ApplyGradients: - """Callable that applies gradients and updates workers. - - This should be used with the .for_each() operator. - - Examples: - >>> apply_op = grads_op.for_each(ApplyGradients(workers)) - >>> print(next(apply_op)) - None - - Updates the STEPS_TRAINED_COUNTER counter in the local iterator context. - """ - - def __init__(self, workers, update_all=True): - """Creates an ApplyGradients instance. - - Arguments: - workers (WorkerSet): workers to apply gradients to. - update_all (bool): If true, updates all workers. Otherwise, only - update the worker that produced the sample batch we are - currently processing (i.e., A3C style). - """ - self.workers = workers - self.update_all = update_all - - def __call__(self, item): - if not isinstance(item, tuple) or len(item) != 2: - raise ValueError( - "Input must be a tuple of (grad_dict, count), got {}".format( - item)) - gradients, count = item - metrics = LocalIterator.get_metrics() - metrics.counters[STEPS_TRAINED_COUNTER] += count - - apply_timer = metrics.timers[APPLY_GRADS_TIMER] - with apply_timer: - self.workers.local_worker().apply_gradients(gradients) - apply_timer.push_units_processed(count) - - # Also update global vars of the local worker. - self.workers.local_worker().set_global_vars(_get_global_vars()) - - if self.update_all: - if self.workers.remote_workers(): - with metrics.timers[WORKER_UPDATE_TIMER]: - weights = ray.put( - self.workers.local_worker().get_weights()) - for e in self.workers.remote_workers(): - e.set_weights.remote(weights, _get_global_vars()) - else: - if metrics.current_actor is None: - raise ValueError( - "Could not find actor to update. When " - "update_all=False, `current_actor` must be set " - "in the iterator context.") - with metrics.timers[WORKER_UPDATE_TIMER]: - weights = self.workers.local_worker().get_weights() - metrics.current_actor.set_weights.remote( - weights, _get_global_vars()) - - -class AverageGradients: - """Callable that averages the gradients in a batch. - - This should be used with the .for_each() operator after a set of gradients - have been batched with .batch(). - - Examples: - >>> batched_grads = grads_op.batch(32) - >>> avg_grads = batched_grads.for_each(AverageGradients()) - >>> print(next(avg_grads)) - {"var_0": ..., ...}, 1600 # averaged grads, summed batch count - """ - - def __call__(self, gradients): - acc = None - sum_count = 0 - for grad, count in gradients: - if acc is None: - acc = grad - else: - acc = [a + b for a, b in zip(acc, grad)] - sum_count += count - logger.info("Computing average of {} microbatch gradients " - "({} samples total)".format(len(gradients), sum_count)) - return acc, sum_count - - -class StoreToReplayBuffer: - """Callable that stores data into a local replay buffer. - - This should be used with the .for_each() operator on a rollouts iterator. - The batch that was stored is returned. - - Examples: - >>> buf = ReplayBuffer(1000) - >>> rollouts = ParallelRollouts(...) - >>> store_op = rollouts.for_each(StoreToReplayBuffer(buf)) - >>> next(store_op) - SampleBatch(...) - """ - - def __init__(self, replay_buffer: ReplayBuffer): - assert isinstance(replay_buffer, ReplayBuffer) - self.replay_buffers = {DEFAULT_POLICY_ID: replay_buffer} - - def __call__(self, batch: SampleBatchType): - # Handle everything as if multiagent - if isinstance(batch, SampleBatch): - batch = MultiAgentBatch({DEFAULT_POLICY_ID: batch}, batch.count) - - for policy_id, s in batch.policy_batches.items(): - for row in s.rows(): - self.replay_buffers[policy_id].add( - pack_if_needed(row["obs"]), - row["actions"], - row["rewards"], - pack_if_needed(row["new_obs"]), - row["dones"], - weight=None) - return batch - - -class StoreToReplayActors: - """Callable that stores data into a replay buffer actors. - - This should be used with the .for_each() operator on a rollouts iterator. - The batch that was stored is returned. - - Examples: - >>> actors = [ReplayActor.remote() for _ in range(4)] - >>> rollouts = ParallelRollouts(...) - >>> store_op = rollouts.for_each(StoreToReplayActors(actors)) - >>> next(store_op) - SampleBatch(...) - """ - - def __init__(self, replay_actors: List["ActorHandle"]): - self.replay_actors = replay_actors - - def __call__(self, batch: SampleBatchType): - actor = random.choice(self.replay_actors) - actor.add_batch.remote(batch) - return batch - - -def ParallelReplay(replay_actors: List["ActorHandle"], async_queue_depth=4): - """Replay experiences in parallel from the given actors. - - This should be combined with the StoreToReplayActors operation using the - Concurrently() operator. - - Arguments: - replay_actors (list): List of replay actors. - async_queue_depth (int): In async mode, the max number of async - requests in flight per actor. - - Examples: - >>> actors = [ReplayActor.remote() for _ in range(4)] - >>> replay_op = ParallelReplay(actors) - >>> next(replay_op) - SampleBatch(...) - """ - replay = from_actors(replay_actors) - return replay.gather_async( - async_queue_depth=async_queue_depth).filter(lambda x: x is not None) - - -def LocalReplay(replay_buffer: ReplayBuffer, train_batch_size: int): - """Replay experiences from a local buffer instance. - - This should be combined with the StoreToReplayBuffer operation using the - Concurrently() operator. - - Arguments: - replay_buffer (ReplayBuffer): Buffer to replay experiences from. - train_batch_size (int): Batch size of fetches from the buffer. - - Examples: - >>> actors = [ReplayActor.remote() for _ in range(4)] - >>> replay_op = ParallelReplay(actors) - >>> next(replay_op) - SampleBatch(...) - """ - assert isinstance(replay_buffer, ReplayBuffer) - replay_buffers = {DEFAULT_POLICY_ID: replay_buffer} - # TODO(ekl) support more options, or combine with ParallelReplay (?) - synchronize_sampling = False - prioritized_replay_beta = None - - def gen_replay(timeout): - while True: - samples = {} - idxes = None - for policy_id, replay_buffer in replay_buffers.items(): - if synchronize_sampling: - if idxes is None: - idxes = replay_buffer.sample_idxes(train_batch_size) - else: - idxes = replay_buffer.sample_idxes(train_batch_size) - - if isinstance(replay_buffer, PrioritizedReplayBuffer): - metrics = LocalIterator.get_metrics() - num_steps_trained = metrics.counters[STEPS_TRAINED_COUNTER] - (obses_t, actions, rewards, obses_tp1, dones, weights, - batch_indexes) = replay_buffer.sample_with_idxes( - idxes, - beta=prioritized_replay_beta.value(num_steps_trained)) - else: - (obses_t, actions, rewards, obses_tp1, - dones) = replay_buffer.sample_with_idxes(idxes) - weights = np.ones_like(rewards) - batch_indexes = -np.ones_like(rewards) - samples[policy_id] = SampleBatch({ - "obs": obses_t, - "actions": actions, - "rewards": rewards, - "new_obs": obses_tp1, - "dones": dones, - "weights": weights, - "batch_indexes": batch_indexes - }) - yield MultiAgentBatch(samples, train_batch_size) - - return LocalIterator(gen_replay, SharedMetrics()) - - -def Concurrently(ops: List[LocalIterator], mode="round_robin"): - """Operator that runs the given parent iterators concurrently. - - Arguments: - mode (str): One of {'round_robin', 'async'}. - - In 'round_robin' mode, we alternate between pulling items from - each parent iterator in order deterministically. - - In 'async' mode, we pull from each parent iterator as fast as - they are produced. This is non-deterministic. - - >>> sim_op = ParallelRollouts(...).for_each(...) - >>> replay_op = LocalReplay(...).for_each(...) - >>> combined_op = Concurrently([sim_op, replay_op]) - """ - - if len(ops) < 2: - raise ValueError("Should specify at least 2 ops.") - if mode == "round_robin": - deterministic = True - elif mode == "async": - deterministic = False - else: - raise ValueError("Unknown mode {}".format(mode)) - return ops[0].union(*ops[1:], deterministic=deterministic) - - -class UpdateTargetNetwork: - """Periodically call policy.update_target() on all trainable policies. - - This should be used with the .for_each() operator after training step - has been taken. - - Examples: - >>> train_op = ParallelRollouts(...).for_each(TrainOneStep(...)) - >>> update_op = train_op.for_each( - ... UpdateTargetIfNeeded(workers, target_update_freq=500)) - >>> print(next(update_op)) - None - - Updates the LAST_TARGET_UPDATE_TS and NUM_TARGET_UPDATES counters in the - local iterator context. The value of the last update counter is used to - track when we should update the target next. - """ - - def __init__(self, workers, target_update_freq, by_steps_trained=False): - self.workers = workers - self.target_update_freq = target_update_freq - if by_steps_trained: - self.metric = STEPS_TRAINED_COUNTER - else: - self.metric = STEPS_SAMPLED_COUNTER - - def __call__(self, _): - metrics = LocalIterator.get_metrics() - cur_ts = metrics.counters[self.metric] - last_update = metrics.counters[LAST_TARGET_UPDATE_TS] - if cur_ts - last_update > self.target_update_freq: - self.workers.local_worker().foreach_trainable_policy( - lambda p, _: p.update_target()) - metrics.counters[NUM_TARGET_UPDATES] += 1 - metrics.counters[LAST_TARGET_UPDATE_TS] = cur_ts - - -class Enqueue: - """Enqueue data items into a queue.Queue instance. - - The enqueue is non-blocking, so Enqueue operations can executed with - Dequeue via the Concurrently() operator. - - Examples: - >>> queue = queue.Queue(100) - >>> write_op = ParallelRollouts(...).for_each(Enqueue(queue)) - >>> read_op = Dequeue(queue) - >>> combined_op = Concurrently([write_op, read_op], mode="async") - >>> next(combined_op) - SampleBatch(...) - """ - - def __init__(self, output_queue: queue.Queue): - if not isinstance(output_queue, queue.Queue): - raise ValueError("Expected queue.Queue, got {}".format( - type(output_queue))) - self.queue = output_queue - - def __call__(self, x): - try: - self.queue.put_nowait(x) - except queue.Full: - return _NextValueNotReady() - - -def Dequeue(input_queue: queue.Queue, check=lambda: True): - """Dequeue data items from a queue.Queue instance. - - The dequeue is non-blocking, so Dequeue operations can executed with - Enqueue via the Concurrently() operator. - - Arguments: - input_queue (Queue): queue to pull items from. - check (fn): liveness check. When this function returns false, - Dequeue() will raise an error to halt execution. - - Examples: - >>> queue = queue.Queue(100) - >>> write_op = ParallelRollouts(...).for_each(Enqueue(queue)) - >>> read_op = Dequeue(queue) - >>> combined_op = Concurrently([write_op, read_op], mode="async") - >>> next(combined_op) - SampleBatch(...) - """ - if not isinstance(input_queue, queue.Queue): - raise ValueError("Expected queue.Queue, got {}".format( - type(input_queue))) - - def base_iterator(timeout=None): - while check(): - try: - item = input_queue.get_nowait() - yield item - except queue.Empty: - yield _NextValueNotReady() - raise RuntimeError("Error raised reading from queue") - - return LocalIterator(base_iterator, SharedMetrics())