forked from djole/IR2L
-
Notifications
You must be signed in to change notification settings - Fork 0
/
policy_training_loop_no_instinct.py
172 lines (141 loc) · 6.43 KB
/
policy_training_loop_no_instinct.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
from math import log
import pickle
import torch
import numpy as np
import gym
from gym.envs.registration import register
from gym.utils import seeding
import safety_gym_mod
from a2c_ppo_acktr import algo, utils
from a2c_ppo_acktr.envs import make_vec_envs
from a2c_ppo_acktr.evaluation import evaluate
from a2c_ppo_acktr.model import init_default_ppo, Policy, custom_weight_init
from a2c_ppo_acktr.storage import RolloutStorage
from arguments import get_args
try:
from exp_dir_util import get_experiment_save_dir
except:
pass
from os.path import join
from torch.utils.tensorboard import SummaryWriter
from double_rl_loop_main import ENV_NAME_BOX, NUM_PROC, reward_cost_combinator, HAZARD_PUNISHMENT, REWARD_SCALE,\
EPISODE_LENGTH
def instinct_loop_ppo(
args,
learning_rate,
num_steps,
num_updates,
inst_on,
visualize,
save_dir
):
torch.set_num_threads(1)
log_writer = SummaryWriter(save_dir, max_queue=1, filename_suffix="log")
device = torch.device("cpu")
env_name = ENV_NAME_BOX #"Safexp-PointGoal1-v0"
envs = make_vec_envs(env_name, np.random.randint(2 ** 32), NUM_PROC,
args.gamma, None, device, allow_early_resets=True, normalize=args.norm_vectors)
eval_envs = make_vec_envs(env_name, np.random.randint(2 ** 32), 1,
args.gamma, None, device, allow_early_resets=True, normalize=args.norm_vectors)
actor_critic_policy = torch.load("pretrained_baseline.pt") # init_default_ppo(envs, log(args.init_sigma))
# Prepare modified observation shape for instinct
obs_shape = envs.observation_space.shape
actor_critic_policy.to(device)
agent_policy = algo.PPO(
actor_critic_policy,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=learning_rate,
eps=args.eps,
max_grad_norm=args.max_grad_norm)
rollouts = RolloutStorage(num_steps, NUM_PROC,
obs_shape, envs.action_space,
actor_critic_policy.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
fitnesses = []
best_fitness_so_far = float("-Inf")
masks = torch.ones(num_steps + 1, NUM_PROC, 1)
for j in range(num_updates):
training_collisions_current_update = 0
for step in range(num_steps):
# Sample actions
with torch.no_grad():
# (value, action, action_log_probs, rnn_hxs), (instinct_value, instinct_action, instinct_outputs_log_prob, i_rnn_hxs), final_action
value, action, action_log_probs, recurrent_hidden_states = actor_critic_policy.act(
rollouts.obs[step],
rollouts.recurrent_hidden_states[step],
rollouts.masks[step],
deterministic=False
)
# Combine two networks
obs, reward, done, infos = envs.step(action)
# envs.render()
training_collisions_current_update += sum([i['cost'] for i in infos])
modded_reward, _ = reward_cost_combinator(reward, infos, NUM_PROC, torch.tensor([[0.0]] * NUM_PROC))
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])
bad_masks = torch.FloatTensor([[0.0] if 'bad_transition' in info.keys() else [1.0] for info in infos])
rollouts.insert(obs, recurrent_hidden_states, action, action_log_probs,
value, modded_reward, masks, bad_masks)
with torch.no_grad():
next_value_policy = actor_critic_policy.get_value(rollouts.obs[-1],
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1].detach())
rollouts.compute_returns(next_value_policy, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
print("training policy")
# Instinct training phase
val_loss, action_loss, dist_entropy = agent_policy.update(rollouts)
rollouts.after_update()
ob_rms = utils.get_vec_normalize(envs)
if ob_rms is not None:
ob_rms = ob_rms.ob_rms
fits, info = evaluate(actor_critic_policy, ob_rms, eval_envs, NUM_PROC,
reward_cost_combinator, device, instinct_on=inst_on, visualise=visualize)
instinct_reward = info['instinct_reward']
hazard_collisions = info['hazard_collisions']
print(
f"Step {j}, Fitness {fits.item()}, value_loss instinct = {val_loss}, action_loss instinct= {action_loss}, "
f"dist_entropy instinct = {dist_entropy}")
print(
f"Step {j}, Cost {instinct_reward}")
print("-----------------------------------------------------------------")
# Tensorboard logging
log_writer.add_scalar("Task reward", fits.item(), j)
log_writer.add_scalar("cost/Training hazard collisions", training_collisions_current_update, j)
log_writer.add_scalar("cost/Instinct reward", instinct_reward, j)
log_writer.add_scalar("cost/Eval hazard collisions", hazard_collisions, j)
log_writer.add_scalar("value loss", val_loss, j)
log_writer.add_scalar("action loss", action_loss, j)
log_writer.add_scalar("dist entropy", dist_entropy, j)
fitnesses.append(fits)
if fits.item() > best_fitness_so_far:
best_fitness_so_far = fits.item()
torch.save(actor_critic_policy, join(save_dir, "model_rl_policy_best.pt"))
torch.save(actor_critic_policy, join(save_dir, f"model_rl_policy_latest_{j}.pt"))
torch.save(actor_critic_policy, join(save_dir, "model_rl_policy_latest.pt"))
pickle.dump(ob_rms, open(join(save_dir, "ob_rms.p"), "wb"))
return (fitnesses[-1]), 0, 0
def main():
args = get_args()
print("start the train function")
args.init_sigma = 0.6
args.lr = 0.001
# plot_weight_histogram(parameters)
exp_save_dir = get_experiment_save_dir(args)
instinct_loop_ppo(
args,
args.lr,
num_steps=9000,
num_updates=300,
inst_on=False,
visualize=False,
save_dir=exp_save_dir
)
if __name__ == "__main__":
main()