Enhance experience replay for priority queue buffer (agentscope-ai#306)

Author: yanxi-chen

tests/buffer/queue_test.py

@@ -93,14 +93,14 @@ def thread_read(reader, result_queue):
         self.assertRaises(StopIteration, reader.read, batch_size=1)
 
     async def test_priority_queue_capacity(self):
-        # test queue capacity
+        # test priority queue capacity
         self.config.train_batch_size = 4
         meta = StorageConfig(
             name="test_buffer_small",
             schema_type="experience",
             storage_type=StorageType.QUEUE,
             max_read_timeout=1,
-            capacity=100,  # priority will use 2 * train_batch_size as capacity (8)
+            capacity=8,
             path=BUFFER_FILE_PATH,
             use_priority_queue=True,
             replay_buffer_kwargs={"priority_fn": "linear_decay", "decay": 0.6},
@@ -177,13 +177,13 @@ def write_blocking_call():
         self.assertFalse(thread.is_alive())
 
     async def test_priority_queue_buffer_reuse(self):
-        # test queue reuse
+        # test experience replay
         meta = StorageConfig(
             name="test_buffer_small",
             schema_type="experience",
             storage_type=StorageType.QUEUE,
             max_read_timeout=3,
-            capacity=4,
+            capacity=4,  # max total number of items; each item is List[Experience]
             path=BUFFER_FILE_PATH,
             use_priority_queue=True,
             reuse_cooldown_time=0.5,
@@ -300,6 +300,109 @@ def replace_call():
         # use_count      5,   4,   2,   1
         # priority      1.0, 0.6, 0.8, 0.4
 
+    async def test_priority_queue_reuse_count_control(self):
+        # test experience replay with linear decay and use count control
+        meta = StorageConfig(
+            name="test_buffer_small",
+            schema_type="experience",
+            storage_type=StorageType.QUEUE,
+            max_read_timeout=3,
+            capacity=4,  # max total number of items; each item is List[Experience]
+            path=BUFFER_FILE_PATH,
+            use_priority_queue=True,
+            reuse_cooldown_time=0.5,
+            replay_buffer_kwargs={
+                "priority_fn": "linear_decay_use_count_control_randomization",
+                "decay": 1.2,
+                "use_count_limit": 2,
+                "sigma": 0.0,
+            },
+        )
+        writer = QueueWriter(meta, self.config)
+        reader = QueueReader(meta, self.config)
+        for i in range(4):
+            writer.write(
+                [
+                    Experience(
+                        tokens=torch.tensor([1, 2, 3]),
+                        prompt_length=2,
+                        info={"model_version": i, "use_count": 0},
+                    ),
+                    Experience(
+                        tokens=torch.tensor([1, 2, 3]),
+                        prompt_length=2,
+                        info={"model_version": i, "use_count": 0},
+                    ),
+                ]
+            )
+
+        # should not be blocked
+        def replace_call():
+            writer.write(
+                [
+                    Experience(
+                        tokens=torch.tensor([1, 2, 3]),
+                        prompt_length=2,
+                        info={"model_version": 4, "use_count": 0},
+                    ),
+                    Experience(
+                        tokens=torch.tensor([1, 2, 3]),
+                        prompt_length=2,
+                        info={"model_version": 4, "use_count": 0},
+                    ),
+                ]
+            )
+
+        thread = threading.Thread(target=replace_call)
+        thread.start()
+        thread.join(timeout=2)
+        self.assertFalse(thread.is_alive())
+
+        exps = reader.read(batch_size=4)
+        self.assertEqual(len(exps), 4)
+        self.assertEqual(exps[0].info["model_version"], 4)
+        self.assertEqual(exps[0].info["use_count"], 1)
+        self.assertEqual(exps[2].info["model_version"], 3)
+        self.assertEqual(exps[2].info["use_count"], 1)
+
+        # model_version  4,   3,   2,   1
+        # use_count      1,   1,   0,   0
+        # priority      2.8, 1.8, 2.0, 1.0
+        # in queue       Y,   Y,   Y,   Y
+
+        time.sleep(1)
+        self.assertEqual(ray.get(reader.queue.length.remote()), 4)
+        exps = reader.read(batch_size=4)
+        self.assertEqual(len(exps), 4)
+        self.assertEqual(exps[0].info["model_version"], 4)
+        self.assertEqual(exps[0].info["use_count"], 2)
+        self.assertEqual(exps[2].info["model_version"], 2)
+        self.assertEqual(exps[2].info["use_count"], 1)
+
+        # model_version  4,   3,   2,   1
+        # use_count      2,   1,   1,   0
+        # priority      1.6, 1.8, 0.8, 1.0
+        # in queue       N,   Y,   Y,   Y
+        # model_version = 4 item is discarded for reaching use_count_limit
+
+        time.sleep(1)
+        self.assertEqual(ray.get(reader.queue.length.remote()), 3)
+        exps = reader.read(batch_size=4)
+        self.assertEqual(len(exps), 4)
+        self.assertEqual(exps[0].info["model_version"], 3)
+        self.assertEqual(exps[0].info["use_count"], 2)
+        self.assertEqual(exps[2].info["model_version"], 1)
+        self.assertEqual(exps[2].info["use_count"], 1)
+
+        # model_version  3,    2,    1
+        # use_count      2,    1,    1
+        # priority      0.6,  0.8, -0.2
+        # in queue       N,    Y,    Y
+        # model_version = 3 item is discarded for reaching use_count_limit
+
+        time.sleep(1)
+        self.assertEqual(ray.get(reader.queue.length.remote()), 2)
+
     def setUp(self):
         self.total_num = 8
         self.put_batch_size = 2

trinity/buffer/storage/queue.py

@@ -5,8 +5,9 @@
 from collections import deque
 from copy import deepcopy
 from functools import partial
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
+import numpy as np
 import ray
 from sortedcontainers import SortedDict
 
@@ -26,11 +27,48 @@ def is_json_file(path: str) -> bool:
 
 
 PRIORITY_FUNC = Registry("priority_fn")
+"""
+Each priority_fn,
+    Args:
+        item: List[Experience], assume that all experiences in it have the same model_version and use_count
+        kwargs: storage_config.replay_buffer_kwargs (except priority_fn)
+    Returns:
+        priority: float
+        put_into_queue: bool, decide whether to put item into queue
+Note that put_into_queue takes effect both for new item from the explorer and for item sampled from the buffer.
+"""
 
 
 @PRIORITY_FUNC.register_module("linear_decay")
-def linear_decay_priority(item: List[Experience], decay: float = 0.1):
-    return item[0].info["model_version"] - decay * item[0].info["use_count"]  # type: ignore
+def linear_decay_priority(
+    item: List[Experience],
+    decay: float = 2.0,
+) -> Tuple[float, bool]:
+    """Calculate priority by linear decay.
+
+    Priority is calculated as `model_version - decay * use_count. The item is always put back into the queue for reuse (as long as `reuse_cooldown_time` is not None).
+    """
+    priority = float(item[0].info["model_version"] - decay * item[0].info["use_count"])
+    put_into_queue = True
+    return priority, put_into_queue
+
+
+@PRIORITY_FUNC.register_module("linear_decay_use_count_control_randomization")
+def linear_decay_use_count_control_priority(
+    item: List[Experience],
+    decay: float = 2.0,
+    use_count_limit: int = 3,
+    sigma: float = 0.0,
+) -> Tuple[float, bool]:
+    """Calculate priority by linear decay, use count control, and randomization.
+
+    Priority is calculated as `model_version - decay * use_count`; if `sigma` is non-zero, priority is further perturbed by random Gaussian noise with standard deviation `sigma`.  The item will be put back into the queue only if use count does not exceed `use_count_limit`.
+    """
+    priority = float(item[0].info["model_version"] - decay * item[0].info["use_count"])
+    if sigma > 0.0:
+        priority += float(np.random.randn() * sigma)
+    put_into_queue = item[0].info["use_count"] < use_count_limit if use_count_limit > 0 else True
+    return priority, put_into_queue
 
 
 class QueueBuffer(ABC):
@@ -61,7 +99,7 @@ def get_queue(cls, storage_config: StorageConfig, config: BufferConfig) -> "Queu
         if storage_config.use_priority_queue:
             reuse_cooldown_time = storage_config.reuse_cooldown_time
             replay_buffer_kwargs = storage_config.replay_buffer_kwargs
-            capacity = min(storage_config.capacity, config.train_batch_size * 2)
+            capacity = storage_config.capacity
             logger.info(
                 f"Using AsyncPriorityQueue with capacity {capacity}, reuse_cooldown_time {reuse_cooldown_time}."
             )
@@ -124,6 +162,7 @@ def __init__(
             kwargs: Additional keyword arguments for the priority function.
         """
         self.capacity = capacity
+        self.item_count = 0
         self.priority_groups = SortedDict()  # Maps priority -> deque of items
         self.priority_fn = partial(PRIORITY_FUNC.get(priority_fn), **kwargs)
         self.reuse_cooldown_time = reuse_cooldown_time
@@ -142,22 +181,28 @@ async def _put(self, item: List[Experience], delay: float = 0) -> None:
             await asyncio.sleep(delay)
         if len(item) == 0:
             return
-        priority = self.priority_fn(item=item)
+
+        priority, put_into_queue = self.priority_fn(item=item)
+        if not put_into_queue:
+            return
+
         async with self._condition:
-            if len(self.priority_groups) == self.capacity:
+            if self.item_count == self.capacity:
                 # If full, only insert if new item has higher or equal priority than the lowest
                 lowest_priority, item_queue = self.priority_groups.peekitem(index=0)
                 if lowest_priority > priority:
                     return  # Skip insertion if lower priority
                 # Remove the lowest priority item
                 item_queue.popleft()
+                self.item_count -= 1
                 if not item_queue:
                     self.priority_groups.popitem(index=0)
 
             # Add the new item
             if priority not in self.priority_groups:
                 self.priority_groups[priority] = deque()
             self.priority_groups[priority].append(item)
+            self.item_count += 1
             self._condition.notify()
 
     async def put(self, item: List[Experience]) -> None:
@@ -181,19 +226,20 @@ async def get(self) -> List[Experience]:
 
             _, item_queue = self.priority_groups.peekitem(index=-1)
             item = item_queue.popleft()
+            self.item_count -= 1
             if not item_queue:
                 self.priority_groups.popitem(index=-1)
 
         for exp in item:
             exp.info["use_count"] += 1
         # Optionally resubmit the item after a cooldown
         if self.reuse_cooldown_time is not None:
-            asyncio.create_task(self._put(item, self.reuse_cooldown_time))
+            asyncio.create_task(self._put(item, delay=self.reuse_cooldown_time))
 
         return item
 
     def qsize(self):
-        return len(self.priority_groups)
+        return self.item_count
 
     async def close(self) -> None:
         """

trinity/common/config.py

@@ -125,7 +125,7 @@ class StorageConfig:
     use_priority_queue: bool = False
     reuse_cooldown_time: Optional[float] = None
     replay_buffer_kwargs: dict = field(
-        default_factory=lambda: {"priority_fn": "linear_decay", "decay": 0.1}
+        default_factory=lambda: {"priority_fn": "linear_decay", "decay": 2.0}
     )
 
     # used for StorageType.SQL