| """Replay buffer with offline-protected FIFO eviction. |
| |
| Stores observation-action |
| windows of fixed length ``seq_len``. Offline data is pinned at the front |
| and never evicted; online samples use FIFO. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
|
|
| class ReplayBuffer: |
| """Fixed-capacity buffer with offline-protected FIFO eviction. |
| |
| Offline samples (loaded once via ``load_offline_data``) are pinned |
| and never evicted. Online samples added via ``add`` are FIFO-evicted |
| when the total count exceeds ``capacity``. |
| |
| Args: |
| capacity: Maximum total number of windows. |
| seq_len: Action-sequence window length. |
| pad_token: Token used to pad short sequences. |
| """ |
|
|
| def __init__( |
| self, |
| capacity: int, |
| seq_len: int, |
| pad_token: int, |
| ) -> None: |
| self._capacity = capacity |
| self._seq_len = seq_len |
| self._pad_token = pad_token |
|
|
| |
| self._offline: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] |
| self._online: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] |
|
|
| |
| self._cache_valid = False |
| self._cached_local: np.ndarray | None = None |
| self._cached_global: np.ndarray | None = None |
| self._cached_actions: np.ndarray | None = None |
|
|
| def load_offline_data( |
| self, |
| data: dict, |
| allowed_envs: list[str], |
| ) -> None: |
| """Load pre-collected trajectories and slice into windows. |
| |
| Expects the dict format ``{"trajectories": [...]}`` where each |
| entry is a dict with ``"local"``, ``"global"``, ``"actions"``, |
| ``"env_id"``. |
| |
| Args: |
| data: Dataset dict. |
| allowed_envs: Only samples from these env IDs are kept. |
| """ |
| if not isinstance(data, dict): |
| raise TypeError( |
| f"Offline dataset must be a dict, got {type(data).__name__}; " |
| "the legacy list format is no longer supported." |
| ) |
|
|
| trajectories = data.get("trajectories", [data]) |
| for traj in trajectories: |
| if traj.get("env_id", "") not in allowed_envs: |
| continue |
| windows = self._slice_trajectory(traj) |
| self._offline.extend(windows) |
| |
| if len(self._offline) > self._capacity: |
| self._offline = self._offline[: self._capacity] |
| self._invalidate_cache() |
|
|
| def _invalidate_cache(self) -> None: |
| """Mark the stacked array cache as stale.""" |
| self._cache_valid = False |
|
|
| def _ensure_cache(self) -> None: |
| """Rebuild stacked arrays from offline + online windows.""" |
| if self._cache_valid: |
| return |
| combined = self._offline + self._online |
| if not combined: |
| return |
| n = len(combined) |
| l0, g0, a0 = combined[0] |
| self._cached_local = np.empty( |
| (n, *l0.shape), |
| dtype=l0.dtype, |
| ) |
| self._cached_global = np.empty( |
| (n, *g0.shape), |
| dtype=g0.dtype, |
| ) |
| self._cached_actions = np.empty( |
| (n, *a0.shape), |
| dtype=a0.dtype, |
| ) |
| for i, (loc, glob, act) in enumerate(combined): |
| self._cached_local[i] = loc |
| self._cached_global[i] = glob |
| self._cached_actions[i] = act |
| self._cache_valid = True |
|
|
| def add(self, trajectory: dict) -> None: |
| """Add a trajectory, sliced into overlapping windows. |
| |
| FIFO-evicts oldest online samples when over capacity. |
| |
| Args: |
| trajectory: Dict with ``"local"`` ``[T,9,9]``, |
| ``"global"`` ``[T,21,79]``, ``"actions"`` ``[T]``. |
| """ |
| windows = self._slice_trajectory(trajectory) |
| self._online.extend(windows) |
| max_online = self._capacity - len(self._offline) |
| if len(self._online) > max_online: |
| excess = len(self._online) - max_online |
| self._online = self._online[excess:] |
| self._invalidate_cache() |
|
|
| def sample( |
| self, |
| batch_size: int, |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None: |
| """Random sample from offline + online combined. |
| |
| Args: |
| batch_size: Number of windows to sample. |
| |
| Returns: |
| ``(local [B,9,9], global [B,21,79], actions [B,seq_len])`` |
| as numpy arrays, or ``None`` if the buffer is empty. |
| """ |
| if len(self) == 0: |
| return None |
| self._ensure_cache() |
| if self._cached_local is None: |
| return None |
| indices = np.random.randint(0, len(self), size=batch_size) |
| return ( |
| self._cached_local[indices], |
| self._cached_global[indices], |
| self._cached_actions[indices], |
| ) |
|
|
| def __len__(self) -> int: |
| """Total number of windows (offline + online).""" |
| return len(self._offline) + len(self._online) |
|
|
| @property |
| def n_offline(self) -> int: |
| """Number of pinned offline windows.""" |
| return len(self._offline) |
|
|
| @property |
| def offline_size(self) -> int: |
| """Number of pinned offline windows (alias).""" |
| return len(self._offline) |
|
|
| def _slice_trajectory( |
| self, |
| traj: dict, |
| ) -> list[tuple[np.ndarray, np.ndarray, np.ndarray]]: |
| """Slice a trajectory into overlapping seq_len windows. |
| |
| Args: |
| traj: Trajectory dict with ``"local"``, ``"global"``, |
| ``"actions"`` arrays. |
| |
| Returns: |
| List of ``(local, global, actions)`` tuples. |
| """ |
| local_arr = np.asarray(traj["local"]) |
| global_arr = np.asarray(traj["global"]) |
| actions_arr = np.asarray(traj["actions"]) |
| T = len(actions_arr) |
| windows: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] |
|
|
| for start in range(T): |
| end = start + self._seq_len |
| if end <= T: |
| a = actions_arr[start:end] |
| else: |
| a = np.full(self._seq_len, self._pad_token, dtype=np.int64) |
| a[: T - start] = actions_arr[start:] |
|
|
| |
| loc = local_arr[min(start, len(local_arr) - 1)] |
| glob = global_arr[min(start, len(global_arr) - 1)] |
| windows.append((loc.copy(), glob.copy(), a)) |
|
|
| return windows |
|
|
