implement core replay buffer logic - #792
Conversation
|
@GHOpenonic can you point me to where this is used in your code? It's a bit hard to understand how these pieces are put together from this alone. |
|
@jder most of it is used/further developed in datasets.py, config.py, base.py, stepper.py, and train.py in replay buffer agent attempt. |
|
ok, thanks. I think unless someone needs this sooner (@amogh-gulati maybe?) I vote we stage this after we land the LLC work #670 and rust data loader #800, since building this on top of those probably makes sense. |
I could work on the branch meanwhile |
|
|
||
| @dataclasses.dataclass | ||
| class ReplayEntry: | ||
| state: torch.Tensor |
There was a problem hiding this comment.
Is there any way we could get more type information for what the state is, maybe via jaxtyping?
| stride: int | ||
| temporal_stride: int | ||
|
|
||
| def advance(self) -> "ReplayCursor": |
There was a problem hiding this comment.
nit:
| def advance(self) -> "ReplayCursor": | |
| def advance(self) -> Self: |
| @dataclasses.dataclass(frozen=True) | ||
| class ReplayBatchSlot: | ||
| replay_index: int | ||
| cursor: ReplayCursor | ||
|
|
||
|
|
||
| @dataclasses.dataclass(frozen=True) | ||
| class ReplaySeedSlot: | ||
| replay_index: int | ||
| cursor: ReplayCursor | ||
| reason: str |
There was a problem hiding this comment.
🐑 IIUC these are related concepts. Maybe the seed slot could inherent from the batch slot, and then add the one additional variable? Then, I wonder if we could use a single tuple instead of two tuples in the ReplayBatchRequests (targeting the base class).
| self.buffer_size = buffer_size | ||
| self.storage_dtype = storage_dtype | ||
| self.generator = generator | ||
| self.pin_memory = pin_memory |
There was a problem hiding this comment.
Is this used similar to a Batch object by torch (e.g. a TrainData)? If so, I would recommend adding to and pin_memory methods. If these two methods exists, the torch framework will try to call them as needed.
| ) | ||
| if len(eligible) >= batch_size: | ||
| draw = torch.randperm( | ||
| len(eligible), | ||
| generator=self.generator, | ||
| device="cpu", | ||
| )[:batch_size] | ||
| else: | ||
| draw = torch.randint( | ||
| len(eligible), | ||
| (batch_size,), | ||
| generator=self.generator, | ||
| device="cpu", | ||
| ) | ||
| return [eligible[int(i)] for i in draw] |
There was a problem hiding this comment.
Is it possible that sample_indicies could call random_indices? Or, can this block (which seems to be shared with that method below) be extracted into a common, private method?
| entry.ready_event.synchronize() | ||
| return entry.state.cpu() | ||
|
|
||
| def _prepare_entry(self, entry: ReplayEntry) -> ReplayEntry: |
There was a problem hiding this comment.
See my note about to and pin_memory -- we may get to omit this method.
4 participants
add replay.py script and test. This does not add replay buffer as a training option, it only implements the core logic. Wiring it into training should be based on the LLC data munging stack #670, #671, #672, #673