image-match/python/py/Lib/site-packages/ray/rllib/algorithms/appo/utils.py

"""
[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