wehub-resource-sync
265 lines
9.1 KiB
Python
265 lines
9.1 KiB
Python
"""
|
|
[1] IMPACT: Importance Weighted Asynchronous Architectures with Clipped Target Networks.
|
|
Luo et al. 2020
|
|
https://arxiv.org/pdf/1912.00167
|
|
"""
|
|
import threading
|
|
import time
|
|
from collections import deque
|
|
from typing import Any, Optional
|
|
|
|
import numpy as np
|
|
|
|
from ray.rllib.models.catalog import ModelCatalog
|
|
from ray.rllib.models.modelv2 import ModelV2
|
|
from ray.rllib.utils.annotations import OldAPIStack
|
|
from ray.rllib.utils.metrics.ray_metrics import (
|
|
DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
|
TimerAndPrometheusLogger,
|
|
)
|
|
from ray.util.metrics import Counter, Histogram
|
|
|
|
POLICY_SCOPE = "func"
|
|
TARGET_POLICY_SCOPE = "target_func"
|
|
|
|
|
|
class CircularBuffer:
|
|
"""A circular batch-wise buffer with Queue-like interface.
|
|
|
|
The buffer holds at most N batches, which are sampled at random (uniformly).
|
|
If full and a new batch is added, the oldest batch is discarded. Each batch
|
|
can be sampled at most K times (after which it is also discarded).
|
|
|
|
This version implements Queue-like put/get methods with blocking support.
|
|
"""
|
|
|
|
def __init__(self, num_batches: int, iterations_per_batch: int):
|
|
"""
|
|
Args:
|
|
num_batches: N from the paper (queue buffer size).
|
|
iterations_per_batch: K ("replay coefficient") from the paper. Defines
|
|
how often a single batch can sampled before being discarded. If a
|
|
new batch is added when the buffer is full, the oldest batch is
|
|
discarded entirely (regardless of how often it has been sampled).
|
|
"""
|
|
self.num_batches = num_batches
|
|
self.iterations_per_batch = iterations_per_batch
|
|
|
|
self._NxK = self.num_batches * self.iterations_per_batch
|
|
self._num_added = 0
|
|
|
|
self._buffer = deque([None for _ in range(self._NxK)], maxlen=self._NxK)
|
|
self._indices = set()
|
|
self._offset = self._NxK
|
|
self._lock = threading.Lock()
|
|
|
|
# Semaphore tracks the number of *available* samples.
|
|
self._items_available = threading.Semaphore(0)
|
|
|
|
self._rng = np.random.default_rng()
|
|
|
|
# Statistics
|
|
self._total_puts = 0
|
|
self._total_gets = 0
|
|
self._total_dropped = 0
|
|
|
|
# Ray metrics
|
|
self._metrics_circular_buffer_put_time = Histogram(
|
|
name="rllib_utils_circular_buffer_put_time",
|
|
description="Time spent in CircularBuffer.put()",
|
|
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
|
tag_keys=("rllib",),
|
|
)
|
|
self._metrics_circular_buffer_put_time.set_default_tags(
|
|
{"rllib": self.__class__.__name__}
|
|
)
|
|
|
|
self._metrics_circular_buffer_put_ts_dropped = Counter(
|
|
name="rllib_utils_circular_buffer_put_ts_dropped_counter",
|
|
description="Total number of env steps dropped by the CircularBuffer.",
|
|
tag_keys=("rllib",),
|
|
)
|
|
self._metrics_circular_buffer_put_ts_dropped.set_default_tags(
|
|
{"rllib": self.__class__.__name__}
|
|
)
|
|
|
|
self._metrics_circular_buffer_get_time = Histogram(
|
|
name="rllib_utils_circular_buffer_get_time",
|
|
description="Time spent in CircularBuffer.get()",
|
|
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
|
tag_keys=("rllib",),
|
|
)
|
|
self._metrics_circular_buffer_get_time.set_default_tags(
|
|
{"rllib": self.__class__.__name__}
|
|
)
|
|
|
|
def put(
|
|
self, item: Any, block: bool = True, timeout: Optional[float] = None
|
|
) -> int:
|
|
"""Add a new batch to the buffer.
|
|
|
|
The batch is added K times (iterations_per_batch) to allow for K samples.
|
|
If full, the oldest batch entries are dropped.
|
|
|
|
Args:
|
|
item: The batch to add
|
|
block: Not used (always non-blocking for puts)
|
|
timeout: Not used
|
|
|
|
Returns:
|
|
Number of dropped entries (0 or iterations_per_batch)
|
|
"""
|
|
with TimerAndPrometheusLogger(self._metrics_circular_buffer_put_time):
|
|
with self._lock:
|
|
self._total_puts += 1
|
|
|
|
# Check if we'll drop old entries
|
|
dropped_entry = self._buffer[0]
|
|
|
|
# Add buffer K times with new indices
|
|
for _ in range(self.iterations_per_batch):
|
|
self._buffer.append(item)
|
|
self._indices.add(self._offset)
|
|
self._indices.discard(self._offset - self._NxK)
|
|
self._offset += 1
|
|
|
|
# Release semaphore for each available sample
|
|
self._items_available.release()
|
|
|
|
self._num_added += 1
|
|
|
|
# A valid entry (w/ a batch whose k has not been reach K yet) was dropped.
|
|
dropped_ts = 0
|
|
if dropped_entry is not None:
|
|
dropped_ts = (
|
|
dropped_entry[0].env_steps()
|
|
if isinstance(dropped_entry, tuple)
|
|
else dropped_entry.env_steps()
|
|
)
|
|
if dropped_ts > 0:
|
|
self._metrics_circular_buffer_put_ts_dropped.inc(
|
|
value=dropped_ts
|
|
)
|
|
|
|
return dropped_ts
|
|
|
|
def put_nowait(self, item: Any) -> int:
|
|
"""Equivalent to self.put(block=False)."""
|
|
return self.put(item, block=False)
|
|
|
|
def get(self, block: bool = True, timeout: Optional[float] = None) -> Any:
|
|
"""Sample a random batch from the buffer.
|
|
|
|
The sampled entry is removed and won't be sampled again.
|
|
Blocks if the buffer is empty (when block=True).
|
|
|
|
Args:
|
|
block: If True, block until an item is available
|
|
timeout: Maximum time to wait (only used when block=True)
|
|
|
|
Returns:
|
|
A randomly sampled batch
|
|
|
|
Raises:
|
|
TimeoutError: If timeout expires while blocking
|
|
IndexError: If buffer is empty and block=False
|
|
"""
|
|
# Only initially, the buffer may be empty -> Just wait for some time.
|
|
with TimerAndPrometheusLogger(self._metrics_circular_buffer_get_time):
|
|
while len(self) == 0:
|
|
time.sleep(0.0001)
|
|
|
|
# Sample a random buffer index.
|
|
with self._lock:
|
|
idx = self._rng.choice(list(self._indices))
|
|
actual_buffer_idx = idx - self._offset + self._NxK
|
|
batch = self._buffer[actual_buffer_idx]
|
|
assert batch is not None, (
|
|
idx,
|
|
actual_buffer_idx,
|
|
self._offset,
|
|
self._indices,
|
|
[b is None for b in self._buffer],
|
|
)
|
|
self._buffer[actual_buffer_idx] = None
|
|
self._indices.discard(idx)
|
|
|
|
# Return the sampled batch.
|
|
return batch
|
|
|
|
def get_nowait(self) -> Any:
|
|
"""Equivalent to self.get(block=False)."""
|
|
return self.get(block=False)
|
|
|
|
@property
|
|
def filled(self) -> bool:
|
|
"""Whether the buffer has been filled once with at least `self.num_batches`."""
|
|
with self._lock:
|
|
return self._num_added >= self.num_batches
|
|
|
|
def qsize(self) -> int:
|
|
"""Returns the number of actually valid (non-expired) batches in the buffer."""
|
|
with self._lock:
|
|
return len(self._indices)
|
|
|
|
def __len__(self) -> int:
|
|
return self.qsize()
|
|
|
|
def task_done(self):
|
|
"""No-op for Queue compatibility."""
|
|
pass
|
|
|
|
def get_stats(self) -> dict:
|
|
"""Get buffer statistics for monitoring."""
|
|
with self._lock:
|
|
return {
|
|
"size": len(self._indices),
|
|
"capacity": self._NxK,
|
|
"num_batches": self.num_batches,
|
|
"iterations_per_batch": self.iterations_per_batch,
|
|
"total_puts": self._total_puts,
|
|
"total_gets": self._total_gets,
|
|
"total_dropped": self._total_dropped,
|
|
"filled": self._num_added >= self.num_batches,
|
|
}
|
|
|
|
|
|
@OldAPIStack
|
|
def make_appo_models(policy) -> ModelV2:
|
|
"""Builds model and target model for APPO.
|
|
|
|
Returns:
|
|
ModelV2: The Model for the Policy to use.
|
|
Note: The target model will not be returned, just assigned to
|
|
`policy.target_model`.
|
|
"""
|
|
# Get the num_outputs for the following model construction calls.
|
|
_, logit_dim = ModelCatalog.get_action_dist(
|
|
policy.action_space, policy.config["model"]
|
|
)
|
|
|
|
# Construct the (main) model.
|
|
policy.model = ModelCatalog.get_model_v2(
|
|
policy.observation_space,
|
|
policy.action_space,
|
|
logit_dim,
|
|
policy.config["model"],
|
|
name=POLICY_SCOPE,
|
|
framework=policy.framework,
|
|
)
|
|
policy.model_variables = policy.model.variables()
|
|
|
|
# Construct the target model.
|
|
policy.target_model = ModelCatalog.get_model_v2(
|
|
policy.observation_space,
|
|
policy.action_space,
|
|
logit_dim,
|
|
policy.config["model"],
|
|
name=TARGET_POLICY_SCOPE,
|
|
framework=policy.framework,
|
|
)
|
|
policy.target_model_variables = policy.target_model.variables()
|
|
|
|
# Return only the model (not the target model).
|
|
return policy.model
|