[RLlib] Fix estimated buffer size in replay buffers. (#24848)

This commit is contained in:
Artur Niederfahrenhorst 2022-05-22 21:03:23 +02:00 committed by GitHub
parent 501d932449
commit cd16dc4dae
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 7 additions and 0 deletions

View file

@ -203,7 +203,10 @@ class ReplayBuffer(ParallelIteratorWorker):
self._storage.append(item)
self._est_size_bytes += item.size_bytes()
else:
item_to_be_removed = self._storage[self._next_idx]
self._est_size_bytes -= item_to_be_removed.size_bytes()
self._storage[self._next_idx] = item
self._est_size_bytes += item.size_bytes()
# Eviction of older samples has already started (buffer is "full").
if self._eviction_started:

View file

@ -74,7 +74,11 @@ class ReservoirBuffer(ReplayBuffer):
self._next_idx = idx
self._evicted_hit_stats.push(self._hit_count[idx])
self._hit_count[idx] = 0
item_to_be_removed = self._storage[idx]
self._est_size_bytes -= item_to_be_removed.size_bytes()
self._storage[idx] = item
self._est_size_bytes += item.size_bytes()
assert item.count > 0, item
warn_replay_capacity(item=item, num_items=self.capacity / item.count)