Coverage for cafl4ds/filters/reservoir.py: 95%

1"""Reservoir sampling + experience replay — the replay knob (Vitter 1985; baseline B1). 

2 

3A fixed-capacity buffer maintained by **Vitter's Algorithm R**: after ``N`` frames have passed, 

4each is present with equal probability ``capacity / N``, so the buffer is a uniform random 

5sample of the whole (single-pass, unbounded) stream — no matter how correlated the arrival 

6order. Each step the model trains on the incoming frames **plus** a random replay draw from the 

7buffer, mixing past eras back into a correlated stream. That interleaving is the "dumb" 

8protective bar (B1): replay alone, no informativeness, partly counters forgetting. 

9 

10The buffer is the **replay** role of the ``A`` factor (see :mod:`cafl4ds.filters.base`): it 

11runs after any admission filters, so what it stores and replays is the *admitted* stream — 

12upstream :class:`~cafl4ds.filters.dedup.SemDeDup` gives reservoir + dedup (baseline B1.5). 

13Sampling uses a private, seeded generator so a run is reproducible independently of whatever 

14else consumes the global RNG. 

15""" 

16 

17from __future__ import annotations 

18 

19import torch 

20 

21from cafl4ds.filters.base import FilterContext, ReplayBuffer 

22 

23 

24class ReservoirReplay(ReplayBuffer): 

25 """Uniform reservoir buffer (Algorithm R) with per-step experience replay.""" 

26 

27 def __init__(self, capacity: int = 256, replay_batch: int = 32, seed: int = 0) -> None: 

28 """Configure the reservoir. 

29 

30 Args: 

31 capacity: Maximum number of frames held in the buffer. 

32 replay_batch: Number of past frames replayed alongside each incoming batch (a 

33 uniform draw, without replacement, from the current buffer; clamped to the 

34 buffer's size, so early steps replay fewer). 

35 seed: Seed for the buffer's private RNG (reservoir replacement + replay sampling). 

36 

37 Raises: 

38 ValueError: If ``capacity`` or ``replay_batch`` is not positive. 

39 """ 

40 if capacity < 1 or replay_batch < 1: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true

41 raise ValueError(f"capacity and replay_batch must be >= 1; got {capacity}, {replay_batch}.") 

42 self.capacity = capacity 

43 self.replay_batch = replay_batch 

44 self._buffer: list[torch.Tensor] = [] 

45 self._seen = 0 # total frames observed so far (Algorithm R's running count N) 

46 self._gen = torch.Generator().manual_seed(seed) # CPU generator; indices only 

47 

48 def observe(self, images: torch.Tensor, ctx: FilterContext) -> torch.Tensor: 

49 """Replay from the buffer, then ingest the incoming frames (Algorithm R). 

50 

51 Replay is sampled from the buffer *before* the incoming frames are ingested, so it is 

52 drawn from genuine history (never the just-arrived frames), then each incoming frame 

53 updates the reservoir. 

54 

55 Args: 

56 images: The admitted images ``[K, C, H, W]`` for this step. 

57 ctx: Unused (reservoir replay is model-agnostic and random). 

58 

59 Returns: 

60 The training batch ``[K + R, C, H, W]``: the incoming frames followed by ``R`` 

61 replayed frames (``R = min(replay_batch, buffer size before this step)``). 

62 """ 

63 del ctx # reservoir selection is model-agnostic 

64 replay = self._sample_replay() 

65 for i in range(images.shape[0]): 

66 self._ingest(images[i]) 

67 return images if replay is None else torch.cat([images, replay], dim=0) 

68 

69 def _sample_replay(self) -> torch.Tensor | None: 

70 """Draw a uniform replay batch from the current buffer (without replacement). 

71 

72 Returns: 

73 A ``[R, C, H, W]`` stack of replayed frames, or ``None`` if the buffer is empty. 

74 """ 

75 if not self._buffer: 

76 return None 

77 r = min(self.replay_batch, len(self._buffer)) 

78 idx = torch.randperm(len(self._buffer), generator=self._gen)[:r].tolist() 

79 return torch.stack([self._buffer[i] for i in idx], dim=0) 

80 

81 def _ingest(self, image: torch.Tensor) -> None: 

82 """Update the reservoir with one frame via Vitter's Algorithm R. 

83 

84 While the buffer is not full every frame is stored; once full, the ``N``-th frame 

85 (0-based index ``N``) replaces a uniformly-random slot with probability 

86 ``capacity / (N + 1)`` — the invariant that keeps the buffer a uniform sample. 

87 

88 Args: 

89 image: A single frame ``[C, H, W]`` to (maybe) store; a detached clone is kept. 

90 """ 

91 item = image.detach().clone() 

92 if len(self._buffer) < self.capacity: 

93 self._buffer.append(item) 

94 else: 

95 j = int(torch.randint(0, self._seen + 1, (1,), generator=self._gen).item()) 

96 if j < self.capacity: 

97 self._buffer[j] = item 

98 self._seen += 1