From a2d121520051de562244705d163cc6629beba6a2 Mon Sep 17 00:00:00 2001 From: Michael Luo Date: Thu, 24 Dec 2020 06:30:33 -0800 Subject: [PATCH] [RLlib] Execution Annotation (#13036) --- rllib/execution/common.py | 8 +++-- rllib/execution/concurrency_ops.py | 15 ++++++---- rllib/execution/learner_thread.py | 13 ++++---- rllib/execution/metric_ops.py | 18 +++++------ rllib/execution/minibatch_buffer.py | 13 ++++++-- rllib/execution/multi_gpu_learner.py | 31 +++++++++---------- rllib/execution/replay_buffer.py | 45 ++++++++++++++-------------- rllib/execution/replay_ops.py | 20 +++++++------ rllib/execution/segment_tree.py | 22 ++++++++------ rllib/execution/train_ops.py | 23 +++++++------- rllib/execution/tree_agg.py | 16 +++++----- 11 files changed, 126 insertions(+), 98 deletions(-) diff --git a/rllib/execution/common.py b/rllib/execution/common.py index 9e15f8b36..b12e557d5 100644 --- a/rllib/execution/common.py +++ b/rllib/execution/common.py @@ -1,5 +1,7 @@ from ray.util.iter import LocalIterator from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch +from ray.rllib.utils.typing import Dict, SampleBatchType +from ray.util.iter_metrics import MetricsContext # Counters for training progress (keys for metrics.counters). STEPS_SAMPLED_COUNTER = "num_steps_sampled" @@ -23,19 +25,19 @@ LEARNER_INFO = "learner" # Asserts that an object is a type of SampleBatch. -def _check_sample_batch_type(batch): +def _check_sample_batch_type(batch: SampleBatchType) -> None: if not isinstance(batch, (SampleBatch, MultiAgentBatch)): 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(): +def _get_global_vars() -> Dict: metrics = LocalIterator.get_metrics() return {"timestep": metrics.counters[STEPS_SAMPLED_COUNTER]} -def _get_shared_metrics(): +def _get_shared_metrics() -> MetricsContext: """Return shared metrics for the training workflow. This only applies if this trainer has an execution plan.""" diff --git a/rllib/execution/concurrency_ops.py b/rllib/execution/concurrency_ops.py index 7b057852e..cfe326ed9 100644 --- a/rllib/execution/concurrency_ops.py +++ b/rllib/execution/concurrency_ops.py @@ -1,15 +1,17 @@ -from typing import List +from typing import List, Optional, Any import queue from ray.util.iter import LocalIterator, _NextValueNotReady from ray.util.iter_metrics import SharedMetrics +from ray.rllib.utils.typing import SampleBatchType def Concurrently(ops: List[LocalIterator], *, - mode="round_robin", - output_indexes=None, - round_robin_weights=None): + mode: str = "round_robin", + output_indexes: Optional[List[int]] = None, + round_robin_weights: Optional[List[int]] = None + ) -> LocalIterator[SampleBatchType]: """Operator that runs the given parent iterators concurrently. Args: @@ -91,7 +93,7 @@ class Enqueue: type(output_queue))) self.queue = output_queue - def __call__(self, x): + def __call__(self, x: Any) -> Any: try: self.queue.put_nowait(x) except queue.Full: @@ -99,7 +101,8 @@ class Enqueue: return x -def Dequeue(input_queue: queue.Queue, check=lambda: True): +def Dequeue(input_queue: queue.Queue, + check=lambda: True) -> LocalIterator[SampleBatchType]: """Dequeue data items from a queue.Queue instance. The dequeue is non-blocking, so Dequeue operations can executed with diff --git a/rllib/execution/learner_thread.py b/rllib/execution/learner_thread.py index 9e905148d..8f5350fa1 100644 --- a/rllib/execution/learner_thread.py +++ b/rllib/execution/learner_thread.py @@ -1,3 +1,4 @@ +from typing import Dict import threading import copy @@ -8,6 +9,7 @@ from ray.rllib.execution.minibatch_buffer import MinibatchBuffer from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.timer import TimerStat from ray.rllib.utils.window_stat import WindowStat +from ray.rllib.evaluation.rollout_worker import RolloutWorker tf1, tf, tfv = try_import_tf() @@ -21,8 +23,9 @@ class LearnerThread(threading.Thread): improves overall throughput. """ - def __init__(self, local_worker, minibatch_buffer_size, num_sgd_iter, - learner_queue_size, learner_queue_timeout): + def __init__(self, local_worker: RolloutWorker, minibatch_buffer_size: int, + num_sgd_iter: int, learner_queue_size: int, + learner_queue_timeout: int): """Initialize the learner thread. Args: @@ -57,14 +60,14 @@ class LearnerThread(threading.Thread): self.stopped = False self.num_steps = 0 - def run(self): + def run(self) -> None: # Switch on eager mode if configured. if self.local_worker.policy_config.get("framework") in ["tf2", "tfe"]: tf1.enable_eager_execution() while not self.stopped: self.step() - def step(self): + def step(self) -> None: with self.queue_timer: batch, _ = self.minibatch_buffer.get() @@ -77,7 +80,7 @@ class LearnerThread(threading.Thread): self.outqueue.put((batch.count, self.stats)) self.learner_queue_size.push(self.inqueue.qsize()) - def add_learner_metrics(self, result): + def add_learner_metrics(self, result: Dict) -> Dict: """Add internal metrics to a trainer result dict.""" def timer_to_ms(timer): diff --git a/rllib/execution/metric_ops.py b/rllib/execution/metric_ops.py index 374ff047d..70ae38e3f 100644 --- a/rllib/execution/metric_ops.py +++ b/rllib/execution/metric_ops.py @@ -1,4 +1,4 @@ -from typing import Any, List +from typing import Any, List, Dict import time from ray.util.iter import LocalIterator @@ -59,9 +59,9 @@ class CollectMetrics: """ def __init__(self, - workers, - min_history=100, - timeout_seconds=180, + workers: WorkerSet, + min_history: int = 100, + timeout_seconds: int = 180, selected_workers: List["ActorHandle"] = None): self.workers = workers self.episode_history = [] @@ -70,7 +70,7 @@ class CollectMetrics: self.timeout_seconds = timeout_seconds self.selected_workers = selected_workers - def __call__(self, _): + def __call__(self, _: Any) -> Dict: # Collect worker metrics. episodes, self.to_be_collected = collect_episodes( self.workers.local_worker(), @@ -124,11 +124,11 @@ class OncePerTimeInterval: 5.00001 # will be greater than 5 seconds """ - def __init__(self, delay): + def __init__(self, delay: int): self.delay = delay self.last_called = 0 - def __call__(self, item): + def __call__(self, item: Any) -> bool: if self.delay <= 0.0: return True now = time.time() @@ -151,11 +151,11 @@ class OncePerTimestepsElapsed: # will only return after 1000 steps have elapsed """ - def __init__(self, delay_steps): + def __init__(self, delay_steps: int): self.delay_steps = delay_steps self.last_called = 0 - def __call__(self, item): + def __call__(self, item: Any) -> bool: if self.delay_steps <= 0: return True metrics = _get_shared_metrics() diff --git a/rllib/execution/minibatch_buffer.py b/rllib/execution/minibatch_buffer.py index 4cd41fcc7..54b5c4a2c 100644 --- a/rllib/execution/minibatch_buffer.py +++ b/rllib/execution/minibatch_buffer.py @@ -1,10 +1,19 @@ +from typing import Any, Tuple +import queue + + class MinibatchBuffer: """Ring buffer of recent data batches for minibatch SGD. This is for use with AsyncSamplesOptimizer. """ - def __init__(self, inqueue, size, timeout, num_passes, init_num_passes=1): + def __init__(self, + inqueue: queue.Queue, + size: int, + timeout: float, + num_passes: int, + init_num_passes: int = 1): """Initialize a minibatch buffer. Args: @@ -23,7 +32,7 @@ class MinibatchBuffer: self.ttl = [0] * size self.idx = 0 - def get(self): + def get(self) -> Tuple[Any, bool]: """Get a new batch from the internal ring buffer. Returns: diff --git a/rllib/execution/multi_gpu_learner.py b/rllib/execution/multi_gpu_learner.py index 4c450c948..f6455dc98 100644 --- a/rllib/execution/multi_gpu_learner.py +++ b/rllib/execution/multi_gpu_learner.py @@ -12,6 +12,7 @@ from ray.rllib.execution.multi_gpu_impl import LocalSyncParallelOptimizer from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.timer import TimerStat +from ray.rllib.evaluation.rollout_worker import RolloutWorker tf1, tf, tfv = try_import_tf() @@ -25,17 +26,17 @@ class TFMultiGPULearner(LearnerThread): """ def __init__(self, - local_worker, - num_gpus=1, - lr=0.0005, - train_batch_size=500, - num_data_loader_buffers=1, - minibatch_buffer_size=1, - num_sgd_iter=1, - learner_queue_size=16, - learner_queue_timeout=300, - num_data_load_threads=16, - _fake_gpus=False): + local_worker: RolloutWorker, + num_gpus: int = 1, + lr: float = 0.0005, + train_batch_size: int = 500, + num_data_loader_buffers: int = 1, + minibatch_buffer_size: int = 1, + num_sgd_iter: int = 1, + learner_queue_size: int = 16, + learner_queue_timeout: int = 300, + num_data_load_threads: int = 16, + _fake_gpus: bool = False): """Initialize a multi-gpu learner thread. Args: @@ -121,7 +122,7 @@ class TFMultiGPULearner(LearnerThread): learner_queue_timeout, num_sgd_iter) @override(LearnerThread) - def step(self): + def step(self) -> None: assert self.loader_thread.is_alive() with self.load_wait_timer: opt, released = self.minibatch_buffer.get() @@ -139,7 +140,7 @@ class TFMultiGPULearner(LearnerThread): class _LoaderThread(threading.Thread): - def __init__(self, learner, share_stats): + def __init__(self, learner: LearnerThread, share_stats: bool): threading.Thread.__init__(self) self.learner = learner self.daemon = True @@ -150,11 +151,11 @@ class _LoaderThread(threading.Thread): self.queue_timer = TimerStat() self.load_timer = TimerStat() - def run(self): + def run(self) -> None: while True: self._step() - def _step(self): + def _step(self) -> None: s = self.learner with self.queue_timer: batch = s.inqueue.get() diff --git a/rllib/execution/replay_buffer.py b/rllib/execution/replay_buffer.py index 8d98e1397..79f4882eb 100644 --- a/rllib/execution/replay_buffer.py +++ b/rllib/execution/replay_buffer.py @@ -3,7 +3,7 @@ import logging import numpy as np import platform import random -from typing import List +from typing import List, Dict # Import ray before psutil will make sure we use psutil's bundled version import ray # noqa F401 @@ -64,11 +64,11 @@ class ReplayBuffer: self._evicted_hit_stats = WindowStat("evicted_hit", 1000) self._est_size_bytes = 0 - def __len__(self): + def __len__(self) -> int: return len(self._storage) @DeveloperAPI - def add(self, item: SampleBatchType, weight: float): + def add(self, item: SampleBatchType, weight: float) -> None: warn_replay_buffer_size( item=item, num_items=self._maxsize / item.count) assert item.count > 0, item @@ -116,7 +116,7 @@ class ReplayBuffer: return self._encode_sample(idxes) @DeveloperAPI - def stats(self, debug=False): + def stats(self, debug=False) -> dict: data = { "added_count": self._num_timesteps_added, "sampled_count": self._num_timesteps_sampled, @@ -156,7 +156,7 @@ class PrioritizedReplayBuffer(ReplayBuffer): self._prio_change_stats = WindowStat("reprio", 1000) @DeveloperAPI - def add(self, item: SampleBatchType, weight: float): + def add(self, item: SampleBatchType, weight: float) -> None: idx = self._next_idx super(PrioritizedReplayBuffer, self).add(item, weight) if weight is None: @@ -164,7 +164,7 @@ class PrioritizedReplayBuffer(ReplayBuffer): self._it_sum[idx] = weight**self._alpha self._it_min[idx] = weight**self._alpha - def _sample_proportional(self, num_items: int): + def _sample_proportional(self, num_items: int) -> List[int]: res = [] for _ in range(num_items): # TODO(szymon): should we ensure no repeats? @@ -215,7 +215,8 @@ class PrioritizedReplayBuffer(ReplayBuffer): return batch @DeveloperAPI - def update_priorities(self, idxes, priorities): + def update_priorities(self, idxes: List[int], + priorities: List[float]) -> None: """Update priorities of sampled transitions. sets priority of transition at index idxes[i] in buffer @@ -242,7 +243,7 @@ class PrioritizedReplayBuffer(ReplayBuffer): self._max_priority = max(self._max_priority, priority) @DeveloperAPI - def stats(self, debug=False): + def stats(self, debug: bool = False) -> Dict: parent = ReplayBuffer.stats(self, debug) if debug: parent.update(self._prio_change_stats.stats()) @@ -260,15 +261,15 @@ class LocalReplayBuffer(ParallelIteratorWorker): may be created to increase parallelism.""" def __init__(self, - num_shards=1, - learning_starts=1000, - buffer_size=10000, - replay_batch_size=1, - prioritized_replay_alpha=0.6, - prioritized_replay_beta=0.4, - prioritized_replay_eps=1e-6, - replay_mode="independent", - replay_sequence_length=1): + num_shards: int = 1, + learning_starts: int = 1000, + buffer_size: int = 10000, + replay_batch_size: int = 1, + prioritized_replay_alpha: float = 0.6, + prioritized_replay_beta: float = 0.4, + prioritized_replay_eps: float = 1e-6, + replay_mode: str = "independent", + replay_sequence_length: int = 1): self.replay_starts = learning_starts // num_shards self.buffer_size = buffer_size // num_shards self.replay_batch_size = replay_batch_size @@ -318,10 +319,10 @@ class LocalReplayBuffer(ParallelIteratorWorker): global _local_replay_buffer return _local_replay_buffer - def get_host(self): + def get_host(self) -> str: return platform.node() - def add_batch(self, batch): + def add_batch(self, batch: SampleBatchType) -> None: # Make a copy so the replay buffer doesn't pin plasma memory. batch = batch.copy() # Handle everything as if multiagent @@ -342,7 +343,7 @@ class LocalReplayBuffer(ParallelIteratorWorker): self.replay_buffers[policy_id].add(s, weight=weight) self.num_added += batch.count - def replay(self): + def replay(self) -> SampleBatchType: if self._fake_batch: fake_batch = SampleBatch(self._fake_batch) return MultiAgentBatch({ @@ -364,7 +365,7 @@ class LocalReplayBuffer(ParallelIteratorWorker): beta=self.prioritized_replay_beta) return MultiAgentBatch(samples, self.replay_batch_size) - def update_priorities(self, prio_dict): + def update_priorities(self, prio_dict: Dict) -> None: with self.update_priorities_timer: for policy_id, (batch_indexes, td_errors) in prio_dict.items(): new_priorities = ( @@ -372,7 +373,7 @@ class LocalReplayBuffer(ParallelIteratorWorker): self.replay_buffers[policy_id].update_priorities( batch_indexes, new_priorities) - def stats(self, debug=False): + def stats(self, debug: bool = False) -> Dict: stat = { "add_batch_time_ms": round(1000 * self.add_batch_timer.mean, 3), "replay_time_ms": round(1000 * self.replay_timer.mean, 3), diff --git a/rllib/execution/replay_ops.py b/rllib/execution/replay_ops.py index 9ed25e9e9..7bfc3a1b9 100644 --- a/rllib/execution/replay_ops.py +++ b/rllib/execution/replay_ops.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Any, Optional import random from ray.util.iter import from_actors, LocalIterator, _NextValueNotReady @@ -55,7 +55,7 @@ class StoreToReplayBuffer: def Replay(*, local_buffer: LocalReplayBuffer = None, actors: List["ActorHandle"] = None, - num_async=4): + num_async: int = 4) -> LocalIterator[SampleBatchType]: """Replay experiences from the given buffer or actors. This should be combined with the StoreToReplayActors operation using the @@ -99,10 +99,10 @@ def Replay(*, class WaitUntilTimestepsElapsed: """Callable that returns True once a given number of timesteps are hit.""" - def __init__(self, target_num_timesteps): + def __init__(self, target_num_timesteps: int): self.target_num_timesteps = target_num_timesteps - def __call__(self, item): + def __call__(self, item: Any) -> bool: metrics = _get_shared_metrics() ts = metrics.counters[STEPS_SAMPLED_COUNTER] return ts > self.target_num_timesteps @@ -112,7 +112,9 @@ class WaitUntilTimestepsElapsed: class SimpleReplayBuffer: """Simple replay buffer that operates over batches.""" - def __init__(self, num_slots, replay_proportion: float = None): + def __init__(self, + num_slots: int, + replay_proportion: Optional[float] = None): """Initialize SimpleReplayBuffer. Args: @@ -122,7 +124,7 @@ class SimpleReplayBuffer: self.replay_batches = [] self.replay_index = 0 - def add_batch(self, sample_batch): + def add_batch(self, sample_batch: SampleBatchType) -> None: warn_replay_buffer_size(item=sample_batch, num_items=self.num_slots) if self.num_slots > 0: if len(self.replay_batches) < self.num_slots: @@ -132,7 +134,7 @@ class SimpleReplayBuffer: self.replay_index += 1 self.replay_index %= self.num_slots - def replay(self): + def replay(self) -> SampleBatchType: return random.choice(self.replay_batches) @@ -145,7 +147,7 @@ class MixInReplay: number of replay slots. """ - def __init__(self, num_slots, replay_proportion: float): + def __init__(self, num_slots: int, replay_proportion: float): """Initialize MixInReplay. Args: @@ -171,7 +173,7 @@ class MixInReplay: self.replay_buffer = SimpleReplayBuffer(num_slots) self.replay_proportion = replay_proportion - def __call__(self, sample_batch): + def __call__(self, sample_batch: SampleBatchType) -> List[SampleBatchType]: # Put in replay buffer if enabled. self.replay_buffer.add_batch(sample_batch) diff --git a/rllib/execution/segment_tree.py b/rllib/execution/segment_tree.py index e436f3a5a..ead3e3188 100644 --- a/rllib/execution/segment_tree.py +++ b/rllib/execution/segment_tree.py @@ -1,4 +1,5 @@ import operator +from typing import Any, Optional class SegmentTree: @@ -28,7 +29,10 @@ class SegmentTree: `tree[0]` accesses `internal_array[4]` in the above example. """ - def __init__(self, capacity, operation, neutral_element=None): + def __init__(self, + capacity: int, + operation: Any, + neutral_element: Optional[Any] = None): """Initializes a Segment Tree object. Args: @@ -52,7 +56,7 @@ class SegmentTree: self.value = [self.neutral_element for _ in range(2 * capacity)] self.operation = operation - def reduce(self, start=0, end=None): + def reduce(self, start: int = 0, end: Optional[int] = None) -> Any: """Applies `self.operation` to subsequence of our values. Subsequence is contiguous, includes `start` and excludes `end`. @@ -122,7 +126,7 @@ class SegmentTree: return result - def __setitem__(self, idx, val): + def __setitem__(self, idx: int, val: float) -> None: """ Inserts/overwrites a value in/into the tree. @@ -147,7 +151,7 @@ class SegmentTree: self.value[update_idx + 1]) idx = idx >> 1 # Divide by 2 (faster than division). - def __getitem__(self, idx): + def __getitem__(self, idx: int) -> Any: assert 0 <= idx < self.capacity return self.value[idx + self.capacity] @@ -155,15 +159,15 @@ class SegmentTree: class SumSegmentTree(SegmentTree): """A SegmentTree with the reduction `operation`=operator.add.""" - def __init__(self, capacity): + def __init__(self, capacity: int): super(SumSegmentTree, self).__init__( capacity=capacity, operation=operator.add) - def sum(self, start=0, end=None): + def sum(self, start: int = 0, end: Optional[Any] = None) -> Any: """Returns the sum over a sub-segment of the tree.""" return self.reduce(start, end) - def find_prefixsum_idx(self, prefixsum): + def find_prefixsum_idx(self, prefixsum: float) -> int: """Finds highest i, for which: sum(arr[0]+..+arr[i - i]) <= prefixsum. Args: @@ -188,9 +192,9 @@ class SumSegmentTree(SegmentTree): class MinSegmentTree(SegmentTree): - def __init__(self, capacity): + def __init__(self, capacity: int): super(MinSegmentTree, self).__init__(capacity=capacity, operation=min) - def min(self, start=0, end=None): + def min(self, start: int = 0, end: Optional[Any] = None) -> Any: """Returns min(arr[start], ..., arr[end])""" return self.reduce(start, end) diff --git a/rllib/execution/train_ops.py b/rllib/execution/train_ops.py index f99fd0a0c..e2411ed32 100644 --- a/rllib/execution/train_ops.py +++ b/rllib/execution/train_ops.py @@ -2,7 +2,7 @@ from collections import defaultdict import logging import numpy as np import math -from typing import List +from typing import List, Tuple, Any import ray from ray.rllib.evaluation.metrics import get_learner_stats, LEARNER_STATS_KEY @@ -18,7 +18,7 @@ from ray.rllib.policy.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \ MultiAgentBatch from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.sgd import do_minibatch_sgd, averaged -from ray.rllib.utils.typing import PolicyID, SampleBatchType +from ray.rllib.utils.typing import PolicyID, SampleBatchType, ModelGradients tf1, tf, tfv = try_import_tf() @@ -242,10 +242,10 @@ class ComputeGradients: Updates the LEARNER_INFO info field in the local iterator context. """ - def __init__(self, workers): + def __init__(self, workers: WorkerSet): self.workers = workers - def __call__(self, samples: SampleBatchType): + def __call__(self, samples: SampleBatchType) -> Tuple[ModelGradients, int]: _check_sample_batch_type(samples) metrics = _get_shared_metrics() with metrics.timers[COMPUTE_GRADS_TIMER]: @@ -283,7 +283,7 @@ class ApplyGradients: self.policies = policies or workers.local_worker().policies_to_train self.update_all = update_all - def __call__(self, item): + def __call__(self, item: Tuple[ModelGradients, int]) -> None: if not isinstance(item, tuple) or len(item) != 2: raise ValueError( "Input must be a tuple of (grad_dict, count), got {}".format( @@ -333,7 +333,8 @@ class AverageGradients: {"var_0": ..., ...}, 1600 # averaged grads, summed batch count """ - def __call__(self, gradients): + def __call__(self, gradients: List[Tuple[ModelGradients, int]] + ) -> Tuple[ModelGradients, int]: acc = None sum_count = 0 for grad, count in gradients: @@ -366,10 +367,10 @@ class UpdateTargetNetwork: """ def __init__(self, - workers, - target_update_freq, - by_steps_trained=False, - policies=frozenset([])): + workers: WorkerSet, + target_update_freq: int, + by_steps_trained: bool = False, + policies: List[PolicyID] = frozenset([])): self.workers = workers self.target_update_freq = target_update_freq self.policies = (policies or workers.local_worker().policies_to_train) @@ -378,7 +379,7 @@ class UpdateTargetNetwork: else: self.metric = STEPS_SAMPLED_COUNTER - def __call__(self, _): + def __call__(self, _: Any) -> None: metrics = _get_shared_metrics() cur_ts = metrics.counters[self.metric] last_update = metrics.counters[LAST_TARGET_UPDATE_TS] diff --git a/rllib/execution/tree_agg.py b/rllib/execution/tree_agg.py index 69e06a4b0..b04bee783 100644 --- a/rllib/execution/tree_agg.py +++ b/rllib/execution/tree_agg.py @@ -1,6 +1,6 @@ import logging import platform -from typing import List +from typing import List, Dict, Any import ray from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER, \ @@ -9,8 +9,9 @@ from ray.rllib.execution.replay_ops import MixInReplay from ray.rllib.execution.rollout_ops import ParallelRollouts, ConcatBatches from ray.rllib.utils.actors import create_colocated from ray.util.iter import ParallelIterator, ParallelIteratorWorker, \ - from_actors -from ray.rllib.utils.typing import SampleBatchType + from_actors, LocalIterator +from ray.rllib.utils.typing import SampleBatchType, ModelWeights +from ray.rllib.evaluation.worker_set import WorkerSet logger = logging.getLogger(__name__) @@ -25,7 +26,7 @@ class Aggregator(ParallelIteratorWorker): work to be offloaded to these actors instead of run in the learner. """ - def __init__(self, config: dict, + def __init__(self, config: Dict, rollout_group: "ParallelIterator[SampleBatchType]"): self.weights = None self.global_vars = None @@ -60,15 +61,16 @@ class Aggregator(ParallelIteratorWorker): super().__init__(generator, repeat=False) - def get_host(self): + def get_host(self) -> str: return platform.node() - def set_weights(self, weights, global_vars): + def set_weights(self, weights: ModelWeights, global_vars: Dict) -> None: self.weights = weights self.global_vars = global_vars -def gather_experiences_tree_aggregation(workers, config): +def gather_experiences_tree_aggregation(workers: WorkerSet, + config: Dict) -> "LocalIterator[Any]": """Tree aggregation version of gather_experiences_directly().""" rollouts = ParallelRollouts(workers, mode="raw")