-
Notifications
You must be signed in to change notification settings - Fork 14
/
PushBlock_ppo_icm.py
151 lines (124 loc) · 5.54 KB
/
PushBlock_ppo_icm.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
from envs import *
from utils import *
from config import *
from ppo_agent import *
from torch.multiprocessing import Pipe
from tensorboardX import SummaryWriter
from mlagents.envs import UnityEnvironment
import numpy as np
import time
env = UnityEnvironment(file_name='pushblock/pushblock')
default_brain = env.brain_names[0]
brain = env.brains[default_brain]
env.reset()
num_worker = 32
input_size = 210
output_size = 5
num_step = 256
gamma = 0.99
pre_obs_norm_step = 10000
reward_rms = RunningMeanStd()
obs_rms = RunningMeanStd(1, input_size)
discounted_reward = RewardForwardFilter(gamma)
agent = MlpICMAgent(input_size, output_size, num_worker,
num_step, gamma, use_cuda=True)
steps = 0
next_obs = []
print('Start to initialize observation normalization ...')
while steps < pre_obs_norm_step:
steps += num_worker
actions = np.random.randint(output_size, size=num_worker)
env_info = env.step(actions)[default_brain]
obs = env_info.vector_observations
for o in obs:
next_obs.append(o)
print('initializing...:', steps, '/', pre_obs_norm_step)
next_obs = np.stack(next_obs)
obs_rms.update(next_obs)
print('End to initialize')
writer = SummaryWriter()
writer_iter = 2000
global_update = 0
global_step = 0
sample_i_rall = 0
sample_episode = 0
sample_env_idx = 0
sample_rall = 0
states = np.zeros([num_worker, input_size])
int_coef = 0.5
large_scale_version = True
while True:
total_state, total_reward, total_done, total_next_state, \
total_action, total_int_reward, total_next_obs, total_values,\
total_policy, total_combine_reward = [], [], [], [], [], [], [], [], [], []
global_step += (num_step * num_worker)
global_update += 1
for _ in range(num_step):
actions, value, policy = agent.get_action((np.float32(states) - obs_rms.mean)/np.sqrt(obs_rms.var))
env_info = env.step(actions)[default_brain]
next_states, rewards, dones, real_dones, next_obs = [], [], [], [], []
obs = env_info.vector_observations
reward = env_info.rewards
reward = np.clip(reward, 0, 1)
done = env_info.local_done
for o, r, d in zip(obs, reward, done):
next_states.append(o)
rewards.append(r)
dones.append(d)
next_states = np.stack(next_states)
rewards = np.hstack(rewards)
dones = np.hstack(dones)
intrinsic_reward = agent.compute_intrinsic_reward(
(states - obs_rms.mean)/np.sqrt(obs_rms.var),
(next_states - obs_rms.mean)/np.sqrt(obs_rms.var),
actions)
intrinsic_reward = np.hstack(intrinsic_reward)
combine_reward = (1-int_coef) * rewards + int_coef * intrinsic_reward
sample_i_rall += intrinsic_reward[sample_env_idx]
sample_rall += rewards[sample_env_idx]
total_combine_reward.append(combine_reward)
total_int_reward.append(intrinsic_reward)
total_state.append(states)
total_next_state.append(next_states)
total_reward.append(rewards)
total_done.append(dones)
total_action.append(actions)
total_values.append(value)
total_policy.append(policy)
states = next_states[:, :]
if dones[sample_env_idx]:
sample_episode += 1
if sample_episode < writer_iter:
writer.add_scalar('data/reward_per_epi', sample_rall, sample_episode)
writer.add_scalar('data/int_reward_per_epi', sample_i_rall, sample_episode)
print("[Episode {}] rall: {} i_: {}".format(
sample_episode, sample_rall, sample_i_rall))
sample_i_rall = 0
sample_rall = 0
_, value, _ = agent.get_action((np.float32(states) - obs_rms.mean) / np.sqrt(obs_rms.var))
total_values.append(value)
total_state = np.stack(total_state).transpose([1, 0, 2]).reshape([-1, input_size])
total_next_state = np.stack(total_next_state).transpose([1, 0, 2]).reshape([-1, input_size])
total_action = np.stack(total_action).transpose().reshape([-1])
total_reward = np.stack(total_reward).transpose()
total_done = np.stack(total_done).transpose()
total_values = np.stack(total_values).transpose()
total_logging_policy = np.vstack(total_policy)
total_combine_reward = np.stack(total_combine_reward).transpose()
total_reward_per_env = np.array([discounted_reward.update(reward_per_step) for reward_per_step in
total_combine_reward.T])
mean, std, count = np.mean(total_reward_per_env), np.std(total_reward_per_env), len(total_combine_reward)
reward_rms.update_from_moments(mean, std ** 2, count)
total_combine_reward /= np.sqrt(reward_rms.var)
writer.add_scalar('data/int_reward_per_epi', np.sum(total_combine_reward)/num_worker, sample_episode)
writer.add_scalar('data/int_reward_per_rollout', np.sum(total_combine_reward) / num_worker, global_update)
writer.add_scalar('data/max_prob', softmax(total_logging_policy).max(1).mean(), sample_episode)
if large_scale_version: flag = np.zeros_like(total_combine_reward)
else: flag = total_done
target ,adv = make_train_data_icm(total_combine_reward, flag, total_values, gamma, num_step, num_worker)
adv = (adv - np.mean(adv)) / (np.std(adv) + 1e-8)
print('training')
agent.train_model((np.float32(total_state) - obs_rms.mean )/ np.sqrt(obs_rms.var),
(np.float32(total_next_state) - obs_rms.mean) / np.sqrt(obs_rms.var),
target, total_action,
adv, total_policy)