mirror of
https://github.com/wassname/ray.git
synced 2026-07-07 16:32:59 +08:00
Eric Liang
281 lines
9.2 KiB
Python
281 lines
9.2 KiB
Python
"""Experimental operators for defining distributed training pipelines.
|
|
|
|
TODO(ekl): describe the concepts."""
|
|
|
|
from typing import List, Any
|
|
import time
|
|
|
|
import ray
|
|
from ray.util.iter import from_actors, LocalIterator
|
|
from ray.rllib.evaluation.metrics import collect_episodes, summarize_episodes
|
|
from ray.rllib.evaluation.worker_set import WorkerSet
|
|
from ray.rllib.policy.sample_batch import SampleBatch
|
|
|
|
|
|
def ParallelRollouts(workers: WorkerSet,
|
|
mode="bulk_sync") -> 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.
|
|
|
|
Returns:
|
|
A local iterator over experiences collected in parallel.
|
|
|
|
Examples:
|
|
>>> rollouts = ParallelRollouts(workers, mode="async")
|
|
>>> batch = next(rollouts)
|
|
>>> print(batch.count)
|
|
50 # config.sample_batch_size
|
|
|
|
>>> rollouts = ParallelRollouts(workers, mode="bulk_sync")
|
|
>>> batch = next(rollouts)
|
|
>>> print(batch.count)
|
|
200 # config.sample_batch_size * config.num_workers
|
|
"""
|
|
|
|
if not workers.remote_workers():
|
|
# Handle the serial sampling case.
|
|
def sampler(_):
|
|
while True:
|
|
yield workers.local_worker().sample()
|
|
|
|
return LocalIterator(sampler)
|
|
|
|
# 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))
|
|
elif mode == "async":
|
|
return rollouts.gather_async()
|
|
else:
|
|
raise ValueError(
|
|
"mode must be one of 'bulk_sync', 'async', got '{}'".format(mode))
|
|
|
|
|
|
def StandardMetricsReporting(train_op: LocalIterator[Any], workers: WorkerSet,
|
|
config: 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(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
|
|
|
|
def __call__(self, batch: SampleBatch) -> List[SampleBatch]:
|
|
if not isinstance(batch, SampleBatch):
|
|
raise ValueError("Expected type SampleBatch, got {}: {}".format(
|
|
type(batch), batch))
|
|
self.buffer.append(batch)
|
|
self.count += batch.count
|
|
if self.count >= self.min_batch_size:
|
|
out = SampleBatch.concat_samples(self.buffer)
|
|
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.
|
|
{"learner_stats": {"policy_loss": ...}}
|
|
"""
|
|
|
|
def __init__(self, workers: WorkerSet):
|
|
self.workers = workers
|
|
|
|
def __call__(self, batch: SampleBatch) -> List[dict]:
|
|
info = self.workers.local_worker().learn_on_batch(batch)
|
|
if self.workers.remote_workers():
|
|
weights = ray.put(self.workers.local_worker().get_weights())
|
|
for e in self.workers.remote_workers():
|
|
e.set_weights.remote(weights)
|
|
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, info):
|
|
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)
|
|
res.update(info=info)
|
|
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": ..., ...}, {"learner_stats": ...} # grads, learner info
|
|
"""
|
|
|
|
def __init__(self, workers):
|
|
self.workers = workers
|
|
|
|
def __call__(self, samples):
|
|
grad, info = self.workers.local_worker().compute_gradients(samples)
|
|
return grad, info
|
|
|
|
|
|
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))
|
|
{"learner_stats": ...} # learner info
|
|
"""
|
|
|
|
def __init__(self, workers):
|
|
self.workers = workers
|
|
|
|
def __call__(self, item):
|
|
gradients, info = item
|
|
self.workers.local_worker().apply_gradients(gradients)
|
|
if self.workers.remote_workers():
|
|
weights = ray.put(self.workers.local_worker().get_weights())
|
|
for e in self.workers.remote_workers():
|
|
e.set_weights.remote(weights)
|
|
return info
|
|
|
|
|
|
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": ..., ...}, {"learner_stats": ...} # avg grads, last info
|
|
"""
|
|
|
|
def __call__(self, gradients):
|
|
acc = None
|
|
for grad, info in gradients:
|
|
if acc is None:
|
|
acc = grad
|
|
else:
|
|
acc = [a + b for a, b in zip(acc, grad)]
|
|
return acc, info
|