-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathtest_serial_entry_il.py
More file actions
executable file
·143 lines (120 loc) · 5.78 KB
/
Copy pathtest_serial_entry_il.py
File metadata and controls
executable file
·143 lines (120 loc) · 5.78 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
from copy import deepcopy
import pytest
import torch.nn.functional as F
from typing import Tuple, List, Dict, Any
import torch
from collections import namedtuple
import os
from ding.torch_utils import Adam, to_device
from ding.config import compile_config
from ding.model import model_wrap
from ding.rl_utils import get_train_sample, get_nstep_return_data
from ding.entry import serial_pipeline_il, collect_demo_data, serial_pipeline
from ding.policy import PPOOffPolicy, ILPolicy
from ding.policy.common_utils import default_preprocess_learn
from ding.utils import POLICY_REGISTRY
from ding.utils.data import default_collate, default_decollate
from dizoo.classic_control.cartpole.config import cartpole_dqn_config, cartpole_dqn_create_config, \
cartpole_ppo_offpolicy_config, cartpole_ppo_offpolicy_create_config
@POLICY_REGISTRY.register('ppo_il')
class PPOILPolicy(PPOOffPolicy):
def _forward_learn(self, data: dict) -> dict:
data = default_preprocess_learn(data, ignore_done=self._cfg.learn.get('ignore_done', False), use_nstep=False)
self._learn_model.train()
output = self._learn_model.forward(data['obs'], mode='compute_actor_critic')
value_loss = F.mse_loss(output['value'], data['value'])
policy_loss = F.smooth_l1_loss(output['logit'], data['logit'])
total_loss = value_loss + policy_loss
self._optimizer.zero_grad()
total_loss.backward()
self._optimizer.step()
return {
'cur_lr': self._optimizer.defaults['lr'],
'total_loss': total_loss.item(),
'policy_loss': policy_loss.item(),
'value_loss': value_loss.item(),
}
def _monitor_vars_learn(self) -> list:
return super()._monitor_vars_learn() + ['policy_loss', 'value_loss']
@pytest.mark.unittest
def test_serial_pipeline_il_ppo():
# train expert policy
train_config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)]
expert_policy = serial_pipeline(train_config, seed=0)
# collect expert demo data
collect_count = 10000
expert_data_path = 'expert_data_ppo.pkl'
state_dict = expert_policy.collect_mode.state_dict()
collect_config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)]
collect_demo_data(
collect_config, seed=0, state_dict=state_dict, expert_data_path=expert_data_path, collect_count=collect_count
)
# il training 1
il_config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)]
il_config[0].policy.learn.train_epoch = 20
il_config[0].policy.type = 'ppo_il'
_, converge_stop_flag = serial_pipeline_il(il_config, seed=314, data_path=expert_data_path)
assert converge_stop_flag
os.popen('rm -rf ' + expert_data_path)
@POLICY_REGISTRY.register('dqn_il')
class DQNILPolicy(ILPolicy):
def _forward_learn(self, data: dict) -> dict:
for d in data:
if isinstance(d['obs'], torch.Tensor):
d['obs'] = {'processed_obs': d['obs']}
else:
assert 'processed_obs' in d['obs']
return super()._forward_learn(data)
def _init_collect(self) -> None:
self._unroll_len = self._cfg.collect.unroll_len
self._gamma = self._cfg.discount_factor # necessary for parallel
self._nstep = self._cfg.nstep # necessary for parallel
self._collect_model = model_wrap(self._model, wrapper_name='argmax_sample')
self._collect_model.reset()
def _forward_collect(self, data: dict):
data_id = list(data.keys())
data = default_collate(list(data.values()))
if self._cuda:
data = to_device(data, self._device)
self._collect_model.eval()
with torch.no_grad():
output = self._collect_model.forward(data)
if self._cuda:
output = to_device(output, 'cpu')
output = default_decollate(output)
return {i: d for i, d in zip(data_id, output)}
def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> Dict[str, Any]:
ret = super()._process_transition(obs, model_output, timestep)
ret['next_obs'] = timestep.obs
return ret
def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
super()._get_train_sample(data)
data = get_nstep_return_data(data, self._nstep, gamma=self._gamma)
return get_train_sample(data, unroll_len=self._unroll_len)
def _forward_eval(self, data: dict) -> dict:
new_data = {id: {'obs': {'processed_obs': t}} for id, t in data.items()}
return super()._forward_eval(new_data)
def default_model(self) -> Tuple[str, List[str]]:
return 'dqn', ['ding.model.template.q_learning']
@pytest.mark.unittest
def test_serial_pipeline_il_dqn():
# train expert policy
train_config = [deepcopy(cartpole_dqn_config), deepcopy(cartpole_dqn_create_config)]
expert_policy = serial_pipeline(train_config, seed=0)
# collect expert demo data
collect_count = 10000
expert_data_path = 'expert_data_dqn.pkl'
state_dict = expert_policy.collect_mode.state_dict()
collect_config = [deepcopy(cartpole_dqn_config), deepcopy(cartpole_dqn_create_config)]
collect_config[0].policy.type = 'dqn_il'
collect_demo_data(
collect_config, seed=0, state_dict=state_dict, expert_data_path=expert_data_path, collect_count=collect_count
)
# il training 2
il_config = [deepcopy(cartpole_dqn_config), deepcopy(cartpole_dqn_create_config)]
il_config[0].policy.learn.train_epoch = 15
il_config[0].policy.type = 'dqn_il'
il_config[0].env.stop_value = 50
_, converge_stop_flag = serial_pipeline_il(il_config, seed=314, data_path=expert_data_path)
assert converge_stop_flag
os.popen('rm -rf ' + expert_data_path)