feat: support staleness-window in ReplayBufferNew (#2458)

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Author: yuki-97

nemo_rl/algorithms/async_utils/__init__.py

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nemo_rl.algorithms.async_utils.replay_buffer import ReplayBuffer, ReplayBufferNew
+from nemo_rl.algorithms.async_utils.replay_buffer import ReplayBuffer
 from nemo_rl.algorithms.async_utils.trajectory_collector import AsyncTrajectoryCollector
 
 __all__ = [
     "ReplayBuffer",
-    "ReplayBufferNew",
     "AsyncTrajectoryCollector",
 ]

nemo_rl/algorithms/async_utils/interfaces.py

@@ -51,10 +51,6 @@ def sample(
         """
         ...
 
-    def evict(self) -> None:
-        """Evict old trajectories."""
-        ...
-
     def size(self) -> int:
         """Return current buffer size."""
         ...

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 import threading as _threading
-from typing import Any, Optional
+from collections import Counter
+from typing import Any, Iterable, Optional
 
 import ray
 
@@ -87,6 +88,13 @@ def get_existing_target_weights(self) -> set[int]:
         with self._lock:
             return set(self.target_weight_versions)
 
+    def _remove_indices(self, indices: Iterable[int]) -> None:
+        """Remove trajectories at the given indices."""
+        for idx in sorted(indices, reverse=True):
+            self.trajectory_versions.pop(idx)
+            self.target_weight_versions.pop(idx)
+            self.trajectories.pop(idx)
+
     def sample(
         self,
         num_prompt_groups: int,
@@ -113,8 +121,6 @@ def sample(
             print(f"   {self.trajectory_versions=}")
 
             # For debugging: check for unexpected old trajectories
-            from collections import Counter
-
             version_counts = Counter(self.trajectory_versions)
             print(f"   {version_counts=}")
 
@@ -180,8 +186,6 @@ def sample(
                 f"   ✅ Selected {len(selected)} trajectories all intended for step {current_weight_version}"
             )
 
-            from collections import Counter
-
             sampled_weights = [self.trajectory_versions[i] for i in selected]
             avg_trajectory_age = current_weight_version - sum(sampled_weights) / len(
                 sampled_weights
@@ -194,13 +198,9 @@ def sample(
                 f"🎯 All selected trajectories target step {current_weight_version} (100% target match)"
             )
 
-            sampled_items = [self.trajectories[i] for i in selected]
-
             # Remove selected items in reverse order to maintain correct indices
-            for idx in sorted(selected, reverse=True):
-                self.trajectory_versions.pop(idx)
-                self.target_weight_versions.pop(idx)
-                self.trajectories.pop(idx)
+            sampled_items = [self.trajectories[i] for i in selected]
+            self._remove_indices(selected)
             print(
                 f"🗑️ Consumed and removed {len(selected)} groups from buffer, old buffer size: {total_trajectories}, new buffer size: {len(self.trajectories)}, new target weight versions {self.target_weight_versions}"
             )
@@ -210,11 +210,6 @@ def sample(
                 "avg_trajectory_age": avg_trajectory_age,
             }
 
-    def evict(self) -> None:
-        """Evict old trajectories."""
-        # Adding for backward compatibility.
-        pass
-
     def size(self) -> int:
         """Return current buffer size."""
         with self._lock:
@@ -233,6 +228,93 @@ class ReplayBuffer(ReplayBufferImpl):
     pass
 
 
+# WIP: DO NOT USE - This class is WIP and may be changed without notice, please DO NOT USE it.
+# Will be replaced by TQReplayBuffer once TQ is ready.
 @ray.remote  # pragma: no cover
 class ReplayBufferNew(ReplayBufferImpl):
-    pass
+    """Staleness-window replay buffer.
+
+    -- WIP: DO NOT USE --
+    This class is WIP and may be changed without notice, please DO NOT USE it.
+
+    Differences from ReplayBuffer:
+    - _evict(): Stale rows (trainer_version - weight_version > max_staleness) are evicted
+      at the start of every sample() call.
+    - sample(): selects trajectories in freshest-first order (default) or FIFO order,
+      controlled by the sample_freshest_first flag, from whatever remains in the buffer
+      after eviction.
+
+    TODO: remove when cleaning up
+    - max_age_steps won't be used in ReplayBufferNew;
+    - self.target_weight_versions won't be used in ReplayBufferNew and will be removed
+      when cleaning up. target_weight_versions gates generation on specific trainer steps,
+      which causes generation pauses; ReplayBufferNew intentionally avoids this.
+    - add this class to nemo_rl/algorithms/async_utils/__init__.py
+    """
+
+    def __init__(
+        self, max_size: int, max_staleness: int, sample_freshest_first: bool = True
+    ):
+        super().__init__(max_size)
+        if max_staleness < 0:
+            raise ValueError(f"max_staleness must be non-negative, got {max_staleness}")
+        self.max_staleness = max_staleness
+        # will move to StalenessSampler when we implement it
+        self.sample_freshest_first = sample_freshest_first
+
+    def _evict(self, current_weight_version: int) -> None:
+        """Evict rows where trainer_version - weight_version > max_staleness.
+
+        Must be called with self._lock held.
+        """
+        min_valid = current_weight_version - self.max_staleness
+        stale = [i for i, v in enumerate(self.trajectory_versions) if v < min_valid]
+        self._remove_indices(stale)
+
+    def sample(
+        self,
+        num_prompt_groups: int,
+        current_weight_version: int,
+        max_age_steps: int,
+    ) -> Optional[dict[str, Any]]:
+        """Sample num_prompt_groups trajectories, freshest-first.
+
+        Will evict stale rows before sampling, so we will get [current_weight_version - self.max_staleness, current_weight_version] valid trajectories.
+
+        Returns:
+            Dictionary with 'trajectories' and 'avg_trajectory_age' keys, or None.
+        """
+        with self._lock:
+            self._evict(current_weight_version)
+
+            if not self.trajectories:
+                return None
+
+            all_indices = range(len(self.trajectory_versions))
+            if self.sample_freshest_first:
+                all_indices = sorted(
+                    all_indices,
+                    key=lambda i: self.trajectory_versions[i],
+                    reverse=True,
+                )
+
+            if len(all_indices) < num_prompt_groups:
+                print(
+                    f"Insufficient trajectories: have {len(all_indices)}, "
+                    f"need {num_prompt_groups}. Waiting."
+                )
+                return None
+
+            selected = all_indices[:num_prompt_groups]
+            sampled_weights = [self.trajectory_versions[i] for i in selected]
+            avg_trajectory_age = current_weight_version - sum(sampled_weights) / len(
+                sampled_weights
+            )
+
+            sampled_items = [self.trajectories[i] for i in selected]
+            self._remove_indices(selected)
+
+            return {
+                "trajectories": sampled_items,
+                "avg_trajectory_age": avg_trajectory_age,
+            }

tests/unit/algorithms/test_async_utils.py

@@ -28,7 +28,11 @@
 os.environ["RAY_TMPDIR"] = _temp_dir  # Alternative env var
 os.environ["TMPDIR"] = _temp_dir  # System temp dir
 
-from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer
+from nemo_rl.algorithms.async_utils import (
+    AsyncTrajectoryCollector,
+    ReplayBuffer,
+)
+from nemo_rl.algorithms.async_utils.replay_buffer import ReplayBufferNew
 from nemo_rl.algorithms.grpo import MasterConfig
 from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
@@ -350,6 +354,160 @@ def test_replay_buffer_clear(self):
         ray.kill(buffer)
 
 
+class TestReplayBufferNew:
+    """Tests for ReplayBufferNew: staleness-window sampling via _evict + sample."""
+
+    def _make_traj(self, label: str) -> dict:
+        return {"batch": {"data": label}, "rollout_metrics": {}}
+
+    def _add(self, buf, label: str, weight_version: int):
+        return ray.get(
+            buf.add.remote(
+                self._make_traj(label),
+                weight_version=weight_version,
+                target_weight_version=0,  # unused in ReplayBufferNew
+            )
+        )
+
+    def _sample(self, buf, num_groups: int, trainer_version: int):
+        return ray.get(
+            buf.sample.remote(
+                num_prompt_groups=num_groups,
+                current_weight_version=trainer_version,
+                max_age_steps=0,  # unused in ReplayBufferNew
+            )
+        )
+
+    # ------------------------------------------------------------------
+    # Construction
+    # ------------------------------------------------------------------
+
+    def test_invalid_max_staleness_raises(self):
+        with pytest.raises(Exception):
+            buf = ReplayBufferNew.remote(max_size=10, max_staleness=-1)
+            ray.get(buf.size.remote())
+
+    # ------------------------------------------------------------------
+    # _evict (via sample)
+    # ------------------------------------------------------------------
+
+    def test_stale_rows_evicted_before_sampling(self):
+        """Rows with age > max_staleness are removed before sample() selects."""
+        buf = ReplayBufferNew.remote(max_size=10, max_staleness=2)
+        # age at trainer=4: gen_v=1 → 3 > 2 (stale), gen_v=3 → 1 ≤ 2 (valid)
+        self._add(buf, "stale", weight_version=1)
+        self._add(buf, "fresh", weight_version=3)
+
+        result = self._sample(buf, num_groups=1, trainer_version=4)
+
+        assert result is not None
+        assert result["trajectories"][0]["batch"]["data"] == "fresh"
+        assert ray.get(buf.size.remote()) == 0  # stale row also gone
+        ray.kill(buf)
+
+    def test_all_stale_returns_none(self):
+        """sample() returns None when all rows are evicted as stale."""
+        buf = ReplayBufferNew.remote(max_size=10, max_staleness=1)
+        self._add(buf, "a", weight_version=0)
+        self._add(buf, "b", weight_version=1)
+
+        # trainer=5: both ages > 1
+        result = self._sample(buf, num_groups=1, trainer_version=5)
+
+        assert result is None
+        assert ray.get(buf.size.remote()) == 0
+        ray.kill(buf)
+
+    def test_eviction_frees_capacity(self):
+        """Evicting a stale row allows a subsequent add() to succeed."""
+        buf = ReplayBufferNew.remote(max_size=1, max_staleness=1)
+        self._add(buf, "x", weight_version=1)
+        assert self._add(buf, "x", weight_version=1) == "full"
+
+        # sample() at trainer=5 evicts the stale row (age 4 > 1)
+        self._sample(buf, num_groups=1, trainer_version=5)
+
+        assert self._add(buf, "y", weight_version=4) == "success"
+        ray.kill(buf)
+
+    def test_within_window_not_evicted(self):
+        """Rows whose age is within max_staleness are not evicted."""
+        buf = ReplayBufferNew.remote(max_size=10, max_staleness=3)
+        self._add(buf, "x", weight_version=4)
+
+        # trainer=6: age = 6 - 4 = 2 ≤ 3 → should survive
+        # should return None since there is only 1 row
+        result = self._sample(buf, num_groups=2, trainer_version=6)
+        assert result is None
+
+        # this sample should still be there
+        assert ray.get(buf.size.remote()) == 1
+        ray.kill(buf)
+
+    # ------------------------------------------------------------------
+    # sample()
+    # ------------------------------------------------------------------
+
+    @pytest.mark.parametrize("sample_freshest_first", [True, False])
+    def test_sample_freshest_first(self, sample_freshest_first):
+        """sample() returns the freshest trajectories first."""
+        buf = ReplayBufferNew.remote(
+            max_size=10, max_staleness=5, sample_freshest_first=sample_freshest_first
+        )
+        for gen_v in [3, 4, 5]:
+            self._add(buf, f"v{gen_v}", weight_version=gen_v)
+
+        result = self._sample(buf, num_groups=2, trainer_version=6)
+
+        assert result is not None
+        data = [t["batch"]["data"] for t in result["trajectories"]]
+        if sample_freshest_first:
+            assert data == ["v5", "v4"]
+        else:
+            assert data == ["v3", "v4"]
+        ray.kill(buf)
+
+    def test_sample_returns_none_when_insufficient(self):
+        """sample() returns None when fewer rows than requested remain after eviction."""
+        buf = ReplayBufferNew.remote(max_size=10, max_staleness=5)
+        self._add(buf, "only", weight_version=1)
+
+        result = self._sample(buf, num_groups=3, trainer_version=2)
+
+        assert result is None
+        ray.kill(buf)
+
+    def test_sample_returns_none_on_empty_buffer(self):
+        buf = ReplayBufferNew.remote(max_size=10, max_staleness=5)
+        result = self._sample(buf, num_groups=1, trainer_version=1)
+        assert result is None
+        ray.kill(buf)
+
+    def test_sample_avg_trajectory_age(self):
+        """avg_trajectory_age is computed from the sampled generation versions."""
+        buf = ReplayBufferNew.remote(max_size=10, max_staleness=5)
+        # freshest first: gen 8 (age 2), gen 6 (age 4) → avg = 3.0
+        for gen_v in [6, 8]:
+            self._add(buf, f"v{gen_v}", weight_version=gen_v)
+
+        result = self._sample(buf, num_groups=2, trainer_version=10)
+
+        assert result is not None
+        assert abs(result["avg_trajectory_age"] - 3.0) < 1e-6
+        ray.kill(buf)
+
+    def test_sample_consumes_selected_rows(self):
+        """Rows returned by sample() are removed from the buffer."""
+        buf = ReplayBufferNew.remote(max_size=10, max_staleness=5)
+        for gen_v in [1, 2, 3]:
+            self._add(buf, f"v{gen_v}", weight_version=gen_v)
+
+        self._sample(buf, num_groups=2, trainer_version=4)
+
+        assert ray.get(buf.size.remote()) == 1
+        ray.kill(buf)
+
+
 class TestAsyncTrajectoryCollector:
     """Test cases for AsyncTrajectoryCollector."""