File size: 4,460 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
from typing import Union, Dict, Any, List
from abc import ABC, abstractmethod
import copy
from easydict import EasyDict
from ding.utils import import_module, BUFFER_REGISTRY
class IBuffer(ABC):
r"""
Overview:
Buffer interface
Interfaces:
default_config, push, update, sample, clear, count, state_dict, load_state_dict
"""
@classmethod
def default_config(cls) -> EasyDict:
r"""
Overview:
Default config of this buffer class.
Returns:
- default_config (:obj:`EasyDict`)
"""
cfg = EasyDict(copy.deepcopy(cls.config))
cfg.cfg_type = cls.__name__ + 'Dict'
return cfg
@abstractmethod
def push(self, data: Union[List[Any], Any], cur_collector_envstep: int) -> None:
r"""
Overview:
Push a data into buffer.
Arguments:
- data (:obj:`Union[List[Any], Any]`): The data which will be pushed into buffer. Can be one \
(in `Any` type), or many(int `List[Any]` type).
- cur_collector_envstep (:obj:`int`): Collector's current env step.
"""
raise NotImplementedError
@abstractmethod
def update(self, info: Dict[str, list]) -> None:
r"""
Overview:
Update data info, e.g. priority.
Arguments:
- info (:obj:`Dict[str, list]`): Info dict. Keys depends on the specific buffer type.
"""
raise NotImplementedError
@abstractmethod
def sample(self, batch_size: int, cur_learner_iter: int) -> list:
r"""
Overview:
Sample data with length ``batch_size``.
Arguments:
- size (:obj:`int`): The number of the data that will be sampled.
- cur_learner_iter (:obj:`int`): Learner's current iteration.
Returns:
- sampled_data (:obj:`list`): A list of data with length `batch_size`.
"""
raise NotImplementedError
@abstractmethod
def clear(self) -> None:
"""
Overview:
Clear all the data and reset the related variables.
"""
raise NotImplementedError
@abstractmethod
def count(self) -> int:
"""
Overview:
Count how many valid datas there are in the buffer.
Returns:
- count (:obj:`int`): Number of valid data.
"""
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]:
"""
Overview:
Provide a state dict to keep a record of current buffer.
Returns:
- state_dict (:obj:`Dict[str, Any]`): A dict containing all important values in the buffer. \
With the dict, one can easily reproduce the buffer.
"""
raise NotImplementedError
@abstractmethod
def load_state_dict(self, _state_dict: Dict[str, Any]) -> None:
"""
Overview:
Load state dict to reproduce the buffer.
Returns:
- state_dict (:obj:`Dict[str, Any]`): A dict containing all important values in the buffer.
"""
raise NotImplementedError
def create_buffer(cfg: EasyDict, *args, **kwargs) -> IBuffer:
r"""
Overview:
Create a buffer according to cfg and other arguments.
Arguments:
- cfg (:obj:`EasyDict`): Buffer config.
ArgumentsKeys:
- necessary: `type`
"""
import_module(cfg.get('import_names', []))
if cfg.type == 'naive':
kwargs.pop('tb_logger', None)
return BUFFER_REGISTRY.build(cfg.type, cfg, *args, **kwargs)
def get_buffer_cls(cfg: EasyDict) -> type:
r"""
Overview:
Get a buffer class according to cfg.
Arguments:
- cfg (:obj:`EasyDict`): Buffer config.
ArgumentsKeys:
- necessary: `type`
"""
import_module(cfg.get('import_names', []))
return BUFFER_REGISTRY.get(cfg.type)
|