wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-06 05:16:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rllib] Pull out experimental dsl into rllib.execution module, add initial unit tests (#7958)

Browse Source
This commit is contained in:
Eric Liang
2020-04-10 00:56:08 -07:00
committed by GitHub
parent 0d713e3eba
commit 31b40b00f6
16 changed files with 967 additions and 758 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/BUILD
+7
View File
@@ -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"],
rllib/agents/a3c/a2c.py
+4 -3
View File
@@ -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,
rllib/agents/a3c/a3c.py
+3 -2
View File
@@ -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__)
rllib/agents/dqn/apex.py
+8 -4
View File
@@ -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
rllib/agents/dqn/dqn.py
+5 -3
View File
@@ -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__)
rllib/agents/pg/pg.py
+3 -2
View File
@@ -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__
rllib/execution/common.py
+40
View File
@@ -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]}
rllib/execution/concurrency_ops.py
+94
View File
@@ -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())
rllib/execution/metric_ops.py
+123
View File
@@ -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
rllib/execution/replay_ops.py
+151
View File
@@ -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())
rllib/execution/rollout_ops.py
+164
View File
@@ -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 []
rllib/execution/train_ops.py
+201
View File
@@ -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
rllib/tests/conftest.py
+1
View File
@@ -0,0 +1 @@
from ray.tests.conftest import ray_start_regular_shared # noqa: F401
rllib/tests/test_execution.py
+161
View File
@@ -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__]))
rllib/tests/test_io.py
+2 -2
View File
@@ -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):
rllib/utils/experimental_dsl.py
-742
View File
@@ -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())