-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathnaive_buffer.py
More file actions
executable file
·352 lines (330 loc) · 14.2 KB
/
Copy pathnaive_buffer.py
File metadata and controls
executable file
·352 lines (330 loc) · 14.2 KB
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
import copy
from typing import Union, Any, Optional, List
import numpy as np
from easydict import EasyDict
from ding.worker.replay_buffer import IBuffer
from ding.utils import LockContext, LockContextType, BUFFER_REGISTRY, build_logger
from .utils import UsedDataRemover, PeriodicThruputMonitor
@BUFFER_REGISTRY.register('naive')
class NaiveReplayBuffer(IBuffer):
r"""
Overview:
Naive replay buffer, can store and sample data.
An naive implementation of replay buffer with no priority or any other advanced features.
This buffer refers to multi-thread/multi-process and guarantees thread-safe, which means that methods like
``sample``, ``push``, ``clear`` are all mutual to each other.
Interface:
start, close, push, update, sample, clear, count, state_dict, load_state_dict, default_config
Property:
replay_buffer_size, push_count
"""
config = dict(
type='naive',
replay_buffer_size=10000,
deepcopy=False,
# default `False` for serial pipeline
enable_track_used_data=False,
)
def __init__(
self,
cfg: 'EasyDict', # noqa
tb_logger: Optional['SummaryWriter'] = None, # noqa
exp_name: Optional[str] = 'default_experiment',
instance_name: Optional[str] = 'buffer',
) -> None:
"""
Overview:
Initialize the buffer
Arguments:
- cfg (:obj:`dict`): Config dict.
- tb_logger (:obj:`Optional['SummaryWriter']`): Outer tb logger. Usually get this argument in serial mode.
- exp_name (:obj:`Optional[str]`): Name of this experiment.
- instance_name (:obj:`Optional[str]`): Name of this instance.
"""
self._exp_name = exp_name
self._instance_name = instance_name
self._cfg = cfg
self._replay_buffer_size = self._cfg.replay_buffer_size
self._deepcopy = self._cfg.deepcopy
# ``_data`` is a circular queue to store data (full data or meta data)
self._data = [None for _ in range(self._replay_buffer_size)]
# Current valid data count, indicating how many elements in ``self._data`` is valid.
self._valid_count = 0
# How many pieces of data have been pushed into this buffer, should be no less than ``_valid_count``.
self._push_count = 0
# Point to the tail position where next data can be inserted, i.e. latest inserted data's next position.
self._tail = 0
# Lock to guarantee thread safe
self._lock = LockContext(type_=LockContextType.THREAD_LOCK)
self._end_flag = False
self._enable_track_used_data = self._cfg.enable_track_used_data
if self._enable_track_used_data:
self._used_data_remover = UsedDataRemover()
if tb_logger is not None:
self._logger, _ = build_logger(
'./{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False
)
self._tb_logger = tb_logger
else:
self._logger, self._tb_logger = build_logger(
'./{}/log/{}'.format(self._exp_name, self._instance_name),
self._instance_name,
)
# Periodic thruput. Here by default, monitor range is 60 seconds. You can modify it for free.
self._periodic_thruput_monitor = PeriodicThruputMonitor(
self._instance_name, EasyDict(seconds=60), self._logger, self._tb_logger
)
def start(self) -> None:
"""
Overview:
Start the buffer's used_data_remover thread if enables track_used_data.
"""
if self._enable_track_used_data:
self._used_data_remover.start()
def close(self) -> None:
"""
Overview:
Clear the buffer; Join the buffer's used_data_remover thread if enables track_used_data.
"""
self.clear()
if self._enable_track_used_data:
self._used_data_remover.close()
self._tb_logger.flush()
self._tb_logger.close()
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. \
Not used in naive buffer, but preserved for compatibility.
"""
if isinstance(data, list):
self._extend(data, cur_collector_envstep)
self._periodic_thruput_monitor.push_data_count += len(data)
else:
self._append(data, cur_collector_envstep)
self._periodic_thruput_monitor.push_data_count += 1
def sample(self, size: int, cur_learner_iter: int, sample_range: slice = None) -> Optional[list]:
"""
Overview:
Sample data with length ``size``.
Arguments:
- size (:obj:`int`): The number of the data that will be sampled.
- cur_learner_iter (:obj:`int`): Learner's current iteration. \
Not used in naive buffer, but preserved for compatibility.
- sample_range (:obj:`slice`): Buffer slice for sampling, such as `slice(-10, None)`, which \
means only sample among the last 10 data
Returns:
- sample_data (:obj:`list`): A list of data with length ``size``.
"""
if size == 0:
return []
can_sample = self._sample_check(size)
if not can_sample:
return None
with self._lock:
indices = self._get_indices(size, sample_range)
sample_data = self._sample_with_indices(indices, cur_learner_iter)
self._periodic_thruput_monitor.sample_data_count += len(sample_data)
return sample_data
def _append(self, ori_data: Any, cur_collector_envstep: int = -1) -> None:
r"""
Overview:
Append a data item into ``self._data``.
Arguments:
- ori_data (:obj:`Any`): The data which will be inserted.
- cur_collector_envstep (:obj:`int`): Not used in this method, but preserved for compatibility.
"""
with self._lock:
if self._deepcopy:
data = copy.deepcopy(ori_data)
else:
data = ori_data
self._push_count += 1
if self._data[self._tail] is None:
self._valid_count += 1
self._periodic_thruput_monitor.valid_count = self._valid_count
elif self._enable_track_used_data:
self._used_data_remover.add_used_data(self._data[self._tail])
self._data[self._tail] = data
self._tail = (self._tail + 1) % self._replay_buffer_size
def _extend(self, ori_data: List[Any], cur_collector_envstep: int = -1) -> None:
r"""
Overview:
Extend a data list into queue.
Add two keys in each data item, you can refer to ``_append`` for details.
Arguments:
- ori_data (:obj:`List[Any]`): The data list.
- cur_collector_envstep (:obj:`int`): Not used in this method, but preserved for compatibility.
"""
with self._lock:
if self._deepcopy:
data = copy.deepcopy(ori_data)
else:
data = ori_data
length = len(data)
# When updating ``_data`` and ``_use_count``, should consider two cases regarding
# the relationship between "tail + data length" and "replay buffer size" to check whether
# data will exceed beyond buffer's max length limitation.
if self._tail + length <= self._replay_buffer_size:
if self._valid_count != self._replay_buffer_size:
self._valid_count += length
self._periodic_thruput_monitor.valid_count = self._valid_count
elif self._enable_track_used_data:
for i in range(length):
self._used_data_remover.add_used_data(self._data[self._tail + i])
self._push_count += length
self._data[self._tail:self._tail + length] = data
else:
new_tail = self._tail
data_start = 0
residual_num = len(data)
while True:
space = self._replay_buffer_size - new_tail
L = min(space, residual_num)
if self._valid_count != self._replay_buffer_size:
self._valid_count += L
self._periodic_thruput_monitor.valid_count = self._valid_count
elif self._enable_track_used_data:
for i in range(L):
self._used_data_remover.add_used_data(self._data[new_tail + i])
self._push_count += L
self._data[new_tail:new_tail + L] = data[data_start:data_start + L]
residual_num -= L
assert residual_num >= 0
if residual_num == 0:
break
else:
new_tail = 0
data_start += L
# Update ``tail`` and ``next_unique_id`` after the whole list is pushed into buffer.
self._tail = (self._tail + length) % self._replay_buffer_size
def _sample_check(self, size: int) -> bool:
r"""
Overview:
Check whether this buffer has more than `size` datas to sample.
Arguments:
- size (:obj:`int`): Number of data that will be sampled.
Returns:
- can_sample (:obj:`bool`): Whether this buffer can sample enough data.
"""
if self._valid_count < size:
print("No enough elements for sampling (expect: {} / current: {})".format(size, self._valid_count))
return False
else:
return True
def update(self, info: dict) -> None:
r"""
Overview:
Naive Buffer does not need to update any info, but this method is preserved for compatibility.
"""
print(
'[BUFFER WARNING] Naive Buffer does not need to update any info, \
but `update` method is preserved for compatibility.'
)
def clear(self) -> None:
"""
Overview:
Clear all the data and reset the related variables.
"""
with self._lock:
for i in range(len(self._data)):
if self._data[i] is not None:
if self._enable_track_used_data:
self._used_data_remover.add_used_data(self._data[i])
self._data[i] = None
self._valid_count = 0
self._periodic_thruput_monitor.valid_count = self._valid_count
self._push_count = 0
self._tail = 0
def __del__(self) -> None:
"""
Overview:
Call ``close`` to delete the object.
"""
self.close()
def _get_indices(self, size: int, sample_range: slice = None) -> list:
r"""
Overview:
Get the sample index list.
Arguments:
- size (:obj:`int`): The number of the data that will be sampled
- sample_range (:obj:`slice`): Buffer slice for sampling, such as `slice(-10, None)`, which \
means only sample among the last 10 data
Returns:
- index_list (:obj:`list`): A list including all the sample indices, whose length should equal to ``size``.
"""
assert self._valid_count <= self._replay_buffer_size
if self._valid_count == self._replay_buffer_size:
tail = self._replay_buffer_size
else:
tail = self._tail
if sample_range is None:
indices = list(np.random.choice(a=tail, size=size, replace=False))
else:
indices = list(range(tail))[sample_range]
indices = list(np.random.choice(indices, size=size, replace=False))
return indices
def _sample_with_indices(self, indices: List[int], cur_learner_iter: int) -> list:
r"""
Overview:
Sample data with ``indices``.
Arguments:
- indices (:obj:`List[int]`): A list including all the sample indices.
- cur_learner_iter (:obj:`int`): Not used in this method, but preserved for compatibility.
Returns:
- data (:obj:`list`) Sampled data.
"""
data = []
for idx in indices:
assert self._data[idx] is not None, idx
if self._deepcopy:
copy_data = copy.deepcopy(self._data[idx])
else:
copy_data = self._data[idx]
data.append(copy_data)
return data
def count(self) -> int:
"""
Overview:
Count how many valid datas there are in the buffer.
Returns:
- count (:obj:`int`): Number of valid data.
"""
return self._valid_count
def state_dict(self) -> dict:
"""
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.
"""
return {
'data': self._data,
'tail': self._tail,
'valid_count': self._valid_count,
'push_count': self._push_count,
}
def load_state_dict(self, _state_dict: dict) -> 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.
"""
assert 'data' in _state_dict
if set(_state_dict.keys()) == set(['data']):
self._extend(_state_dict['data'])
else:
for k, v in _state_dict.items():
setattr(self, '_{}'.format(k), v)
@property
def replay_buffer_size(self) -> int:
return self._replay_buffer_size
@property
def push_count(self) -> int:
return self._push_count