-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathr2d2.py
More file actions
executable file
·434 lines (395 loc) · 21 KB
/
Copy pathr2d2.py
File metadata and controls
executable file
·434 lines (395 loc) · 21 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
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
from typing import List, Dict, Any, Tuple, Union, Optional
from collections import namedtuple
import torch
import copy
from ding.torch_utils import Adam, to_device
from ding.rl_utils import q_nstep_td_data, q_nstep_td_error, q_nstep_td_error_with_rescale, get_nstep_return_data, \
get_train_sample
from ding.model import model_wrap
from ding.utils import POLICY_REGISTRY
from ding.utils.data import timestep_collate, default_collate, default_decollate
from .base_policy import Policy
@POLICY_REGISTRY.register('r2d2')
class R2D2Policy(Policy):
r"""
Overview:
Policy class of R2D2, from paper `Recurrent Experience Replay in Distributed Reinforcement Learning` .
R2D2 proposes that several tricks should be used to improve upon DRQN,
namely some recurrent experience replay tricks such as burn-in.
Config:
== ==================== ======== ============== ======================================== =======================
ID Symbol Type Default Value Description Other(Shape)
== ==================== ======== ============== ======================================== =======================
1 ``type`` str dqn | RL policy register name, refer to | This arg is optional,
| registry ``POLICY_REGISTRY`` | a placeholder
2 ``cuda`` bool False | Whether to use cuda for network | This arg can be diff-
| erent from modes
3 ``on_policy`` bool False | Whether the RL algorithm is on-policy
| or off-policy
4 ``priority`` bool False | Whether use priority(PER) | Priority sample,
| update priority
5 | ``priority_IS`` bool False | Whether use Importance Sampling Weight
| ``_weight`` | to correct biased update. If True,
| priority must be True.
6 | ``discount_`` float 0.997, | Reward's future discount factor, aka. | May be 1 when sparse
| ``factor`` [0.95, 0.999] | gamma | reward env
7 ``nstep`` int 3, | N-step reward discount sum for target
[3, 5] | q_value estimation
8 ``burnin_step`` int 2 | The timestep of burnin operation,
| which is designed to RNN hidden state
| difference caused by off-policy
9 | ``learn.update`` int 1 | How many updates(iterations) to train | This args can be vary
| ``per_collect`` | after collector's one collection. Only | from envs. Bigger val
| valid in serial training | means more off-policy
10 | ``learn.batch_`` int 64 | The number of samples of an iteration
| ``size``
11 | ``learn.learning`` float 0.001 | Gradient step length of an iteration.
| ``_rate``
12 | ``learn.value_`` bool True | Whether use value_rescale function for
| ``rescale`` | predicted value
13 | ``learn.target_`` int 100 | Frequence of target network update. | Hard(assign) update
| ``update_freq``
14 | ``learn.ignore_`` bool False | Whether ignore done for target value | Enable it for some
| ``done`` | calculation. | fake termination env
15 ``collect.n_sample`` int [8, 128] | The number of training samples of a | It varies from
| call of collector. | different envs
16 | ``collect.unroll`` int 1 | unroll length of an iteration | In RNN, unroll_len>1
| ``_len``
== ==================== ======== ============== ======================================== =======================
"""
config = dict(
# (str) RL policy register name (refer to function "POLICY_REGISTRY").
type='r2d2',
# (bool) Whether to use cuda for network.
cuda=False,
# (bool) Whether the RL algorithm is on-policy or off-policy.
on_policy=False,
# (bool) Whether use priority(priority sample, IS weight, update priority)
priority=False,
# (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True.
priority_IS_weight=False,
# ==============================================================
# The following configs are algorithm-specific
# ==============================================================
# (float) Reward's future discount factor, aka. gamma.
discount_factor=0.997,
# (int) N-step reward for target q_value estimation
nstep=5,
# (int) the timestep of burnin operation, which is designed to RNN hidden state difference
# caused by off-policy
burnin_step=2,
# (int) the trajectory length to unroll the RNN network minus
# the timestep of burnin operation
unroll_len=80,
learn=dict(
# (bool) Whether to use multi gpu
multi_gpu=False,
update_per_collect=1,
batch_size=64,
learning_rate=0.0001,
# ==============================================================
# The following configs are algorithm-specific
# ==============================================================
# (int) Frequence of target network update.
target_update_freq=100,
# (bool) whether use value_rescale function for predicted value
value_rescale=True,
ignore_done=False,
),
collect=dict(
# (int) Only one of [n_sample, n_episode] shoule be set
n_sample=64,
# `env_num` is used in hidden state, should equal to that one in env config.
# User should specify this value in user config.
env_num=None,
),
eval=dict(
# `env_num` is used in hidden state, should equal to that one in env config.
# User should specify this value in user config.
env_num=None,
),
other=dict(
eps=dict(
type='exp',
start=0.95,
end=0.05,
decay=10000,
),
replay_buffer=dict(replay_buffer_size=10000, ),
),
)
def _init_learn(self) -> None:
r"""
Overview:
Init the learner model of R2D2Policy
Arguments:
.. note::
The _init_learn method takes the argument from the self._cfg.learn in the config file
- learning_rate (:obj:`float`): The learning rate fo the optimizer
- gamma (:obj:`float`): The discount factor
- nstep (:obj:`int`): The num of n step return
- value_rescale (:obj:`bool`): Whether to use value rescaled loss in algorithm
- burnin_step (:obj:`int`): The num of step of burnin
"""
self._priority = self._cfg.priority
self._priority_IS_weight = self._cfg.priority_IS_weight
self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate)
self._gamma = self._cfg.discount_factor
self._nstep = self._cfg.nstep
self._burnin_step = self._cfg.burnin_step
self._value_rescale = self._cfg.learn.value_rescale
self._target_model = copy.deepcopy(self._model)
self._target_model = model_wrap(
self._target_model,
wrapper_name='target',
update_type='assign',
update_kwargs={'freq': self._cfg.learn.target_update_freq}
)
self._target_model = model_wrap(
self._target_model,
wrapper_name='hidden_state',
state_num=self._cfg.learn.batch_size,
)
self._learn_model = model_wrap(
self._model,
wrapper_name='hidden_state',
state_num=self._cfg.learn.batch_size,
)
self._learn_model = model_wrap(self._learn_model, wrapper_name='argmax_sample')
self._learn_model.reset()
self._target_model.reset()
def _data_preprocess_learn(self, data: List[Dict[str, Any]]) -> dict:
r"""
Overview:
Preprocess the data to fit the required data format for learning
Arguments:
- data (:obj:`List[Dict[str, Any]]`): the data collected from collect function
Returns:
- data (:obj:`Dict[str, Any]`): the processed data, including at least \
['main_obs', 'target_obs', 'burnin_obs', 'action', 'reward', 'done', 'weight']
- data_info (:obj:`dict`): the data info, such as replay_buffer_idx, replay_unique_id
"""
# data preprocess
data = timestep_collate(data)
if self._cuda:
data = to_device(data, self._device)
bs = self._burnin_step
# data['done'], data['weight'], data['value_gamma'] is used in def _forward_learn() to calculate
# the q_nstep_td_error, should be length of [self._unroll_len_add_burnin_step-self._burnin_step-self._nstep]
ignore_done = self._cfg.learn.ignore_done
if ignore_done:
data['done'] = [None for _ in range(self._unroll_len_add_burnin_step - bs - self._nstep)]
else:
data['done'] = data['done'][bs + self._nstep:].float() # for computation of online model self._learn_model
# if the data don't include 'weight' or 'value_gamma' then fill in None in a list
# with length of [self._unroll_len_add_burnin_step-self._burnin_step-self._nstep],
# below is two different implementation ways
if 'value_gamma' not in data:
data['value_gamma'] = [None for _ in range(self._unroll_len_add_burnin_step - bs - self._nstep)]
else:
data['value_gamma'] = data['value_gamma'][bs + self._nstep:]
data['weight'] = data.get('weight', [None for _ in range(self._unroll_len_add_burnin_step - bs - self._nstep)])
data['action'] = data['action'][bs:-self._nstep]
data['reward'] = data['reward'][bs:-self._nstep]
# the burnin_nstep_obs is used to calculate the init hidden state of rnn for the calculation of the q_value,
# target_q_value, and target_q_action
data['burnin_nstep_obs'] = data['obs'][:bs + self._nstep]
# the main_obs is used to calculate the q_value, the [bs:-self._nstep] means using the data from
# [bs] timestep to [self._unroll_len_add_burnin_step-self._nstep] timestep
data['main_obs'] = data['obs'][bs:-self._nstep]
# the target_obs is used to calculate the target_q_value
data['target_obs'] = data['obs'][bs + self._nstep:]
return data
def _forward_learn(self, data: dict) -> Dict[str, Any]:
r"""
Overview:
Forward and backward function of learn mode.
Acquire the data, calculate the loss and optimize learner model.
Arguments:
- data (:obj:`dict`): Dict type data, including at least \
['main_obs', 'target_obs', 'burnin_obs', 'action', 'reward', 'done', 'weight']
Returns:
- info_dict (:obj:`Dict[str, Any]`): Including cur_lr and total_loss
- cur_lr (:obj:`float`): Current learning rate
- total_loss (:obj:`float`): The calculated loss
"""
# forward
data = self._data_preprocess_learn(data)
self._learn_model.train()
self._target_model.train()
self._learn_model.reset(data_id=None, state=data['prev_state'][0]) # take out timestep=0
self._target_model.reset(data_id=None, state=data['prev_state'][0])
if len(data['burnin_nstep_obs']) != 0:
with torch.no_grad():
inputs = {'obs': data['burnin_nstep_obs'], 'enable_fast_timestep': True}
burnin_output = self._learn_model.forward(
inputs, saved_hidden_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep]
)
burnin_output_target = self._target_model.forward(
inputs, saved_hidden_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep]
)
self._learn_model.reset(data_id=None, state=burnin_output['saved_hidden_state'][0])
inputs = {'obs': data['main_obs'], 'enable_fast_timestep': True}
q_value = self._learn_model.forward(inputs)['logit']
self._learn_model.reset(data_id=None, state=burnin_output['saved_hidden_state'][1])
self._target_model.reset(data_id=None, state=burnin_output_target['saved_hidden_state'][1])
next_inputs = {'obs': data['target_obs'], 'enable_fast_timestep': True}
with torch.no_grad():
target_q_value = self._target_model.forward(next_inputs)['logit']
# argmax_action double_dqn
target_q_action = self._learn_model.forward(next_inputs)['action']
action, reward, done, weight = data['action'], data['reward'], data['done'], data['weight']
# T, B, nstep -> T, nstep, B
reward = reward.permute(0, 2, 1).contiguous()
loss = []
td_error = []
value_gamma = data['value_gamma']
for t in range(self._unroll_len_add_burnin_step - self._burnin_step - self._nstep):
td_data = q_nstep_td_data(
q_value[t], target_q_value[t], action[t], target_q_action[t], reward[t], done[t], weight[t]
)
if self._value_rescale:
l, e = q_nstep_td_error_with_rescale(td_data, self._gamma, self._nstep, value_gamma=value_gamma[t])
loss.append(l)
td_error.append(e.abs())
else:
l, e = q_nstep_td_error(td_data, self._gamma, self._nstep, value_gamma=value_gamma[t])
loss.append(l)
td_error.append(e.abs())
loss = sum(loss) / (len(loss) + 1e-8)
td_error_per_sample = sum(td_error) / (len(td_error) + 1e-8)
# update
self._optimizer.zero_grad()
loss.backward()
self._optimizer.step()
# after update
self._target_model.update(self._learn_model.state_dict())
return {
'cur_lr': self._optimizer.defaults['lr'],
'total_loss': loss.item(),
'priority': td_error_per_sample.abs().tolist(),
}
def _reset_learn(self, data_id: Optional[List[int]] = None) -> None:
self._learn_model.reset(data_id=data_id)
def _state_dict_learn(self) -> Dict[str, Any]:
return {
'model': self._learn_model.state_dict(),
'optimizer': self._optimizer.state_dict(),
}
def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None:
self._learn_model.load_state_dict(state_dict['model'])
self._optimizer.load_state_dict(state_dict['optimizer'])
def _init_collect(self) -> None:
r"""
Overview:
Collect mode init method. Called by ``self.__init__``.
Init traj and unroll length, collect model.
"""
assert 'unroll_len' not in self._cfg.collect, "r2d2 use default unroll_len"
self._nstep = self._cfg.nstep
self._burnin_step = self._cfg.burnin_step
self._gamma = self._cfg.discount_factor
self._unroll_len_add_burnin_step = self._cfg.unroll_len + self._cfg.burnin_step
self._unroll_len = self._unroll_len_add_burnin_step # for compatibility
self._collect_model = model_wrap(
self._model, wrapper_name='hidden_state', state_num=self._cfg.collect.env_num, save_prev_state=True
)
self._collect_model = model_wrap(self._collect_model, wrapper_name='eps_greedy_sample')
self._collect_model.reset()
def _forward_collect(self, data: dict, eps: float) -> dict:
r"""
Overview:
Forward function for collect mode with eps_greedy
Arguments:
- data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \
values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer.
- eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step.
Returns:
- output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs.
ReturnsKeys
- necessary: ``action``
"""
data_id = list(data.keys())
data = default_collate(list(data.values()))
if self._cuda:
data = to_device(data, self._device)
data = {'obs': data}
self._collect_model.eval()
with torch.no_grad():
# in collect phase, inference=True means that each time we only pass one timestep data,
# so the we can get the hidden state of rnn: <prev_state> at each timestep.
output = self._collect_model.forward(data, data_id=data_id, eps=eps, inference=True)
if self._cuda:
output = to_device(output, 'cpu')
output = default_decollate(output)
return {i: d for i, d in zip(data_id, output)}
def _reset_collect(self, data_id: Optional[List[int]] = None) -> None:
self._collect_model.reset(data_id=data_id)
def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict:
r"""
Overview:
Generate dict type transition data from inputs.
Arguments:
- obs (:obj:`Any`): Env observation
- model_output (:obj:`dict`): Output of collect model, including at least ['action', 'prev_state']
- timestep (:obj:`namedtuple`): Output after env step, including at least ['reward', 'done'] \
(here 'obs' indicates obs after env step).
Returns:
- transition (:obj:`dict`): Dict type transition data.
"""
transition = {
'obs': obs,
'action': model_output['action'],
'prev_state': model_output['prev_state'],
'reward': timestep.reward,
'done': timestep.done,
}
return transition
def _get_train_sample(self, data: list) -> Union[None, List[Any]]:
r"""
Overview:
Get the trajectory and the n step return data, then sample from the n_step return data
Arguments:
- data (:obj:`list`): The trajectory's cache
Returns:
- samples (:obj:`dict`): The training samples generated
"""
data = get_nstep_return_data(data, self._nstep, gamma=self._gamma)
return get_train_sample(data, self._unroll_len_add_burnin_step)
def _init_eval(self) -> None:
r"""
Overview:
Evaluate mode init method. Called by ``self.__init__``.
Init eval model with argmax strategy.
"""
self._eval_model = model_wrap(self._model, wrapper_name='hidden_state', state_num=self._cfg.eval.env_num)
self._eval_model = model_wrap(self._eval_model, wrapper_name='argmax_sample')
self._eval_model.reset()
def _forward_eval(self, data: dict) -> dict:
r"""
Overview:
Forward function of eval mode, similar to ``self._forward_collect``.
Arguments:
- data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \
values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer.
Returns:
- output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env.
ReturnsKeys
- necessary: ``action``
"""
data_id = list(data.keys())
data = default_collate(list(data.values()))
if self._cuda:
data = to_device(data, self._device)
data = {'obs': data}
self._eval_model.eval()
with torch.no_grad():
output = self._eval_model.forward(data, data_id=data_id, inference=True)
if self._cuda:
output = to_device(output, 'cpu')
output = default_decollate(output)
return {i: d for i, d in zip(data_id, output)}
def _reset_eval(self, data_id: Optional[List[int]] = None) -> None:
self._eval_model.reset(data_id=data_id)
def default_model(self) -> Tuple[str, List[str]]:
return 'drqn', ['ding.model.template.q_learning']