feature(zjow): add load and save method for replaybuffer

ding/data/buffer/buffer.py

@@ -129,6 +129,26 @@ def delete(self, index: str):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def save_data(self, file_name: str):
+        """
+        Overview:
+            Save buffer data into a file.
+        Arguments:
+            - file_name (:obj:`str`): file name of buffer data
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_data(self, file_name: str):
+        """
+        Overview:
+            Load buffer data from a file.
+        Arguments:
+            - file_name (:obj:`str`): file name of buffer data
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def count(self) -> int:
         raise NotImplementedError

ding/data/buffer/deque_buffer.py

@@ -1,7 +1,9 @@
+import os
 import itertools
 import random
 import uuid
 from ditk import logging
+import hickle
 from typing import Any, Iterable, List, Optional, Tuple, Union
 from collections import Counter
 from collections import defaultdict, deque, OrderedDict
@@ -209,6 +211,22 @@ def delete(self, indices: Union[str, Iterable[str]]) -> None:
         key_value_pairs = zip(remain_indices, range(len(indices)))
         self.indices = BufferIndex(self.storage.maxlen, key_value_pairs)
 
+    def save_data(self, file_name: str):
+        if not os.path.exists(os.path.dirname(file_name)):
+            # If the folder for the specified file does not exist, it will be created.
+            if os.path.dirname(file_name) != "":
+                os.makedirs(os.path.dirname(file_name))
+        hickle.dump(
+            py_obj=(
+                self.storage,
+                self.indices,
+                self.meta_index,
+            ), file_obj=file_name
+        )
+
+    def load_data(self, file_name: str):
+        self.storage, self.indices, self.meta_index = hickle.load(file_name)
+
     def count(self) -> int:
         """
         Overview:
@@ -233,25 +251,6 @@ def clear(self) -> None:
         self.indices.clear()
         self.meta_index = {}
 
-    def import_data(self, data_with_meta: List[Tuple[Any, dict]]) -> None:
-        """
-        Overview:
-            The method that push data by sequence.
-        Arguments:
-            data_with_meta (List[Tuple[Any, dict]]): The sequence of (data, meta) tuples.
-        """
-        for data, meta in data_with_meta:
-            self._push(data, meta)
-
-    def export_data(self) -> List[BufferedData]:
-        """
-        Overview:
-            The method that export all data in the buffer by sequence.
-        Returns:
-            storage (List[BufferedData]): All ``BufferedData`` objects stored in the buffer.
-        """
-        return list(self.storage)
-
     def _push(self, data: Any, meta: Optional[dict] = None) -> BufferedData:
         index = uuid.uuid1().hex
         if meta is None:

ding/data/buffer/deque_buffer_wrapper.py

@@ -1,7 +1,9 @@
+import os
 from typing import Optional
 import copy
 from easydict import EasyDict
 import numpy as np
+import hickle
 
 from ding.data.buffer import DequeBuffer
 from ding.data.buffer.middleware import use_time_check, PriorityExperienceReplay
@@ -112,3 +114,9 @@ def update(self, meta: dict) -> None:
 
     def count(self) -> int:
         return self.buffer.count()
+
+    def save_data(self, file_name):
+        self.buffer.save_data(file_name)
+
+    def load_data(self, file_name: str):
+        self.buffer.load_data(file_name)

ding/data/buffer/tests/test_buffer.py

@@ -1,7 +1,9 @@
+import os
 import pytest
 import time
 import random
 import functools
+import tempfile
 from typing import Callable
 from ding.data.buffer import DequeBuffer
 from ding.data.buffer.buffer import BufferedData
@@ -87,6 +89,25 @@ def test_rate_limit_push_sample():
     assert 5 not in buffer.sample(5)
 
 
+@pytest.mark.unittest
+def test_load_and_save():
+    buffer = DequeBuffer(size=10).use(RateLimit(max_rate=5))
+    buffer.meta_index = {"label": []}
+    for i in range(10):
+        buffer.push(i, meta={"label": i})
+    assert buffer.count() == 5
+    assert 5 not in buffer.sample(5)
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_file = os.path.join(tmpdirname, "data.hkl")
+        buffer.save_data(test_file)
+        buffer_new = DequeBuffer(size=10).use(RateLimit(max_rate=5))
+        buffer_new.load_data(test_file)
+        assert buffer_new.count() == 5
+        assert 5 not in buffer_new.sample(5)
+        assert len(buffer.meta_index["label"]) == 5
+        assert all([index < 5 for index in buffer.meta_index["label"]])
+
+
 @pytest.mark.unittest
 def test_buffer_view():
     buf1 = DequeBuffer(size=10)
@@ -250,17 +271,6 @@ def test_groupby():
     assert len(sampled_data) == 2
 
 
-@pytest.mark.unittest
-def test_import_export():
-    buffer = DequeBuffer(size=10)
-    data_with_meta = [(i, {}) for i in range(10)]
-    buffer.import_data(data_with_meta)
-    assert buffer.count() == 10
-
-    sampled_data = buffer.export_data()
-    assert len(sampled_data) == 10
-
-
 @pytest.mark.unittest
 def test_dataset():
     buffer = DequeBuffer(size=10)

ding/framework/middleware/functional/__init__.py

@@ -1,6 +1,6 @@
 from .trainer import trainer, multistep_trainer
 from .data_processor import offpolicy_data_fetcher, data_pusher, offline_data_fetcher, offline_data_saver, \
-    sqil_data_pusher
+    sqil_data_pusher, buffer_saver
 from .collector import inferencer, rolloutor, TransitionList
 from .evaluator import interaction_evaluator
 from .termination_checker import termination_checker

ding/framework/middleware/functional/data_processor.py

@@ -1,3 +1,4 @@
+import os
 from typing import TYPE_CHECKING, Callable, List, Union, Tuple, Dict, Optional
 from easydict import EasyDict
 from ditk import logging
@@ -45,6 +46,42 @@ def _push(ctx: "OnlineRLContext"):
     return _push
 
 
+def buffer_saver(cfg: EasyDict, buffer_: Buffer, every_envstep: int = 1000, replace: bool = False):
+    """
+    Overview:
+        Save current buffer data.
+    Arguments:
+        - cfg (:obj:`EasyDict`): Config.
+        - buffer (:obj:`Buffer`): Buffer to push the data in.
+        - every_envstep (:obj:`int`): save at every env step.
+        - replace (:obj:`bool`): Whether replace the last file.
+    """
+
+    buffer_saver_env_counter = -every_envstep
+
+    def _save(ctx: "OnlineRLContext"):
+        """
+        Overview:
+            In ctx, `ctx.env_step` should not be None.
+        Input of ctx:
+            - env_step (:obj:`int`): env step.
+        """
+        nonlocal buffer_saver_env_counter
+        if ctx.env_step is not None:
+            if ctx.env_step >= every_envstep + buffer_saver_env_counter:
+                buffer_saver_env_counter = ctx.env_step
+                if replace:
+                    buffer_.save_data(os.path.join(cfg.exp_name, "replaybuffer", "data_latest.hkl"))
+                else:
+                    buffer_.save_data(
+                        os.path.join(cfg.exp_name, "replaybuffer", "data_envstep_{}.hkl".format(ctx.env_step))
+                    )
+        else:
+            raise RuntimeError("buffer_saver only supports collecting data by step rather than episode.")
+
+    return _save
+
+
 def offpolicy_data_fetcher(
         cfg: EasyDict,
         buffer_: Union[Buffer, List[Tuple[Buffer, float]], Dict[str, Buffer]],

ding/framework/middleware/tests/test_advantage_estimator.py

@@ -62,7 +62,7 @@ def call_gae_estimator(batch_size: int = 32, trajectory_end_idx_size: int = 5, b
         gae_estimator(cfg, MockPolicy(TheModelClass()), buffer)(ctx)
 
     if buffer is not None:
-        train_data = [d.data for d in buffer.export_data()]
+        train_data = [d.data for d in list(buffer.storage)]
         for d in train_data:
             d.logit = d.logit[0]
             d.next_obs = d.next_obs[0]

ding/framework/middleware/tests/test_data_processor.py

@@ -1,10 +1,11 @@
+import tempfile
 import pytest
 
 from ding.data.buffer import DequeBuffer
 
 from ding.framework import Context, OnlineRLContext, OfflineRLContext
 from ding.framework.middleware.functional.data_processor import \
-    data_pusher, offpolicy_data_fetcher, offline_data_fetcher, offline_data_saver, sqil_data_pusher
+    data_pusher, offpolicy_data_fetcher, offline_data_fetcher, offline_data_saver, sqil_data_pusher, buffer_saver
 
 from ding.data.buffer.middleware import PriorityExperienceReplay
 
@@ -61,14 +62,14 @@ def offpolicy_data_fetcher_type_buffer_helper(priority=0.5, use_list=True):
     assert [d['reward'] for d in ctx.train_data] == [1 for i in range(20)]
     assert [d['info'] for d in ctx.train_data] == ['xxx' for i in range(20)]
     assert [d['priority_IS'] for d in ctx.train_data] == [torch.tensor([1]) for i in range(20)]
-    assert buffer.export_data()[0].meta['priority'] == 1.0
+    assert list(buffer.storage)[0].meta['priority'] == 1.0
     # assert sorted(ctx.train_data) == [i for i in range(20)]
 
     try:
         next(func_generator)
     except StopIteration:
         pass
-    assert buffer.export_data()[0].meta['priority'] == priority
+    assert list(buffer.storage)[0].meta['priority'] == priority
 
 
 def call_offpolicy_data_fetcher_type_buffer():
@@ -224,12 +225,34 @@ def test_sqil_data_pusher():
     buffer = DequeBuffer(size=10)
     sqil_data_pusher(cfg=None, buffer_=buffer, expert=True)(ctx)
     assert buffer.count() == 5
-    assert all(t.data.reward == 1 for t in buffer.export_data())
+    assert all(t.data.reward == 1 for t in list(buffer.storage))
 
     # expert = False
     ctx = OnlineRLContext()
     ctx.trajectories = copy.deepcopy(fake_data)
     buffer = DequeBuffer(size=10)
     sqil_data_pusher(cfg=None, buffer_=buffer, expert=False)(ctx)
     assert buffer.count() == 5
-    assert all(t.data.reward == 0 for t in buffer.export_data())
+    assert all(t.data.reward == 0 for t in list(buffer.storage))
+
+
+@pytest.mark.unittest
+def test_buffer_saver():
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_folder = os.path.join(tmpdirname, "test_buffer_saver")
+        cfg = EasyDict({"exp_name": test_folder})
+        os.makedirs(test_folder)
+        buffer_ = DequeBuffer(size=10)
+        ctx = OnlineRLContext()
+        ctx.trajectories = [i for i in range(5)]
+        ctx.env_step = 0
+        data_pusher(cfg=cfg, buffer_=buffer_)(ctx)
+        assert buffer_.count() == 5
+        buffer_saver(cfg=cfg, buffer_=buffer_, replace=False)(ctx)
+        buffer_saver(cfg=cfg, buffer_=buffer_, replace=True)(ctx)
+        buffer_1 = DequeBuffer(size=10)
+        buffer_1.load_data(os.path.join(test_folder, "replaybuffer", "data_latest.hkl"))
+        assert buffer_1.count() == 5
+        buffer_2 = DequeBuffer(size=10)
+        buffer_2.load_data(os.path.join(test_folder, "replaybuffer", "data_envstep_0.hkl"))
+        assert buffer_2.count() == 5

ding/worker/replay_buffer/advanced_buffer.py

@@ -1,7 +1,9 @@
+import os
 import copy
 import time
 from typing import Union, NoReturn, Any, Optional, List, Dict, Tuple
 import numpy as np
+import hickle
 
 from ding.worker.replay_buffer import IBuffer
 from ding.utils import SumSegmentTree, MinSegmentTree, BUFFER_REGISTRY
@@ -284,6 +286,15 @@ def push(self, data: Union[List[Any], Any], cur_collector_envstep: int) -> None:
         else:
             self._append(data, cur_collector_envstep)
 
+    def save_data(self, file_name: str):
+        if not os.path.exists(os.path.dirname(file_name)):
+            if os.path.dirname(file_name) != "":
+                os.makedirs(os.path.dirname(file_name))
+        hickle.dump(py_obj=self._data, file_obj=file_name)
+
+    def load_data(self, file_name: str):
+        self.push(hickle.load(file_name), 0)
+
     def _sample_check(self, size: int, cur_learner_iter: int) -> Tuple[bool, str]:
         r"""
         Overview:

ding/worker/replay_buffer/base_buffer.py

@@ -79,6 +79,26 @@ def count(self) -> int:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def save_data(self, file_name: str):
+        """
+        Overview:
+            Save buffer data into a file.
+        Arguments:
+            - file_name (:obj:`str`): file name of buffer data
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_data(self, file_name: str):
+        """
+        Overview:
+            Load buffer data from a file.
+        Arguments:
+            - file_name (:obj:`str`): file name of buffer data
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def state_dict(self) -> Dict[str, Any]:
         """

ding/worker/replay_buffer/naive_buffer.py

@@ -1,6 +1,8 @@
+import os
 import copy
 from typing import Union, Any, Optional, List
 import numpy as np
+import hickle
 from easydict import EasyDict
 
 from ding.worker.replay_buffer import IBuffer
@@ -146,6 +148,15 @@ def sample(self,
         self._periodic_thruput_monitor.sample_data_count += len(sample_data)
         return sample_data
 
+    def save_data(self, file_name: str):
+        if not os.path.exists(os.path.dirname(file_name)):
+            if os.path.dirname(file_name) != "":
+                os.makedirs(os.path.dirname(file_name))
+        hickle.dump(py_obj=self._data, file_obj=file_name)
+
+    def load_data(self, file_name: str):
+        self.push(hickle.load(file_name), 0)
+
     def _append(self, ori_data: Any, cur_collector_envstep: int = -1) -> None:
         r"""
         Overview: