chore: import upstream snapshot with attribution

rllib/algorithms/appo/utils.py

@@ -0,0 +1,264 @@
+"""
+[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