Created ClearOptimizerBuffersCallback (#222)

Author: jamesbraza

ldp/alg/__init__.py

@@ -8,6 +8,7 @@
 from .callbacks import (
     Callback,
     ClearContextCallback,
+    ClearOptimizerBuffersCallback,
     ComputeTrajectoryMetricsMixin,
     LoggingCallback,
     MeanMetricsCallback,
@@ -33,6 +34,7 @@
     "BeamSearchRollout",
     "Callback",
     "ClearContextCallback",
+    "ClearOptimizerBuffersCallback",
     "ComputeTrajectoryMetricsMixin",
     "Evaluator",
     "EvaluatorConfig",

ldp/alg/callbacks.py

@@ -6,7 +6,7 @@
 from collections import defaultdict
 from collections.abc import Callable, Collection, Iterable, Sequence
 from pathlib import Path
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import aiofiles
 from aviary.core import (
@@ -27,6 +27,9 @@
 except ImportError:
     wandb = None  # type: ignore[assignment]
 
+if TYPE_CHECKING:
+    from ldp.alg.optimizer.replay_buffers import ReplayBuffer
+
 logger = logging.getLogger(__name__)
 
 
@@ -567,3 +570,14 @@ async def after_env_step(
         print("\nObservation:")
         pprint(obs, expand_all=True)
         print(f"Elapsed time: {elapsed_time:.2f} seconds")
+
+
+class ClearOptimizerBuffersCallback(Callback):
+    """Invoke the clear method on buffer(s) after each optimizer update."""
+
+    def __init__(self, *buffers: "ReplayBuffer"):
+        self._buffers = list(buffers)
+
+    async def after_update(self) -> None:
+        for b in self._buffers:
+            b.clear()

tests/test_buffers.py

@@ -7,7 +7,7 @@
 
 
 @pytest.mark.asyncio
-async def test_circular_buffer():
+async def test_circular_buffer() -> None:
     buf = CircularReplayBuffer()
 
     samples = [{"state": 1, "action": 2, "reward": 3, "t": t} for t in range(5)]
@@ -27,6 +27,9 @@ async def test_circular_buffer():
     ):
         next(buf.batched_iter(batch_size=4))
 
+    buf.clear()
+    assert not buf, "Failed to clear data"
+
 
 async def _dummy_q_function(*args, **kwargs) -> float:  # noqa: ARG001
     return 1.0