-
Notifications
You must be signed in to change notification settings - Fork 16
Support dp_size in replay buffer #93
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 8 commits
494654b
88f0672
4647052
f55b4a4
98649c1
0b008ac
4a64f9a
3cb5d23
1371400
9357cab
16a1b97
7605132
2a4a7ac
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -17,8 +17,9 @@ | |
class ReplayBuffer(ForgeActor): | ||
"""Simple in-memory replay buffer implementation.""" | ||
|
||
batch_size: int = 4 | ||
max_policy_age: int = 0 | ||
batch_size: int | ||
max_policy_age: int | ||
dp_size: int = 1 | ||
seed: int | None = None | ||
|
||
@endpoint | ||
|
@@ -43,23 +44,32 @@ async def sample(self, curr_policy_version: int, batch_size: int | None = None): | |
passed in at initialization. | ||
|
||
Returns: | ||
A list of sampled episodes or None if there are not enough episodes in the buffer. | ||
A list of sampled episodes with shape (dp_size, bsz, ...) or None if there are not enough episodes in the buffer. | ||
""" | ||
bsz = batch_size if batch_size is not None else self.batch_size | ||
total_samples = self.dp_size * bsz | ||
|
||
# Evict old episodes | ||
self._evict(curr_policy_version) | ||
|
||
if bsz > len(self.buffer): | ||
if total_samples > len(self.buffer): | ||
return None | ||
|
||
# TODO: Make this more efficient | ||
idx_to_sample = self.sampler(range(len(self.buffer)), k=bsz) | ||
idx_to_sample = self.sampler(range(len(self.buffer)), k=total_samples) | ||
sorted_idxs = sorted( | ||
idx_to_sample, reverse=True | ||
) # Sort in desc order to avoid shifting idxs | ||
sampled_episodes = [self.buffer.pop(i) for i in sorted_idxs] | ||
return sampled_episodes | ||
|
||
# Reshape to (dp_size, bsz, ...) | ||
reshaped_episodes = [] | ||
for dp_idx in range(self.dp_size): | ||
|
||
start_idx = dp_idx * bsz | ||
end_idx = start_idx + bsz | ||
reshaped_episodes.append(sampled_episodes[start_idx:end_idx]) | ||
|
||
return reshaped_episodes | ||
|
||
@endpoint | ||
async def evict(self, curr_policy_version: int) -> None: | ||
|
Uh oh!
There was an error while loading. Please reload this page.