prioritized experience replay bug (#527)

openai · Sep 20, 2018 · 7859f60 · 7859f60
1 parent 0f4ae2f
commit 7859f60
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/baselines/deepq/replay_buffer.py b/baselines/deepq/replay_buffer.py
@@ -106,9 +106,10 @@ def add(self, *args, **kwargs):
 
     def _sample_proportional(self, batch_size):
         res = []
-        for _ in range(batch_size):
-            # TODO(szymon): should we ensure no repeats?
-            mass = random.random() * self._it_sum.sum(0, len(self._storage) - 1)
+        p_total = self._it_sum.sum(0, len(self._storage) - 1)
+        every_range_len = p_total / batch_size
+        for i in range(batch_size):
+            mass = random.random() * every_range_len + i * every_range_len
             idx = self._it_sum.find_prefixsum_idx(mass)
             res.append(idx)
         return res