-
Notifications
You must be signed in to change notification settings - Fork 11
/
hybrid_sac_platform.py
414 lines (361 loc) · 18.9 KB
/
hybrid_sac_platform.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
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
# https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/sac_continuous_action.py
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions.categorical import Categorical
from torch.distributions.normal import Normal
from torch.utils.tensorboard import SummaryWriter
from utils import to_gym_action, to_torch_action, gym_to_buffer
from wrappers.platfrom_wrappers import ScaledStateWrapper, PlatformFlattenedActionWrapper, ScaledParameterisedActionWrapper
import gym_platform
import argparse
from distutils.util import strtobool
import collections
import numpy as np
import gym
from gym.wrappers import TimeLimit, Monitor
# import pybullet_envs
from gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space
import time
import random
import os
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='SAC with 2 Q functions, Online updates')
# Common arguments
parser.add_argument('--exp-name',
type=str,
default=os.path.basename(__file__).rstrip(".py"),
help='the name of this experiment')
parser.add_argument('--gym-id', type=str, default="HopperBulletEnv-v0", help='the id of the gym environment')
parser.add_argument('--learning-rate', type=float, default=7e-4, help='the learning rate of the optimizer')
parser.add_argument('--seed', type=int, default=2, help='seed of the experiment')
parser.add_argument('--episode-length', type=int, default=0, help='the maximum length of each episode')
parser.add_argument('--total-timesteps', type=int, default=4000000, help='total timesteps of the experiments')
parser.add_argument('--torch-deterministic',
type=lambda x: bool(strtobool(x)),
default=True,
nargs='?',
const=True,
help='if toggled, `torch.backends.cudnn.deterministic=False`')
parser.add_argument('--cuda',
type=lambda x: bool(strtobool(x)),
default=True,
nargs='?',
const=True,
help='if toggled, cuda will not be enabled by default')
parser.add_argument('--prod-mode',
type=lambda x: bool(strtobool(x)),
default=False,
nargs='?',
const=True,
help='run the script in production mode and use wandb to log outputs')
parser.add_argument('--capture-video',
type=lambda x: bool(strtobool(x)),
default=False,
nargs='?',
const=True,
help='weather to capture videos of the agent performances (check out `videos` folder)')
parser.add_argument('--wandb-project-name', type=str, default="cleanRL", help="the wandb's project name")
parser.add_argument('--wandb-entity', type=str, default=None, help="the entity (team) of wandb's project")
parser.add_argument('--autotune',
type=lambda x: bool(strtobool(x)),
default=True,
nargs='?',
const=True,
help='automatic tuning of the entropy coefficient.')
# Algorithm specific arguments
parser.add_argument('--buffer-size', type=int, default=10000, help='the replay memory buffer size')
parser.add_argument('--gamma', type=float, default=0.9, help='the discount factor gamma')
parser.add_argument(
'--target-network-frequency',
type=int,
default=1, # Denis Yarats' implementation delays this by 2.
help="the timesteps it takes to update the target network")
parser.add_argument('--max-grad-norm', type=float, default=10.0, help='the maximum norm for the gradient clipping')
parser.add_argument(
'--batch-size',
type=int,
default=128, # Worked better in my experiments, still have to do ablation on this. Please remind me
help="the batch size of sample from the reply memory")
parser.add_argument('--tau', type=float, default=0.1, help="target smoothing coefficient (default: 0.005)")
parser.add_argument('--alpha', type=float, default=0.2, help="Entropy regularization coefficient.")
parser.add_argument('--learning-starts', type=int, default=5e3, help="timestep to start learning")
# Additional hyper parameters for tweaks
## Separating the learning rate of the policy and value commonly seen: (Original implementation, Denis Yarats)
parser.add_argument('--policy-lr',
type=float,
default=1e-4,
help='the learning rate of the policy network optimizer')
parser.add_argument('--q-lr', type=float, default=1e-3, help='the learning rate of the Q network network optimizer')
parser.add_argument('--policy-frequency',
type=int,
default=1,
help='delays the update of the actor, as per the TD3 paper.')
# NN Parameterization
parser.add_argument('--weights-init',
default='xavier',
const='xavier',
nargs='?',
choices=['xavier', "orthogonal", 'uniform'],
help='weight initialization scheme for the neural networks.')
parser.add_argument('--bias-init',
default='zeros',
const='xavier',
nargs='?',
choices=['zeros', 'uniform'],
help='weight initialization scheme for the neural networks.')
parser.add_argument('--ent-c',
default=-0.25,
type=float,
help='target entropy of continuous component.')
parser.add_argument('--ent-d',
default=0.25,
type=float,
help='target entropy of discrete component.')
args = parser.parse_args()
if not args.seed:
args.seed = int(time.time())
# TRY NOT TO MODIFY: setup the environment
experiment_name = f"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
writer = SummaryWriter(f"runs/{experiment_name}")
writer.add_text('hyperparameters',
"|param|value|\n|-|-|\n%s" % ('\n'.join([f"|{key}|{value}|" for key, value in vars(args).items()])))
if args.prod_mode:
import wandb
wandb.init(project=args.wandb_project_name,
entity=args.wandb_entity,
sync_tensorboard=True,
config=vars(args),
name=experiment_name,
monitor_gym=True,
save_code=True)
writer = SummaryWriter(f"/tmp/{experiment_name}")
# TRY NOT TO MODIFY: seeding
device = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')
env = gym.make('Platform-v0')
env = ScaledStateWrapper(env)
env = PlatformFlattenedActionWrapper(env)
env = ScaledParameterisedActionWrapper(env)
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.backends.cudnn.deterministic = args.torch_deterministic
env.seed(args.seed)
env.action_space.seed(args.seed)
env.observation_space.seed(args.seed)
input_shape = 9
out_c = out_d = 3
if args.capture_video:
env = Monitor(env, f'videos/{experiment_name}')
# ALGO LOGIC: initialize agent here:
LOG_STD_MAX = 0.0
LOG_STD_MIN = -3.0
def layer_init(layer, weight_gain=1, bias_const=0):
if isinstance(layer, nn.Linear):
if args.weights_init == "xavier":
torch.nn.init.xavier_uniform_(layer.weight, gain=weight_gain)
elif args.weights_init == "orthogonal":
torch.nn.init.orthogonal_(layer.weight, gain=weight_gain)
if args.bias_init == "zeros":
torch.nn.init.constant_(layer.bias, bias_const)
class Policy(nn.Module):
def __init__(self, input_shape, out_c, out_d, env):
super(Policy, self).__init__()
self.fc1 = nn.Linear(input_shape, 128) # Better result with slightly wider networks.
self.mean = nn.Linear(128, out_c)
self.logstd = nn.Linear(128, out_c)
self.pi_d = nn.Linear(128, out_d)
self.apply(layer_init)
def forward(self, x, device):
x = torch.Tensor(x).to(device)
x = F.relu(self.fc1(x))
mean = torch.tanh(self.mean(x))
log_std = torch.tanh(self.logstd(x))
pi_d = self.pi_d(x)
log_std = LOG_STD_MIN + 0.5 * (LOG_STD_MAX - LOG_STD_MIN) * (log_std + 1) # From SpinUp / Denis Yarats
return mean, log_std, pi_d
def get_action(self, x, device):
mean, log_std, pi_d = self.forward(x, device)
std = log_std.exp()
normal = Normal(mean, std)
x_t = normal.rsample() # for reparameterization trick (mean + std * N(0,1))
action_c = torch.tanh(x_t)
log_prob_c = normal.log_prob(x_t)
log_prob_c -= torch.log(1.0 - action_c.pow(2) + 1e-8)
dist = Categorical(logits=pi_d)
action_d = dist.sample()
prob_d = dist.probs
log_prob_d = torch.log(prob_d + 1e-8)
return action_c, action_d, log_prob_c, log_prob_d, prob_d
def to(self, device):
return super(Policy, self).to(device)
class SoftQNetwork(nn.Module):
def __init__(self, input_shape, out_c, out_d, layer_init):
super(SoftQNetwork, self).__init__()
self.fc1 = nn.Linear(input_shape + out_c, 128)
self.fc2 = nn.Linear(128, out_d)
self.apply(layer_init)
def forward(self, x, a, device):
x = torch.Tensor(x).to(device)
x = torch.cat([x, a], 1)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
# modified from https://github.com/seungeunrho/minimalRL/blob/master/dqn.py#
class ReplayBuffer():
def __init__(self, buffer_limit):
self.buffer = collections.deque(maxlen=buffer_limit)
def put(self, transition):
self.buffer.append(transition)
def sample(self, n):
mini_batch = random.sample(self.buffer, n)
s_lst, a_lst, r_lst, s_prime_lst, done_mask_lst = [], [], [], [], []
for transition in mini_batch:
s, a, r, s_prime, done_mask = transition
s_lst.append(s)
a_lst.append(a)
r_lst.append(r)
s_prime_lst.append(s_prime)
done_mask_lst.append(done_mask)
return np.array(s_lst), np.array(a_lst), \
np.array(r_lst), np.array(s_prime_lst), \
np.array(done_mask_lst)
rb = ReplayBuffer(args.buffer_size)
pg = Policy(input_shape, out_c, out_d, env).to(device)
qf1 = SoftQNetwork(input_shape, out_c, out_d, layer_init).to(device)
qf2 = SoftQNetwork(input_shape, out_c, out_d, layer_init).to(device)
qf1_target = SoftQNetwork(input_shape, out_c, out_d, layer_init).to(device)
qf2_target = SoftQNetwork(input_shape, out_c, out_d, layer_init).to(device)
qf1_target.load_state_dict(qf1.state_dict())
qf2_target.load_state_dict(qf2.state_dict())
values_optimizer = optim.Adam(list(qf1.parameters()) + list(qf2.parameters()), lr=args.q_lr)
policy_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)
loss_fn = nn.MSELoss()
# Automatic entropy tuning
if args.autotune:
# target_entropy = -float(out_c)
target_entropy = args.ent_c
log_alpha = torch.zeros(1, requires_grad=True, device=device)
alpha = log_alpha.exp().detach().cpu().item()
a_optimizer = optim.Adam([log_alpha], lr=1e-4)
# target_entropy_d = -0.98 * np.log(1/out_d)
target_entropy_d = args.ent_d
log_alpha_d = torch.zeros(1, requires_grad=True, device=device)
alpha_d = log_alpha_d.exp().detach().cpu().item()
a_d_optimizer = optim.Adam([log_alpha_d], lr=1e-4)
else:
alpha = args.alpha
alpha_d = args.alpha
# TRY NOT TO MODIFY: start the game
global_episode = 0
(obs, _), done = env.reset(), False
episode_reward, episode_length = 0., 0
for global_step in range(1, args.total_timesteps + 1):
# ALGO LOGIC: put action logic here
if global_step < args.learning_starts:
action_ = env.action_space.sample()
action_ = gym_to_buffer(action_)
action = [action_[0], action_[1:]]
else:
action_c, action_d, _, _, _ = pg.get_action([obs], device)
action = to_gym_action(action_c, action_d)
# TRY NOT TO MODIFY: execute the game and log data.
(next_obs, _), reward, done, _ = env.step(action)
rb.put((obs, gym_to_buffer(action), reward, next_obs, done))
episode_reward += reward
episode_length += 1
obs = np.array(next_obs)
# ALGO LOGIC: training.
if len(rb.buffer) > args.batch_size: # starts update as soon as there is enough data.
s_obs, s_actions, s_rewards, s_next_obses, s_dones = rb.sample(args.batch_size)
with torch.no_grad():
next_state_actions_c, next_state_actions_d, next_state_log_pi_c, next_state_log_pi_d, next_state_prob_d = pg.get_action(s_next_obses, device)
qf1_next_target = qf1_target.forward(s_next_obses, next_state_actions_c, device)
qf2_next_target = qf2_target.forward(s_next_obses, next_state_actions_c, device)
min_qf_next_target = next_state_prob_d * (torch.min(qf1_next_target, qf2_next_target) - alpha * next_state_prob_d * next_state_log_pi_c - alpha_d * next_state_log_pi_d)
next_q_value = torch.Tensor(s_rewards).to(
device) + (1 - torch.Tensor(s_dones).to(device)) * args.gamma * (min_qf_next_target.sum(1)).view(-1)
s_actions_c, s_actions_d = to_torch_action(s_actions, device)
qf1_a_values = qf1.forward(s_obs, s_actions_c, device).gather(1, s_actions_d.long().view(-1, 1).to(device)).squeeze().view(-1)
qf2_a_values = qf2.forward(s_obs, s_actions_c, device).gather(1, s_actions_d.long().view(-1, 1).to(device)).squeeze().view(-1)
qf1_loss = loss_fn(qf1_a_values, next_q_value)
qf2_loss = loss_fn(qf2_a_values, next_q_value)
qf_loss = (qf1_loss + qf2_loss) / 2
values_optimizer.zero_grad()
qf_loss.backward()
values_optimizer.step()
if global_step % args.policy_frequency == 0: # TD 3 Delayed update support
for _ in range(
args.policy_frequency): # compensate for the delay by doing 'actor_update_interval' instead of 1
actions_c, actions_d, log_pi_c, log_pi_d, prob_d = pg.get_action(s_obs, device)
qf1_pi = qf1.forward(s_obs, actions_c, device)
qf2_pi = qf2.forward(s_obs, actions_c, device)
min_qf_pi = torch.min(qf1_pi, qf2_pi)
policy_loss_d = (prob_d * (alpha_d * log_pi_d - min_qf_pi)).sum(1).mean()
policy_loss_c = (prob_d * (alpha * prob_d * log_pi_c - min_qf_pi)).sum(1).mean()
policy_loss = policy_loss_d + policy_loss_c
policy_optimizer.zero_grad()
policy_loss.backward()
policy_optimizer.step()
if args.autotune:
with torch.no_grad():
a_c, a_d, lpi_c, lpi_d, p_d = pg.get_action(s_obs, device)
alpha_loss = (-log_alpha * p_d * (p_d * lpi_c + target_entropy)).sum(1).mean()
alpha_d_loss = (-log_alpha_d * p_d * (lpi_d + target_entropy_d)).sum(1).mean()
a_optimizer.zero_grad()
alpha_loss.backward()
a_optimizer.step()
alpha = log_alpha.exp().detach().cpu().item()
a_d_optimizer.zero_grad()
alpha_d_loss.backward()
a_d_optimizer.step()
alpha_d = log_alpha_d.exp().detach().cpu().item()
# update the target network
if global_step % args.target_network_frequency == 0:
for param, target_param in zip(qf1.parameters(), qf1_target.parameters()):
target_param.data.copy_(args.tau * param.data + (1 - args.tau) * target_param.data)
for param, target_param in zip(qf2.parameters(), qf2_target.parameters()):
target_param.data.copy_(args.tau * param.data + (1 - args.tau) * target_param.data)
if len(rb.buffer) > args.batch_size and global_step % 100 == 0:
writer.add_scalar("losses/soft_q_value_1_loss", qf1_loss.item(), global_step)
writer.add_scalar("losses/soft_q_value_2_loss", qf2_loss.item(), global_step)
writer.add_scalar("losses/qf_loss", qf_loss.item(), global_step)
writer.add_scalar("losses/policy_loss", policy_loss.item(), global_step)
writer.add_scalar("losses/alpha", alpha, global_step)
# NOTE: additional changes from cleanrl
writer.add_scalar("losses/alpha_d", alpha_d, global_step)
writer.add_histogram("actions/discrete", action[0]+1, global_step)
writer.add_histogram("actions/run", action[1][0], global_step)
writer.add_histogram("actions/hop", action[1][1], global_step)
writer.add_histogram("actions/leap", action[1][2], global_step)
writer.add_scalar("debug/ent_bonus", (- alpha * next_state_prob_d * next_state_log_pi_c - alpha_d * next_state_log_pi_d).sum(1).mean().item(), global_step)
writer.add_scalar("debug/next_q", next_q_value.mean().item(), global_step)
writer.add_scalar("debug/policy_loss_c", policy_loss_c.item(), global_step)
writer.add_scalar("debug/policy_loss_d", policy_loss_d.item(), global_step)
writer.add_scalar("debug/policy_ent_d", -(prob_d*log_pi_d).sum(1).mean().item(), global_step)
writer.add_scalar("debug/mean_q", min_qf_pi.mean().item(), global_step)
writer.add_scalar("debug/mean_r", s_rewards.mean(), global_step)
writer.add_histogram("debug_q/q_0", min_qf_pi[:, 0].mean().item(), global_step)
writer.add_histogram("debug_q/q_1", min_qf_pi[:, 1].mean().item(), global_step)
writer.add_histogram("debug_q/q_2", min_qf_pi[:, 2].mean().item(), global_step)
writer.add_histogram("debug_pi/pi_0", prob_d[:, 0].mean().item(), global_step)
writer.add_histogram("debug_pi/pi_1", prob_d[:, 1].mean().item(), global_step)
writer.add_histogram("debug_pi/pi_2", prob_d[:, 2].mean().item(), global_step)
if args.autotune:
writer.add_scalar("losses/alpha_loss", alpha_loss.item(), global_step)
writer.add_scalar("losses/alpha_d_loss", alpha_d_loss.item(), global_step)
if done:
global_episode += 1 # Outside the loop already means the epsiode is done
writer.add_scalar("charts/episode_reward", episode_reward, global_step)
writer.add_scalar("charts/episode_length", episode_length, global_step)
# Terminal verbosity
if global_episode % 10 == 0:
print(f"Episode: {global_episode} Step: {global_step}, Ep. Reward: {episode_reward}")
# Reseting what need to be
(obs, _), done = env.reset(), False
episode_reward, episode_length = 0., 0
writer.close()
env.close()
# save the policy
torch.save(pg.state_dict(), 'platform.pth')