-
Notifications
You must be signed in to change notification settings - Fork 2
/
task_mujoco_visual_reacher.py
187 lines (156 loc) · 6.81 KB
/
task_mujoco_visual_reacher.py
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
import torch
import argparse
import time
import os
import relod.utils as utils
from relod.algo.local_wrapper import LocalWrapper
from relod.algo.sac_rad_agent import SACRADPerformer, SACRADLearner
from relod.envs.dm_reacher import ReacherWrapper
from relod.algo.comm import MODE
from relod.logger import Logger
config = {
'conv': [
# in_channel, out_channel, kernel_size, stride
[-1, 32, 3, 2],
[32, 32, 3, 2],
[32, 32, 3, 2],
[32, 32, 3, 1],
],
'latent': 50,
'mlp': [
[-1, 1024],
[1024, 1024],
[1024, -1]
],
}
def parse_args():
parser = argparse.ArgumentParser()
# environment
parser.add_argument('--target_type', default='visual_reacher', type=str)
parser.add_argument('--image_height', default=125, type=int)
parser.add_argument('--image_width', default=200, type=int)
parser.add_argument('--stack_frames', default=3, type=int)
parser.add_argument('--tol', default=0.036, type=float)
parser.add_argument('--image_period', default=1, type=int)
parser.add_argument('--episode_length_time', default=50, type=int)
parser.add_argument('--dt', default=1, type=int)
# replay buffer
parser.add_argument('--replay_buffer_capacity', default=10000, type=int)
parser.add_argument('--rad_offset', default=0.01, type=float)
# train
parser.add_argument('--init_steps', default=1000, type=int)
parser.add_argument('--env_steps', default=25000, type=int)
parser.add_argument('--batch_size', default=256, type=int)
parser.add_argument('--async_mode', default=True, action='store_true')
parser.add_argument('--async_buffer', default=False, action='store_true')
parser.add_argument('--max_updates_per_step', default=1, type=float)
parser.add_argument('--update_every', default=50, type=int)
parser.add_argument('--update_epochs', default=50, type=int)
# critic
parser.add_argument('--critic_lr', default=1e-3, type=float)
parser.add_argument('--critic_tau', default=0.01, type=float)
parser.add_argument('--critic_target_update_freq', default=1, type=int)
parser.add_argument('--bootstrap_terminal', default=0, type=int)
# actor
parser.add_argument('--actor_lr', default=1e-3, type=float)
parser.add_argument('--actor_update_freq', default=1, type=int)
# encoder
parser.add_argument('--encoder_tau', default=0.05, type=float)
# sac
parser.add_argument('--discount', default=1., type=float)
parser.add_argument('--init_temperature', default=0.1, type=float)
parser.add_argument('--alpha_lr', default=1e-4, type=float)
# agent
parser.add_argument('--remote_ip', default='localhost', type=str)
parser.add_argument('--port', default=9876, type=int)
parser.add_argument('--mode', default='l', type=str, help="Modes in ['r', 'l', 'rl'] ")
# misc
parser.add_argument('--seed', default=0, type=int)
parser.add_argument('--work_dir', default='.', type=str)
parser.add_argument('--save_tb', default=False, action='store_true')
parser.add_argument('--save_model', default=False, action='store_true')
parser.add_argument('--save_model_freq', default=1000, type=int)
parser.add_argument('--load_model', default=-1, type=int)
parser.add_argument('--device', default='cuda:0', type=str)
parser.add_argument('--lock', default=False, action='store_true')
parser.add_argument('--save_path', default='', type=str, help="For saving SAC buffer")
parser.add_argument('--load_path', default='', type=str, help="Path to SAC buffer file")
args = parser.parse_args()
return args
def main():
args = parse_args()
if args.mode == 'r':
mode = MODE.REMOTE_ONLY
elif args.mode == 'l':
mode = MODE.LOCAL_ONLY
elif args.mode == 'rl':
mode = MODE.REMOTE_LOCAL
else:
raise NotImplementedError()
if args.device == '':
args.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
if not 'conv' in config:
image_shape = (0, 0, 0)
else:
image_shape = (3*args.stack_frames, args.image_height, args.image_width)
args.work_dir += f'/results/{args.target_type}_' \
f'seed={args.seed}_' \
f'tol={args.tol}/'
utils.make_dir(args.work_dir)
model_dir = utils.make_dir(os.path.join(args.work_dir, 'model'))
args.model_dir = model_dir
L = Logger(args.work_dir, use_tb=args.save_tb)
env = ReacherWrapper(mode="hard", use_image=True, seed=args.seed, penalty=-1)
utils.set_seed_everywhere(args.seed, env)
args.image_shape = env.image_space.shape
args.proprioception_shape = env.proprioception_space.shape
args.action_shape = env.action_space.shape
args.net_params = config
args.env_action_space = env.action_space
episode_length_step = int(args.episode_length_time / args.dt)
agent = LocalWrapper(episode_length_step, mode, remote_ip=args.remote_ip, port=args.port)
agent.send_data(args)
agent.init_performer(SACRADPerformer, args)
agent.init_learner(SACRADLearner, args, agent.performer)
# sync initial weights with remote
agent.apply_remote_policy(block=True)
episode, episode_reward, episode_step, done = 0, 0, 0, True
obs = env.reset()
agent.send_init_ob((obs.images, obs.proprioception))
start_time = time.time()
for step in range(args.env_steps + args.init_steps):
action = agent.sample_action((obs.images, obs.proprioception))
next_obs, reward, done, _ = env.step(action)
episode_reward += reward
episode_step += 1
agent.push_sample((obs.images, obs.proprioception), action, reward, (next_obs.images, next_obs.proprioception), done)
if done or (episode_step == episode_length_step): # set time out here
if mode == MODE.LOCAL_ONLY:
L.log('train/duration', time.time() - start_time, step)
L.log('train/episode_reward', episode_reward, step)
L.dump(step)
L.log('train/episode', episode+1, step)
next_obs = env.reset()
agent.send_init_ob((next_obs.images, next_obs.proprioception))
episode_reward = 0
episode_step = 0
episode += 1
start_time = time.time()
stat = agent.update_policy(step)
if stat is not None:
for k, v in stat.items():
L.log(k, v, step)
obs = next_obs
if args.save_model and (step+1) % args.save_model_freq == 0:
agent.save_policy_to_file(args.model_dir, step)
agent.save_buffer()
time.sleep(0.04)
if args.save_model:
agent.save_policy_to_file(args.model_dir, step)
# Clean up
agent.close()
env.close()
print('Train finished')
if __name__ == '__main__':
torch.multiprocessing.set_start_method('spawn')
main()