Append-only step log as a prev_id linked list (design note 5). RAM holds only the active chunk; full chunks flush to /chunk_NNNNN.npz (compressed — rewards/dones are almost all zero). With dir=None (probes/tests) full chunks stay in RAM instead.
| 90 | |
| 91 | |
| 92 | class ExperienceLog: |
| 93 | """Append-only step log as a prev_id linked list (design note 5). |
| 94 | |
| 95 | RAM holds only the active chunk; full chunks flush to |
| 96 | <dir>/chunk_NNNNN.npz (compressed — rewards/dones are almost all zero). |
| 97 | With dir=None (probes/tests) full chunks stay in RAM instead.""" |
| 98 | |
| 99 | def __init__(self, log_dir, chunk_size=EXPLOG_CHUNK, ancestor_dir=None): |
| 100 | self.dir = log_dir |
| 101 | if log_dir: |
| 102 | os.makedirs(log_dir, exist_ok=True) |
| 103 | self.chunk_size = chunk_size |
| 104 | self.ancestor = ancestor_dir # explog dir of the run we resumed FROM: |
| 105 | self.count = 0 # chunks flushed before the resume live there |
| 106 | self.n_flushed = 0 |
| 107 | self._ram_chunks = [] # dir=None mode only |
| 108 | self._cache = {} # chunk_idx -> loaded arrays (reconstruction) |
| 109 | self._new_chunk() |
| 110 | |
| 111 | def _new_chunk(self): |
| 112 | n = self.chunk_size |
| 113 | self.prev = np.empty(n, dtype=np.int64) |
| 114 | self.act = np.empty(n, dtype=np.uint8) |
| 115 | self.rew = np.empty(n, dtype=np.float32) |
| 116 | self.done = np.empty(n, dtype=np.uint8) |
| 117 | self.fill = 0 |
| 118 | |
| 119 | def append(self, prev_id, action, reward, done): |
| 120 | i = self.fill |
| 121 | self.prev[i], self.act[i], self.rew[i], self.done[i] = prev_id, action, reward, done |
| 122 | self.fill += 1 |
| 123 | idx = self.count |
| 124 | self.count += 1 |
| 125 | if self.fill == self.chunk_size: |
| 126 | self._flush() |
| 127 | return idx |
| 128 | |
| 129 | def _flush(self): |
| 130 | arrays = {"prev": self.prev[:self.fill], "act": self.act[:self.fill], |
| 131 | "rew": self.rew[:self.fill], "done": self.done[:self.fill]} |
| 132 | if self.dir: |
| 133 | tmp = os.path.join(self.dir, f"chunk_{self.n_flushed:05d}.tmp") |
| 134 | np.savez_compressed(tmp, **arrays) |
| 135 | os.replace(f"{tmp}.npz", os.path.join(self.dir, f"chunk_{self.n_flushed:05d}.npz")) |
| 136 | else: |
| 137 | self._ram_chunks.append({k: v.copy() for k, v in arrays.items()}) |
| 138 | self.n_flushed += 1 |
| 139 | self._new_chunk() |
| 140 | |
| 141 | def _chunk_path(self, chunk_idx): |
| 142 | """A flushed chunk lives in our own dir, or (after a cross-run-dir |
| 143 | resume) in the ancestor run's explog dir.""" |
| 144 | own = os.path.join(self.dir, f"chunk_{chunk_idx:05d}.npz") |
| 145 | if os.path.exists(own): |
| 146 | return own |
| 147 | if self.ancestor: |
| 148 | anc = os.path.join(self.ancestor, f"chunk_{chunk_idx:05d}.npz") |
| 149 | if os.path.exists(anc): |