Single-frame uint8 buffer — stacks of 4 are reconstructed at sample time, cutting RAM ~4x vs. storing the full stack per slot.
| 56 | |
| 57 | |
| 58 | class ReplayBuffer: |
| 59 | """Single-frame uint8 buffer — stacks of 4 are reconstructed at sample time, |
| 60 | cutting RAM ~4x vs. storing the full stack per slot.""" |
| 61 | |
| 62 | def __init__(self, capacity, frame_shape=(84, 84), stack=4): |
| 63 | self.capacity = capacity |
| 64 | self.stack = stack |
| 65 | self.frames = np.zeros((capacity, *frame_shape), dtype=np.uint8) |
| 66 | self.actions = np.zeros(capacity, dtype=np.int64) |
| 67 | self.rewards = np.zeros(capacity, dtype=np.float32) |
| 68 | self.dones = np.zeros(capacity, dtype=np.float32) |
| 69 | self.idx = 0 |
| 70 | self.size = 0 |
| 71 | |
| 72 | def push(self, frame, action, reward, done): |
| 73 | self.frames[self.idx] = frame |
| 74 | self.actions[self.idx] = action |
| 75 | self.rewards[self.idx] = reward |
| 76 | self.dones[self.idx] = float(done) |
| 77 | self.idx = (self.idx + 1) % self.capacity |
| 78 | self.size = min(self.size + 1, self.capacity) |
| 79 | |
| 80 | def _stack(self, idx): |
| 81 | # Gather frames[idx-stack+1 .. idx]; newest at last channel. |
| 82 | offsets = np.arange(self.stack) |
| 83 | gather = (idx[:, None] - (self.stack - 1) + offsets[None, :]) % self.capacity |
| 84 | out = self.frames[gather] |
| 85 | # Zero out frames sitting before an episode boundary inside the stack. |
| 86 | # dones at the (stack-1) older positions mark where a prior episode ended. |
| 87 | older = self.dones[gather[:, :-1]].astype(bool) |
| 88 | # Once we cross any done walking newest→oldest, everything older is invalid. |
| 89 | invalid = np.cumsum(older[:, ::-1], axis=1)[:, ::-1] > 0 |
| 90 | mask = np.concatenate([~invalid, np.ones((idx.shape[0], 1), dtype=bool)], axis=1) |
| 91 | return out * mask[:, :, None, None] |
| 92 | |
| 93 | def sample(self, batch_size, device): |
| 94 | # Reject indices whose stack would straddle the write head (stale frames). |
| 95 | while True: |
| 96 | if self.size < self.capacity: |
| 97 | if self.size < self.stack + 2: |
| 98 | raise RuntimeError("buffer too small to sample yet") |
| 99 | idx = np.random.randint(self.stack - 1, self.size - 1, size=batch_size) |
| 100 | break |
| 101 | idx = np.random.randint(0, self.capacity, size=batch_size) |
| 102 | dist = (self.idx - 1 - idx) % self.capacity |
| 103 | if np.all(dist >= self.stack): |
| 104 | break |
| 105 | states = self._stack(idx) |
| 106 | next_states = self._stack((idx + 1) % self.capacity) |
| 107 | return ( |
| 108 | torch.as_tensor(states, device=device), |
| 109 | torch.as_tensor(self.actions[idx], device=device), |
| 110 | torch.as_tensor(self.rewards[idx], device=device), |
| 111 | torch.as_tensor(next_states, device=device), |
| 112 | torch.as_tensor(self.dones[idx], device=device), |
| 113 | ) |
| 114 | |
| 115 |