def train(model, dataset, grad_func, lr, kfac, num_epochs):
batch_size = 100
if kfac:
optim = KFACOptimizer(model)
optim.acc_stats = True
else:
optim = torch.optim.SGD(model.parameters(), lr)
train_loss = np.zeros(num_epochs)
train_acc = np.zeros(num_epochs)
test_acc = np.zeros(num_epochs)
times = np.zeros(num_epochs)
t = time.time()
for epoch in range(num_epochs):
dataset.reset_and_shuffle(batch_size)
losses = np.zeros(len(dataset))
for i in range(len(dataset)):
x, y = dataset.next_batch()
x = torch.autograd.Variable(torch.from_numpy(x)).cuda()
y = torch.autograd.Variable(torch.from_numpy(y).long()).cuda()
logits = model(x)
loss = nn.functional.cross_entropy(logits, y)
losses[i] = loss.data[0]
grad_func(loss, model)
optim.step()
optim.zero_grad()
train_loss[epoch] = losses.mean()
train_acc[epoch] = eval_(model, dataset.train_xs, dataset.train_ys)
test_acc[epoch] = eval_(model, dataset.test_xs, dataset.test_ys)
times[epoch] = time.time() - t
print 'epoch: %i, loss: %.4f' % (epoch + 1, train_loss[epoch])
print 'accumulate time', times[epoch]
print 'train acc:', train_acc[epoch]
print 'eval acc:', test_acc[epoch]
print '----------------'
return train_loss, train_acc, test_acc, times
def main():
print("#######")
print(
"WARNING: All rewards are clipped or normalized so you need to use a monitor (see envs.py) or visdom plot to get true rewards"
)
print("#######")
os.environ['OMP_NUM_THREADS'] = '1'
if args.vis:
from visdom import Visdom
viz = Visdom()
win = None
envs = [
make_env(args.env_name, args.seed, i, args.log_dir)
for i in range(args.num_processes)
]
if args.num_processes > 1:
envs = SubprocVecEnv(envs)
else:
envs = DummyVecEnv(envs)
if len(envs.observation_space.shape) == 1:
envs = VecNormalize(envs)
obs_shape = envs.observation_space.shape
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])
obs_numel = reduce(operator.mul, obs_shape, 1)
if len(obs_shape) == 3 and obs_numel > 1024:
actor_critic = CNNPolicy(obs_shape[0], envs.action_space,
args.recurrent_policy)
else:
assert not args.recurrent_policy, \
"Recurrent policy is not implemented for the MLP controller"
actor_critic = MLPPolicy(obs_numel, envs.action_space)
if envs.action_space.__class__.__name__ == "Discrete":
action_shape = 1
else:
action_shape = envs.action_space.shape[0]
if args.cuda:
actor_critic.cuda()
if args.algo == 'a2c':
optimizer = optim.RMSprop(actor_critic.parameters(),
args.lr,
eps=args.eps,
alpha=args.alpha)
elif args.algo == 'ppo':
optimizer = optim.Adam(actor_critic.parameters(),
args.lr,
eps=args.eps)
elif args.algo == 'acktr':
optimizer = KFACOptimizer(actor_critic)
rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape,
envs.action_space, actor_critic.state_size)
current_obs = torch.zeros(args.num_processes, *obs_shape)
def update_current_obs(obs):
shape_dim0 = envs.observation_space.shape[0]
obs = torch.from_numpy(obs).float()
if args.num_stack > 1:
current_obs[:, :-shape_dim0] = current_obs[:, shape_dim0:]
current_obs[:, -shape_dim0:] = obs
obs = envs.reset()
update_current_obs(obs)
rollouts.observations[0].copy_(current_obs)
# These variables are used to compute average rewards for all processes.
episode_rewards = torch.zeros([args.num_processes, 1])
final_rewards = torch.zeros([args.num_processes, 1])
if args.cuda:
current_obs = current_obs.cuda()
rollouts.cuda()
start = time.time()
for j in range(num_updates):
for step in range(args.num_steps):
# Sample actions
value, action, action_log_prob, states = actor_critic.act(
Variable(rollouts.observations[step], volatile=True),
Variable(rollouts.states[step], volatile=True),
Variable(rollouts.masks[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy()
# Obser reward and next obs
obs, reward, done, info = envs.step(cpu_actions)
reward = torch.from_numpy(np.expand_dims(np.stack(reward),
1)).float()
episode_rewards += reward
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
if args.cuda:
masks = masks.cuda()
if current_obs.dim() == 4:
current_obs *= masks.unsqueeze(2).unsqueeze(2)
else:
current_obs *= masks
update_current_obs(obs)
rollouts.insert(step, current_obs, states.data, action.data,
action_log_prob.data, value.data, reward, masks)
next_value = actor_critic(
Variable(rollouts.observations[-1], volatile=True),
Variable(rollouts.states[-1], volatile=True),
Variable(rollouts.masks[-1], volatile=True))[0].data
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
if args.algo in ['a2c', 'acktr']:
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(rollouts.observations[:-1].view(-1, *obs_shape)),
Variable(rollouts.states[0].view(-1, actor_critic.state_size)),
Variable(rollouts.masks[:-1].view(-1, 1)),
Variable(rollouts.actions.view(-1, action_shape)))
values = values.view(args.num_steps, args.num_processes, 1)
action_log_probs = action_log_probs.view(args.num_steps,
args.num_processes, 1)
advantages = Variable(rollouts.returns[:-1]) - values
value_loss = advantages.pow(2).mean()
action_loss = -(Variable(advantages.data) *
action_log_probs).mean()
if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(values -
Variable(sample_values.data)).pow(2).mean()
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss -
dist_entropy * args.entropy_coef).backward()
if args.algo == 'a2c':
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
elif args.algo == 'ppo':
advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-5)
for e in range(args.ppo_epoch):
if args.recurrent_policy:
data_generator = rollouts.recurrent_generator(
advantages, args.num_mini_batch)
else:
data_generator = rollouts.feed_forward_generator(
advantages, args.num_mini_batch)
for sample in data_generator:
observations_batch, states_batch, actions_batch, \
return_batch, masks_batch, old_action_log_probs_batch, \
adv_targ = sample
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(observations_batch), Variable(states_batch),
Variable(masks_batch), Variable(actions_batch))
adv_targ = Variable(adv_targ)
ratio = torch.exp(action_log_probs -
Variable(old_action_log_probs_batch))
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - args.clip_param,
1.0 + args.clip_param) * adv_targ
action_loss = -torch.min(
surr1,
surr2).mean() # PPO's pessimistic surrogate (L^CLIP)
value_loss = (Variable(return_batch) -
values).pow(2).mean()
optimizer.zero_grad()
(value_loss + action_loss -
dist_entropy * args.entropy_coef).backward()
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
rollouts.after_update()
if j % args.save_interval == 0 and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
hasattr(envs, 'ob_rms') and envs.ob_rms or None
]
torch.save(save_model,
os.path.join(save_path, args.env_name + ".pt"))
if j % args.log_interval == 0:
end = time.time()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
print(
"Updates {}, num timesteps {}, FPS {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
final_rewards.mean(), final_rewards.median(),
final_rewards.min(), final_rewards.max(),
dist_entropy.data[0], value_loss.data[0],
action_loss.data[0]))
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, args.env_name,
args.algo)
except IOError:
pass
def main():
print("######")
print("HELLO! Returns start with infinity values")
print("######")
os.environ['OMP_NUM_THREADS'] = '1'
if args.random_task:
env_params = {
'wt': np.round(np.random.uniform(0.5, 1.0), 2),
'x': np.round(np.random.uniform(-0.1, 0.1), 2),
'y': np.round(np.random.uniform(-0.1, 0.1), 2),
'z': np.round(np.random.uniform(0.15, 0.2), 2),
}
else:
env_params = {
'wt': args.euclidean_weight,
'x': args.goal_x,
'y': args.goal_y,
'z': args.goal_z,
}
envs = [make_env(args.env_name, args.seed, i, args.log_dir, **env_params)
for i in range(args.num_processes)]
if args.num_processes > 1:
envs = SubprocVecEnv(envs)
else:
envs = DummyVecEnv(envs)
envs = VecNormalize(envs, ob=False)
obs_shape = envs.observation_space.shape
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])
if len(envs.observation_space.shape) == 3:
actor_critic = CNNPolicy(obs_shape[0], envs.action_space, args.recurrent_policy)
else:
assert not args.recurrent_policy, \
"Recurrent policy is not implemented for the MLP controller"
actor_critic = MLPPolicy(obs_shape[0], envs.action_space)
if envs.action_space.__class__.__name__ == "Discrete":
action_shape = 1
else:
action_shape = envs.action_space.shape[0]
if args.cuda:
actor_critic.cuda()
if args.algo == 'a2c':
optimizer = optim.RMSprop(actor_critic.parameters(), args.lr, eps=args.eps, alpha=args.alpha)
elif args.algo == 'ppo':
optimizer = optim.Adam(actor_critic.parameters(), args.lr, eps=args.eps)
elif args.algo == 'acktr':
optimizer = KFACOptimizer(actor_critic)
rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape, envs.action_space, actor_critic.state_size)
current_obs = torch.zeros(args.num_processes, *obs_shape)
def update_current_obs(obs):
shape_dim0 = envs.observation_space.shape[0]
obs = torch.from_numpy(obs).float()
if args.num_stack > 1:
current_obs[:, :-shape_dim0] = current_obs[:, shape_dim0:]
current_obs[:, -shape_dim0:] = obs
obs = envs.reset()
update_current_obs(obs)
rollouts.observations[0].copy_(current_obs)
# These variables are used to compute average rewards for all processes.
episode_rewards = torch.zeros([args.num_processes, 1])
final_rewards = torch.zeros([args.num_processes, 1])
if args.cuda:
current_obs = current_obs.cuda()
rollouts.cuda()
actor_critic.input_norm.update(rollouts.observations[0])
last_return = -np.inf
best_return = -np.inf
best_models = None
start = time.time()
for j in range(num_updates):
for step in range(args.num_steps):
# Sample actions
value, action, action_log_prob, states = actor_critic.act(Variable(rollouts.observations[step], volatile=True),
Variable(rollouts.states[step], volatile=True),
Variable(rollouts.masks[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy()
# Obser reward and next obs
obs, reward, done, info = envs.step(cpu_actions)
reward = torch.from_numpy(np.expand_dims(np.stack(reward), 1)).float()
episode_rewards += reward
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
if args.cuda:
masks = masks.cuda()
if current_obs.dim() == 4:
current_obs *= masks.unsqueeze(2).unsqueeze(2)
else:
current_obs *= masks
update_current_obs(obs)
rollouts.insert(step, current_obs, states.data, action.data, action_log_prob.data, value.data, reward, masks)
actor_critic.input_norm.update(rollouts.observations[step + 1])
next_value = actor_critic(Variable(rollouts.observations[-1], volatile=True),
Variable(rollouts.states[-1], volatile=True),
Variable(rollouts.masks[-1], volatile=True))[0].data
rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau)
if args.algo in ['a2c', 'acktr']:
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(Variable(rollouts.observations[:-1].view(-1, *obs_shape)),
Variable(rollouts.states[0].view(-1, actor_critic.state_size)),
Variable(rollouts.masks[:-1].view(-1, 1)),
Variable(rollouts.actions.view(-1, action_shape)))
values = values.view(args.num_steps, args.num_processes, 1)
action_log_probs = action_log_probs.view(args.num_steps, args.num_processes, 1)
advantages = Variable(rollouts.returns[:-1]) - values
value_loss = advantages.pow(2).mean()
action_loss = -(Variable(advantages.data) * action_log_probs).mean()
if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(values - Variable(sample_values.data)).pow(2).mean()
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss - dist_entropy * args.entropy_coef).backward()
if args.algo == 'a2c':
nn.utils.clip_grad_norm(actor_critic.parameters(), args.max_grad_norm)
optimizer.step()
elif args.algo == 'ppo':
advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)
for e in range(args.ppo_epoch):
if args.recurrent_policy:
data_generator = rollouts.recurrent_generator(advantages,
args.num_mini_batch)
else:
data_generator = rollouts.feed_forward_generator(advantages,
args.num_mini_batch)
for sample in data_generator:
observations_batch, states_batch, actions_batch, \
return_batch, masks_batch, old_action_log_probs_batch, \
adv_targ = sample
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(Variable(observations_batch),
Variable(states_batch),
Variable(masks_batch),
Variable(actions_batch))
adv_targ = Variable(adv_targ)
ratio = torch.exp(action_log_probs - Variable(old_action_log_probs_batch))
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) * adv_targ
action_loss = -torch.min(surr1, surr2).mean() # PPO's pessimistic surrogate (L^CLIP)
value_loss = (Variable(return_batch) - values).pow(2).mean()
optimizer.zero_grad()
(value_loss + action_loss - dist_entropy * args.entropy_coef).backward()
nn.utils.clip_grad_norm(actor_critic.parameters(), args.max_grad_norm)
optimizer.step()
rollouts.after_update()
if args.vis and j % args.vis_interval == 0:
last_return = plot(logger, args.log_dir)
if last_return > best_return:
best_return = last_return
try:
os.makedirs(os.path.dirname(args.save_path))
except OSError:
pass
info = {
'return': best_return,
'reward_norm': np.sqrt(envs.ret_rms.var + envs.epsilon)
}
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
torch.save((save_model, env_params, info), args.save_path)
if j % args.log_interval == 0:
end = time.time()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
print("Updates {}, num timesteps {}, FPS {}, average return {:.5f}, best_return {:.5f}, value loss {:.5f}, policy loss {:.5f}".
format(j, total_num_steps,
int(total_num_steps / (end - start)),
last_return, best_return,
value_loss.data[0], action_loss.data[0]))
def main():
print("###############################################################")
print("#################### VISDOOM LEARNER START ####################")
print("###############################################################")
os.environ['OMP_NUM_THREADS'] = '1'
if args.vis:
from visdom import Visdom
viz = Visdom()
win = None
global envs
envs = VecEnv(
[make_env(i, args.config_path) for i in range(args.num_processes)],
logging=True,
log_dir=args.log_dir)
obs_shape = envs.observation_space_shape
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])
if args.algo == 'a2c' or args.algo == 'acktr':
actor_critic = CNNPolicy(obs_shape[0], envs.action_space_shape)
elif args.algo == 'a2t':
source_models = []
files = glob.glob(os.path.join(args.source_models_path, '*.pt'))
for file in files:
print(file, 'loading model...')
source_models.append(torch.load(file))
actor_critic = A2TPolicy(obs_shape[0], envs.action_space_shape,
source_models)
elif args.algo == 'resnet':
# args.num_stack = 3
actor_critic = ResnetPolicy(obs_shape[0], envs.action_space_shape)
action_shape = 1
if args.cuda:
actor_critic.cuda()
if args.algo == 'a2c' or args.algo == 'resnet':
optimizer = optim.RMSprop(actor_critic.parameters(),
args.lr,
eps=args.eps,
alpha=args.alpha)
elif args.algo == 'a2t':
a2t_params = [p for p in actor_critic.parameters() if p.requires_grad]
optimizer = optim.RMSprop(a2t_params,
args.lr,
eps=args.eps,
alpha=args.alpha)
elif args.algo == 'acktr':
optimizer = KFACOptimizer(actor_critic)
rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape,
envs.action_space_shape)
current_obs = torch.zeros(args.num_processes, *obs_shape)
def update_current_obs(obs):
shape_dim0 = envs.observation_space_shape[0]
obs = torch.from_numpy(obs).float()
if args.num_stack > 1:
current_obs[:, :-shape_dim0] = current_obs[:, shape_dim0:]
current_obs[:, -shape_dim0:] = obs
obs = envs.reset()
update_current_obs(obs)
rollouts.observations[0].copy_(current_obs)
# These variables are used to compute average rewards for all processes.
episode_rewards = torch.zeros([args.num_processes, 1])
final_rewards = torch.zeros([args.num_processes, 1])
if args.cuda:
current_obs = current_obs.cuda()
rollouts.cuda()
start = time.time()
for j in range(num_updates):
for step in range(args.num_steps):
# Sample actions
value, action = actor_critic.act(
Variable(rollouts.observations[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy()
# Obser reward and next obs
obs, reward, done, info = envs.step(cpu_actions)
# print ('Actions:', cpu_actions, 'Rewards:', reward)
reward = torch.from_numpy(np.expand_dims(np.stack(reward),
1)).float()
episode_rewards += reward
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
if args.cuda:
masks = masks.cuda()
if current_obs.dim() == 4:
current_obs *= masks.unsqueeze(2).unsqueeze(2)
else:
current_obs *= masks
update_current_obs(obs)
rollouts.insert(step, current_obs, action.data, value.data, reward,
masks)
next_value = actor_critic(
Variable(rollouts.observations[-1], volatile=True))[0].data
if hasattr(actor_critic, 'obs_filter'):
actor_critic.obs_filter.update(rollouts.observations[:-1].view(
-1, *obs_shape))
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
if args.algo in ['a2c', 'acktr', 'a2t', 'resnet']:
values, action_log_probs, dist_entropy = actor_critic.evaluate_actions(
Variable(rollouts.observations[:-1].view(-1, *obs_shape)),
Variable(rollouts.actions.view(-1, action_shape)))
values = values.view(args.num_steps, args.num_processes, 1)
action_log_probs = action_log_probs.view(args.num_steps,
args.num_processes, 1)
advantages = Variable(rollouts.returns[:-1]) - values
value_loss = advantages.pow(2).mean()
action_loss = -(Variable(advantages.data) *
action_log_probs).mean()
if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(values -
Variable(sample_values.data)).pow(2).mean()
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss -
dist_entropy * args.entropy_coef).backward()
if args.algo == 'a2c' or args.algo == 'resnet':
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
elif args.algo == 'a2t':
nn.utils.clip_grad_norm(a2t_params, args.max_grad_norm)
optimizer.step()
rollouts.observations[0].copy_(rollouts.observations[-1])
if j % args.save_interval == 0 and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
torch.save(save_model,
os.path.join(save_path, args.env_name + ".pt"))
if j % args.log_interval == 0:
envs.log()
end = time.time()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
print(
"Updates {}, num timesteps {}, FPS {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
final_rewards.mean(), final_rewards.median(),
final_rewards.min(), final_rewards.max(),
-dist_entropy.data[0], value_loss.data[0],
action_loss.data[0]))
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, 'VizDoom', args.algo)
except IOError:
pass
envs.close()
time.sleep(5)
def main():
print("#######")
print(
"WARNING: All rewards are clipped so you need to use a monitor (see envs.py) or visdom plot to get true rewards"
)
print("#######")
os.environ['OMP_NUM_THREADS'] = '1'
if args.vis:
from visdom import Visdom
viz = Visdom()
win = None
envs = SubprocVecEnv([
make_env(args.env_name, args.seed, i, args.log_dir)
for i in range(args.num_processes)
])
obs_shape = envs.observation_space.shape
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])
if len(envs.observation_space.shape) == 3:
actor_critic = CNNPolicy(obs_shape[0], envs.action_space)
else:
actor_critic = MLPPolicy(obs_shape[0], envs.action_space)
if envs.action_space.__class__.__name__ == "Discrete":
action_shape = 1
else:
action_shape = envs.action_space.shape[0]
if args.cuda:
actor_critic.cuda()
if args.algo == 'a2c':
optimizer = optim.RMSprop(actor_critic.parameters(),
args.lr,
eps=args.eps,
alpha=args.alpha)
elif args.algo == 'ppo':
optimizer = optim.Adam(actor_critic.parameters(),
args.lr,
eps=args.eps)
elif args.algo == 'acktr':
optimizer = KFACOptimizer(actor_critic)
rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape,
envs.action_space)
current_state = torch.zeros(args.num_processes, *obs_shape)
def update_current_state(state):
shape_dim0 = envs.observation_space.shape[0]
state = torch.from_numpy(state).float()
if args.num_stack > 1:
current_state[:, :-shape_dim0] = current_state[:, shape_dim0:]
current_state[:, -shape_dim0:] = state
state = envs.reset()
update_current_state(state)
rollouts.states[0].copy_(current_state)
# These variables are used to compute average rewards for all processes.
episode_rewards = torch.zeros([args.num_processes, 1])
final_rewards = torch.zeros([args.num_processes, 1])
if args.cuda:
current_state = current_state.cuda()
rollouts.cuda()
if args.algo == 'ppo':
old_model = copy.deepcopy(actor_critic)
for j in range(num_updates):
for step in range(args.num_steps):
# Sample actions
value, action = actor_critic.act(
Variable(rollouts.states[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy()
# Obser reward and next state
state, reward, done, info = envs.step(cpu_actions)
reward = torch.from_numpy(np.expand_dims(np.stack(reward),
1)).float()
episode_rewards += reward
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
if args.cuda:
masks = masks.cuda()
if current_state.dim() == 4:
current_state *= masks.unsqueeze(2).unsqueeze(2)
else:
current_state *= masks
update_current_state(state)
rollouts.insert(step, current_state, action.data, value.data,
reward, masks)
next_value = actor_critic(Variable(rollouts.states[-1],
volatile=True))[0].data
if hasattr(actor_critic, 'obs_filter'):
actor_critic.obs_filter.update(rollouts.states[:-1].view(
-1, *obs_shape))
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
if args.algo in ['a2c', 'acktr']:
values, action_log_probs, dist_entropy = actor_critic.evaluate_actions(
Variable(rollouts.states[:-1].view(-1, *obs_shape)),
Variable(rollouts.actions.view(-1, action_shape)))
values = values.view(args.num_steps, args.num_processes, 1)
action_log_probs = action_log_probs.view(args.num_steps,
args.num_processes, 1)
advantages = Variable(rollouts.returns[:-1]) - values
value_loss = advantages.pow(2).mean()
action_loss = -(Variable(advantages.data) *
action_log_probs).mean()
if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(values -
Variable(sample_values.data)).pow(2).mean()
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss -
dist_entropy * args.entropy_coef).backward()
if args.algo == 'a2c':
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
elif args.algo == 'ppo':
advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-5)
old_model.load_state_dict(actor_critic.state_dict())
if hasattr(actor_critic, 'obs_filter'):
old_model.obs_filter = actor_critic.obs_filter
for _ in range(args.ppo_epoch):
sampler = BatchSampler(SubsetRandomSampler(
range(args.num_processes * args.num_steps)),
args.batch_size * args.num_processes,
drop_last=False)
for indices in sampler:
indices = torch.LongTensor(indices)
if args.cuda:
indices = indices.cuda()
states_batch = rollouts.states[:-1].view(
-1, *obs_shape)[indices]
actions_batch = rollouts.actions.view(
-1, action_shape)[indices]
return_batch = rollouts.returns[:-1].view(-1, 1)[indices]
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy = actor_critic.evaluate_actions(
Variable(states_batch), Variable(actions_batch))
_, old_action_log_probs, _ = old_model.evaluate_actions(
Variable(states_batch, volatile=True),
Variable(actions_batch, volatile=True))
ratio = torch.exp(action_log_probs -
Variable(old_action_log_probs.data))
adv_targ = Variable(advantages.view(-1, 1)[indices])
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - args.clip_param,
1.0 + args.clip_param) * adv_targ
action_loss = -torch.min(
surr1,
surr2).mean() # PPO's pessimistic surrogate (L^CLIP)
value_loss = (Variable(return_batch) -
values).pow(2).mean()
optimizer.zero_grad()
(value_loss + action_loss -
dist_entropy * args.entropy_coef).backward()
optimizer.step()
rollouts.states[0].copy_(rollouts.states[-1])
if j % args.save_interval == 0 and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
torch.save(save_model,
os.path.join(save_path, args.env_name + ".pt"))
if j % args.log_interval == 0:
print(
"Updates {}, num frames {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}"
.format(j, (j + 1) * args.num_processes * args.num_steps,
final_rewards.mean(), final_rewards.median(),
final_rewards.min(), final_rewards.max(),
-dist_entropy.data[0], value_loss.data[0],
action_loss.data[0]))
if j % args.vis_interval == 0:
win = visdom_plot(viz, win, args.log_dir, args.env_name, args.algo)
def main():
print("#######")
print(
"WARNING: All rewards are clipped or normalized so you need to use a monitor (see envs.py) or visdom plot to get true rewards"
)
print("#######")
os.environ['OMP_NUM_THREADS'] = '1'
if args.vis:
from visdom import Visdom
viz = Visdom()
win = None
envs = [
make_env(args.env_name, args.seed, i, args.log_dir)
for i in range(args.num_processes)
]
if args.num_processes > 1:
envs = SubprocVecEnv(envs)
else:
envs = DummyVecEnv(envs)
if len(envs.observation_space.shape) == 1:
envs = VecNormalize(envs)
obs_shape = envs.observation_space.shape
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:]
) # I guess the obs_shape[0] is channel number
if len(envs.observation_space.shape) == 3:
actor_critic = CNNPolicy(obs_shape[0], envs.action_space,
args.recurrent_policy)
else:
assert not args.recurrent_policy, \
"Recurrent policy is not implemented for the MLP controller"
actor_critic = MLPPolicy(obs_shape[0], envs.action_space)
if envs.action_space.__class__.__name__ == "Discrete":
action_shape = 1
else:
action_shape = envs.action_space.shape[0]
if args.cuda:
actor_critic.cuda()
if args.algo == 'a2c':
optimizer = optim.RMSprop(actor_critic.parameters(),
args.lr,
eps=args.eps,
alpha=args.alpha)
elif args.algo == 'ppo':
optimizer = optim.Adam(actor_critic.parameters(),
args.lr,
eps=args.eps)
elif args.algo == 'acktr':
optimizer = KFACOptimizer(actor_critic)
rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape,
envs.action_space, actor_critic.state_size)
current_obs = torch.zeros(args.num_processes, *obs_shape)
def update_current_obs(obs):
shape_dim0 = envs.observation_space.shape[0]
obs = torch.from_numpy(obs).float()
if args.num_stack > 1:
current_obs[:, :-shape_dim0] = current_obs[:, shape_dim0:]
current_obs[:, -shape_dim0:] = obs
obs = envs.reset()
update_current_obs(obs)
rollouts.observations[0].copy_(current_obs)
# These variables are used to compute average rewards for all processes.
episode_rewards = torch.zeros([args.num_processes, 1])
final_rewards = torch.zeros([args.num_processes, 1])
if args.cuda:
current_obs = current_obs.cuda()
rollouts.cuda()
start = time.time()
for j in range(num_updates):
for step in range(args.num_steps):
# args.num_steps should be the length of interactions before each updating/training
# Sample actions
value, action, action_log_prob, states = actor_critic.act(
Variable(rollouts.observations[step], volatile=True),
Variable(rollouts.states[step], volatile=True),
Variable(rollouts.masks[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy(
) # returns are state value, sampled action, act_log_prob, hidden states
# Obser reward and next obs
obs, reward, done, info = envs.step(cpu_actions)
reward = torch.from_numpy(np.expand_dims(np.stack(reward),
1)).float()
episode_rewards += reward
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
if args.cuda:
masks = masks.cuda()
if current_obs.dim() == 4:
current_obs *= masks.unsqueeze(2).unsqueeze(2)
else:
current_obs *= masks
update_current_obs(obs)
rollouts.insert(
step, current_obs, states.data, action.data,
action_log_prob.data, value.data, reward, masks
) # so the rollout stores one batch of interaction sequences, each sequence has length of args.num_steps
next_value = actor_critic(
Variable(rollouts.observations[-1], volatile=True),
Variable(rollouts.states[-1], volatile=True),
Variable(rollouts.masks[-1], volatile=True))[0].data
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
if args.algo in ['a2c', 'acktr']:
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(rollouts.observations[:-1].view(-1, *obs_shape)),
Variable(rollouts.states[0].view(-1, actor_critic.state_size)),
Variable(rollouts.masks[:-1].view(-1, 1)),
Variable(rollouts.actions.view(-1, action_shape)))
# values should be values of observations, states are the hidden states used in rnn module, by pwang8
values = values.view(
args.num_steps, args.num_processes,
1) # values are estimated current state values
action_log_probs = action_log_probs.view(args.num_steps,
args.num_processes, 1)
# rollouts.returns are current "Action" value calculted following Bellmans' eqaution gamma * State_value(t+1) + reward(t)
advantages = Variable(
rollouts.returns[:-1]
) - values # This is also the definition of advantage value (action_value - state_value).
value_loss = advantages.pow(
2).mean() # values are estimated current state_value(t)
action_loss = -(Variable(advantages.data) *
action_log_probs).mean()
# If ACKTR is utilized, it is not only a different optimizer is used, they also added some new loss source
if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(
values - Variable(sample_values.data)
).pow(2).mean(
) # don't know what is the difference between this and just randomly sample some noise
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss -
dist_entropy * args.entropy_coef).backward()
if args.algo == 'a2c':
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
elif args.algo == 'ppo':
advantages = rollouts.returns[:-1] - rollouts.value_preds[:
-1] # calculating the advantage value of an action
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-5)
# The difference from this ppo optimization to the optimization above is that: it updates params for
# multiple epochs in ppo optimization. Because of this, it samples from the rollouts storage a minibatch
# every time to calculate gradient. Sampling is conducted for optimization purpose.
for e in range(args.ppo_epoch):
if args.recurrent_policy:
data_generator = rollouts.recurrent_generator(
advantages, args.num_mini_batch)
else:
data_generator = rollouts.feed_forward_generator(
advantages, args.num_mini_batch)
for sample in data_generator:
observations_batch, states_batch, actions_batch, \
return_batch, masks_batch, old_action_log_probs_batch, \
adv_targ = sample
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(observations_batch), Variable(states_batch),
Variable(masks_batch), Variable(actions_batch))
# For the 1st epoch of updating, I guess the action_log_probls is the same as old_action_log_probs_batch
# because params of the NN have not been updated at that time. But later, in other updating epochs,
# this ratio will generate some error. The old_action_log_probs_batch will not be updated during
# these param updating epochs.
# action_log_probs is the log prob of that action taken by the agent. So it's one value here, not
# log_prob for all actions with certain input observation/state. By pwang8, Dec 31, 2017
adv_targ = Variable(adv_targ)
ratio = torch.exp(action_log_probs -
Variable(old_action_log_probs_batch))
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - args.clip_param,
1.0 + args.clip_param) * adv_targ
action_loss = -torch.min(
surr1,
surr2).mean() # PPO's pessimistic surrogate (L^CLIP)
# compared to a2c, the major difference for ppo is that action_loss is calculated in controlled way
value_loss = (Variable(return_batch) -
values).pow(2).mean()
optimizer.zero_grad()
(value_loss + action_loss -
dist_entropy * args.entropy_coef).backward()
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
rollouts.after_update()
if j % args.save_interval == 0 and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
hasattr(envs, 'ob_rms') and envs.ob_rms or None
]
torch.save(save_model,
os.path.join(save_path, args.env_name + ".pt"))
if j % args.log_interval == 0:
end = time.time()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
print(
"Updates {}, num timesteps {}, FPS {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
final_rewards.mean(), final_rewards.median(),
final_rewards.min(), final_rewards.max(),
dist_entropy.data[0], value_loss.data[0],
action_loss.data[0]))
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, args.env_name,
args.algo)
except IOError:
pass
def main():
print("#######")
print(
"WARNING: All rewards are clipped so you need to use a monitor (see envs.py) or visdom plot to get true rewards"
)
print("#######")
os.environ['OMP_NUM_THREADS'] = '1'
# T choose whetehr to visualize
if args.vis:
from visdom import Visdom
viz = Visdom()
win = None
envs = SubprocVecEnv([
make_env(args.env_name, args.seed, i, args.log_dir)
for i in range(args.num_processes)
])
# T get shape of observation array of the environment
obs_shape = envs.observation_space.shape
# T adjusting the shape; not sure what the * is
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])
#T initialize the actor critic; MLP and CNN classes imported from model.py
if len(envs.observation_space.shape) == 3:
actor_critic = CNNPolicy(obs_shape[0], envs.action_space)
else:
actor_critic = MLPPolicy(obs_shape[0], envs.action_space)
#T - some kind of setup with the actor_critic
if args.finetune:
checkpoint_path = save_path = os.path.join(args.save_dir, args.algo,
args.checkpoint)
state_dict = torch.load(checkpoint_path)
print("Finetuning from checkpoint: %s, at step: %d" %
(checkpoint_path, state_dict['update']))
actor_critic.load_state_dict(state_dict['model_state_dict'])
keep_layers = [
'v_fc3.weight', 'v_fc3.bias', 'a_fc2.weight', 'a_fc2.bias',
'dist.fc_mean.weight', 'dist.fc_mean.bias', 'dist.logstd._bias'
]
for name, param in actor_critic.named_parameters():
if name not in keep_layers:
param.requires_grad = False
for name, param in actor_critic.named_parameters():
print('Param name: %s, requires_grad: %d' %
(name, param.requires_grad))
# T set up dimensions of the action space
if envs.action_space.__class__.__name__ == "Discrete":
action_shape = 1
else:
action_shape = envs.action_space.shape[0]
# T all arguments imported from arguments.py
# T enable cuda pythorch tensor support
if args.cuda:
actor_critic.cuda()
# T - pull arguments and choose algorithm and optimizer
if args.algo == 'a2c':
optimizer = optim.RMSprop(filter(lambda p: p.requires_grad,
actor_critic.parameters()),
args.lr,
eps=args.eps,
alpha=args.alpha)
elif args.algo == 'ppo':
optimizer = optim.Adam(actor_critic.parameters(),
args.lr,
eps=args.eps)
elif args.algo == 'acktr':
optimizer = KFACOptimizer(actor_critic)
#TO-DO figure out how to restore optimizer parameters when freezing some weights
rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape,
envs.action_space)
# return all zeros, so nothing observed
current_obs = torch.zeros(args.num_processes, *obs_shape)
# T-not sure what this function is doing??
def update_current_obs(obs):
shape_dim0 = envs.observation_space.shape[0]
obs = torch.from_numpy(obs).float()
if args.num_stack > 1:
current_obs[:, :-shape_dim0] = current_obs[:, shape_dim0:]
current_obs[:, -shape_dim0:] = obs
# T - reset the environment; call function to update observation
obs = envs.reset()
update_current_obs(obs)
rollouts.observations[0].copy_(current_obs)
# These variables are used to compute average rewards for all processes.
# T - initialize rewards to be zero
episode_rewards = torch.zeros([args.num_processes, 1])
final_rewards = torch.zeros([args.num_processes, 1])
if args.cuda:
current_obs = current_obs.cuda()
rollouts.cuda()
if args.algo == 'ppo':
old_model = copy.deepcopy(actor_critic)
start = time.time()
# T - begin iterative loop
for j in range(num_updates):
# T - take steps through single instance
# T - this is the loop where action/critic happens
for step in range(args.num_steps):
# Sample actions
# T - buried by the action method ultimately comes from torch.nn.Module
value, action = actor_critic.act(
Variable(rollouts.observations[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy()
# Obser reward and next obs
obs, reward, done, info = envs.step(cpu_actions)
reward = torch.from_numpy(np.expand_dims(np.stack(reward),
1)).float()
episode_rewards += reward
# T done bool returned by steps; indicates if failure occurred (done)
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
if args.cuda:
masks = masks.cuda()
if current_obs.dim() == 4:
current_obs *= masks.unsqueeze(2).unsqueeze(2)
else:
current_obs *= masks
#T - now update the observation matrix
update_current_obs(obs)
#T - store what happened in this step
rollouts.insert(step, current_obs, action.data, value.data, reward,
masks)
next_value = actor_critic(
Variable(rollouts.observations[-1], volatile=True))[0].data
if hasattr(actor_critic, 'obs_filter'):
actor_critic.obs_filter.update(rollouts.observations[:-1].view(
-1, *obs_shape))
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
if args.algo in ['a2c', 'acktr']:
values, action_log_probs, dist_entropy = actor_critic.evaluate_actions(
Variable(rollouts.observations[:-1].view(-1, *obs_shape)),
Variable(rollouts.actions.view(-1, action_shape)))
values = values.view(args.num_steps, args.num_processes, 1)
action_log_probs = action_log_probs.view(args.num_steps,
args.num_processes, 1)
advantages = Variable(rollouts.returns[:-1]) - values
value_loss = advantages.pow(2).mean()
action_loss = -(Variable(advantages.data) *
action_log_probs).mean()
if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(values -
Variable(sample_values.data)).pow(2).mean()
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss -
dist_entropy * args.entropy_coef).backward()
if args.algo == 'a2c':
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
elif args.algo == 'ppo':
advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-5)
old_model.load_state_dict(actor_critic.state_dict())
for _ in range(args.ppo_epoch):
sampler = BatchSampler(SubsetRandomSampler(
range(args.num_processes * args.num_steps)),
args.batch_size * args.num_processes,
drop_last=False)
for indices in sampler:
indices = torch.LongTensor(indices)
if args.cuda:
indices = indices.cuda()
observations_batch = rollouts.observations[:-1].view(
-1, *obs_shape)[indices]
actions_batch = rollouts.actions.view(
-1, action_shape)[indices]
return_batch = rollouts.returns[:-1].view(-1, 1)[indices]
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy = actor_critic.evaluate_actions(
Variable(observations_batch), Variable(actions_batch))
_, old_action_log_probs, _ = old_model.evaluate_actions(
Variable(observations_batch, volatile=True),
Variable(actions_batch, volatile=True))
ratio = torch.exp(action_log_probs -
Variable(old_action_log_probs.data))
adv_targ = Variable(advantages.view(-1, 1)[indices])
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - args.clip_param,
1.0 + args.clip_param) * adv_targ
action_loss = -torch.min(
surr1,
surr2).mean() # PPO's pessimistic surrogate (L^CLIP)
value_loss = (Variable(return_batch) -
values).pow(2).mean()
optimizer.zero_grad()
(value_loss + action_loss -
dist_entropy * args.entropy_coef).backward()
optimizer.step()
rollouts.observations[0].copy_(rollouts.observations[-1])
if j % args.save_interval == 0 and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
file_name = FILE_PREFIX + '.pt'
#torch.save(save_model, os.path.join(save_path, file_name))
data = {
'update': j,
'model_state_dict': save_model.state_dict(),
'optim_state_dict': optimizer.state_dict()
}
torch.save(data, os.path.join(save_path, file_name))
# T - write out some log information (not important for us)
if j % args.log_interval == 0:
end = time.time()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
print(
"Updates {}, num timesteps {}, FPS {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
final_rewards.mean(), final_rewards.median(),
final_rewards.min(), final_rewards.max(),
-dist_entropy.data[0], value_loss.data[0],
action_loss.data[0]))
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, args.env_name,
args.algo)
except IOError:
pass
def main():
num_updates = int(
config.max_num_frames) // args.num_steps // config.a2c.num_processes
n_times_is_converging = 0
print("num_updates: " + str(num_updates))
print("stop_learning: " + str(config.a2c.stop_learning))
# Initializing evaluation
evaluator = Evaluator(evaluation_id)
os.environ['OMP_NUM_THREADS'] = '1'
envs = [
make_env(config.env_name, args.seed, i, evaluation_id)
for i in range(config.a2c.num_processes)
]
if config.a2c.num_processes > 1:
envs = SubprocVecEnv(envs)
else:
envs = DummyVecEnv(envs)
obs_shape = envs.observation_space.shape
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])
obs_numel = reduce(operator.mul, obs_shape, 1)
actor_critic = Policy(obs_numel, envs.action_space)
# Maxime: log some info about the model and its size
modelSize = 0
for p in actor_critic.parameters():
pSize = reduce(operator.mul, p.size(), 1)
modelSize += pSize
print(str(actor_critic))
print('Total model size: %d' % modelSize)
if envs.action_space.__class__.__name__ == "Discrete":
action_shape = 1
else:
action_shape = envs.action_space.shape[0]
if args.cuda:
actor_critic.cuda()
if config.a2c.algorithm == 'a2c':
optimizer = optim.RMSprop(actor_critic.parameters(),
args.lr,
eps=args.eps,
alpha=args.alpha)
elif config.a2c.algorithm == 'ppo':
optimizer = optim.Adam(actor_critic.parameters(),
args.lr,
eps=args.eps)
elif config.a2c.algorithm == 'acktr':
optimizer = KFACOptimizer(actor_critic)
rollouts = RolloutStorage(args.num_steps, config.a2c.num_processes,
obs_shape, envs.action_space,
actor_critic.state_size)
current_obs = torch.zeros(config.a2c.num_processes, *obs_shape)
def update_current_obs(obs):
shape_dim0 = envs.observation_space.shape[0]
obs = torch.from_numpy(obs).float()
if args.num_stack > 1:
current_obs[:, :-shape_dim0] = current_obs[:, shape_dim0:]
current_obs[:, -shape_dim0:] = obs
obs = envs.reset()
update_current_obs(obs)
rollouts.observations[0].copy_(current_obs)
# These variables are used to compute average rewards for all processes.
episode_rewards = torch.zeros([config.a2c.num_processes, 1])
final_rewards = torch.zeros([config.a2c.num_processes, 1])
if args.cuda:
current_obs = current_obs.cuda()
rollouts.cuda()
start = time.time()
send_env_name = False
for j in range(num_updates):
if n_times_is_converging > 1:
print("Converged...")
break
for step in range(args.num_steps):
# Sample actions
value, action, action_log_prob, states = actor_critic.act(
Variable(rollouts.observations[step], volatile=True),
Variable(rollouts.states[step], volatile=True),
Variable(rollouts.masks[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy()
# Obser reward and next obs
obs, reward, done, info = envs.step(cpu_actions)
evaluator.update(done, info)
reward = torch.from_numpy(np.expand_dims(np.stack(reward),
1)).float()
episode_rewards += reward
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
if args.cuda:
masks = masks.cuda()
if current_obs.dim() == 4:
current_obs *= masks.unsqueeze(2).unsqueeze(2)
elif current_obs.dim() == 3:
current_obs *= masks.unsqueeze(2)
else:
current_obs *= masks
update_current_obs(obs)
rollouts.insert(step, current_obs, states.data, action.data,
action_log_prob.data, value.data, reward, masks)
next_value = actor_critic(
Variable(rollouts.observations[-1], volatile=True),
Variable(rollouts.states[-1], volatile=True),
Variable(rollouts.masks[-1], volatile=True))[0].data
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
if config.a2c.algorithm in ['a2c', 'acktr']:
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(rollouts.observations[:-1].view(-1, *obs_shape)),
Variable(rollouts.states[:-1].view(-1,
actor_critic.state_size)),
Variable(rollouts.masks[:-1].view(-1, 1)),
Variable(rollouts.actions.view(-1, action_shape)))
values = values.view(args.num_steps, config.a2c.num_processes, 1)
action_log_probs = action_log_probs.view(args.num_steps,
config.a2c.num_processes,
1)
advantages = Variable(rollouts.returns[:-1]) - values
value_loss = advantages.pow(2).mean()
action_loss = -(Variable(advantages.data) *
action_log_probs).mean()
if config.a2c.algorithm == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(values -
Variable(sample_values.data)).pow(2).mean()
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss -
dist_entropy * args.entropy_coef).backward()
if config.a2c.algorithm == 'a2c':
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
elif config.a2c.algorithm == 'ppo':
advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-5)
for e in range(args.ppo_epoch):
if args.recurrent_policy:
data_generator = rollouts.recurrent_generator(
advantages, args.num_mini_batch)
else:
data_generator = rollouts.feed_forward_generator(
advantages, args.num_mini_batch)
for sample in data_generator:
observations_batch, states_batch, actions_batch, \
return_batch, masks_batch, old_action_log_probs_batch, \
adv_targ = sample
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(observations_batch), Variable(states_batch),
Variable(masks_batch), Variable(actions_batch))
adv_targ = Variable(adv_targ)
ratio = torch.exp(action_log_probs -
Variable(old_action_log_probs_batch))
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - args.clip_param,
1.0 + args.clip_param) * adv_targ
action_loss = -torch.min(
surr1,
surr2).mean() # PPO's pessimistic surrogate (L^CLIP)
value_loss = (Variable(return_batch) -
values).pow(2).mean()
optimizer.zero_grad()
(value_loss + action_loss -
dist_entropy * args.entropy_coef).backward()
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
rollouts.after_update()
save_dir = "../a2c_trained_model/"
if j % config.a2c.save_model_interval == 0:
save_path = save_dir
try:
os.makedirs(save_path)
except OSError:
pass
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
hasattr(envs, 'ob_rms') and envs.ob_rms or None
]
if j % config.a2c.save_evaluation_interval == 0:
end = time.time()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
# if the environment name and the envelope state was not send
if not send_env_name:
evaluator.save(j, total_num_steps, final_rewards, dist_entropy,
value_loss, action_loss, config.env_name,
config.envelope)
send_env_name = True
else:
evaluator.save(j, total_num_steps, final_rewards, dist_entropy,
value_loss, action_loss)
if evaluator.is_converging:
n_times_is_converging += 1
else:
n_times_is_converging = 0
print(
"Updates {}, num timesteps {}, FPS {}, mean/median reward {:.2f}/{:.2f}, min/max reward {:.2f}/{:.2f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
final_rewards.mean(), final_rewards.median(),
final_rewards.min(), final_rewards.max(),
dist_entropy.data[0], value_loss.data[0],
action_loss.data[0]))
if config.visdom and j % config.visdom_interval == 0:
win = visdom_plot(total_num_steps, final_rewards.mean())
def main():
os.environ['OMP_NUM_THREADS'] = '1'
if args.vis:
from visdom import Visdom
viz = Visdom()
win = None
envs = [
make_env(args.env_name, args.seed, i, args.log_dir,
args.start_container) for i in range(args.num_processes)
]
test_envs = [
make_env(args.env_name, args.seed, i, args.log_dir,
args.start_container) for i in range(args.num_processes)
]
if args.num_processes > 1:
envs = SubprocVecEnv(envs)
test_envs = SubprocVecEnv(test_envs)
else:
envs = DummyVecEnv(envs)
test_envs = DummyVecEnv(test_envs)
obs_shape = envs.observation_space.shape
obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])
if args.saved_encoder_model:
obs_shape = (args.num_stack, args.latent_space_size)
obs_numel = reduce(operator.mul, obs_shape, 1)
if len(obs_shape) == 3 and obs_numel > 1024:
actor_critic = CNNPolicy(obs_shape[0], envs.action_space,
args.recurrent_policy)
else:
assert not args.recurrent_policy, \
"Recurrent policy is not implemented for the MLP controller"
actor_critic = MLPPolicy(obs_numel, envs.action_space)
modelSize = 0
for p in actor_critic.parameters():
pSize = reduce(operator.mul, p.size(), 1)
modelSize += pSize
print(str(actor_critic))
print('Total model size: %d' % modelSize)
if envs.action_space.__class__.__name__ == "Discrete":
action_shape = 1
else:
action_shape = envs.action_space.shape[0]
if args.resume_experiment:
print("\n############## Loading saved model ##############\n")
actor_critic, ob_rms = torch.load(
os.path.join(save_path, args.env_name + args.save_tag + ".pt"))
tr.load(os.path.join(log_path, args.env_name + args.save_tag + ".p"))
if args.cuda:
actor_critic.cuda()
if args.algo == 'a2c':
optimizer = optim.RMSprop(actor_critic.parameters(),
args.lr,
eps=args.eps,
alpha=args.alpha)
elif args.algo == 'ppo':
optimizer = optim.Adam(actor_critic.parameters(),
args.lr,
eps=args.eps)
elif args.algo == 'acktr':
optimizer = KFACOptimizer(actor_critic)
print(obs_shape)
rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape,
envs.action_space, actor_critic.state_size)
rollouts_test = RolloutStorage(args.num_steps_test, args.num_processes,
obs_shape, envs.action_space,
actor_critic.state_size)
current_obs = torch.zeros(args.num_processes, *obs_shape)
current_obs_test = torch.zeros(args.num_processes, *obs_shape)
def update_current_obs(obs, test=False):
shape_dim0 = envs.observation_space.shape[0]
if args.saved_encoder_model:
shape_dim0 = 1
obs, _ = vae.encode(Variable(torch.cuda.FloatTensor(obs)))
obs = obs.data.cpu().numpy()
obs = torch.from_numpy(obs).float()
if not test:
if args.num_stack > 1:
current_obs[:, :-shape_dim0] = current_obs[:, shape_dim0:]
current_obs[:, -shape_dim0:] = obs
else:
if args.num_stack > 1:
current_obs_test[:, :
-shape_dim0] = current_obs_test[:,
shape_dim0:]
current_obs_test[:, -shape_dim0:] = obs
obs = envs.reset()
update_current_obs(obs)
rollouts.observations[0].copy_(current_obs)
# These variables are used to compute average rewards for all processes.
episode_rewards = torch.zeros([args.num_processes, 1])
final_rewards = torch.zeros([args.num_processes, 1])
reward_avg = 0
if args.cuda:
current_obs = current_obs.cuda()
current_obs_test = current_obs_test.cuda()
rollouts.cuda()
rollouts_test.cuda()
start = time.time()
for j in range(num_updates):
for step in range(args.num_steps):
# Sample actions
value, action, action_log_prob, states = actor_critic.act(
Variable(rollouts.observations[step], volatile=True),
Variable(rollouts.states[step], volatile=True),
Variable(rollouts.masks[step], volatile=True))
cpu_actions = action.data.squeeze(1).cpu().numpy()
# Observation, reward and next obs
obs, reward, done, info = envs.step(cpu_actions)
# Maxime: clip the reward within [0,1] for more reliable training
# This code deals poorly with large reward values
reward = np.clip(reward, a_min=0, a_max=None) / 400
reward = torch.from_numpy(np.expand_dims(np.stack(reward),
1)).float()
episode_rewards += reward
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
tr.episodes_done += args.num_processes - masks.sum()
if args.cuda:
masks = masks.cuda()
if current_obs.dim() == 4:
current_obs *= masks.unsqueeze(2).unsqueeze(2)
else:
current_obs *= masks
update_current_obs(obs)
rollouts.insert(step, current_obs, states.data, action.data,
action_log_prob.data, value.data, reward, masks)
next_value = actor_critic(
Variable(rollouts.observations[-1], volatile=True),
Variable(rollouts.states[-1], volatile=True),
Variable(rollouts.masks[-1], volatile=True))[0].data
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
tr.iterations_done += 1
if args.algo in ['a2c', 'acktr']:
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(rollouts.observations[:-1].view(-1, *obs_shape)),
Variable(rollouts.states[0].view(-1, actor_critic.state_size)),
Variable(rollouts.masks[:-1].view(-1, 1)),
Variable(rollouts.actions.view(-1, action_shape)))
values = values.view(args.num_steps, args.num_processes, 1)
action_log_probs = action_log_probs.view(args.num_steps,
args.num_processes, 1)
advantages = Variable(rollouts.returns[:-1]) - values
value_loss = advantages.pow(2).mean()
action_loss = -(Variable(advantages.data) *
action_log_probs).mean()
if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
# Sampled fisher, see Martens 2014
actor_critic.zero_grad()
pg_fisher_loss = -action_log_probs.mean()
value_noise = Variable(torch.randn(values.size()))
if args.cuda:
value_noise = value_noise.cuda()
sample_values = values + value_noise
vf_fisher_loss = -(values -
Variable(sample_values.data)).pow(2).mean()
fisher_loss = pg_fisher_loss + vf_fisher_loss
optimizer.acc_stats = True
fisher_loss.backward(retain_graph=True)
optimizer.acc_stats = False
optimizer.zero_grad()
(value_loss * args.value_loss_coef + action_loss -
dist_entropy * args.entropy_coef).backward()
if args.algo == 'a2c':
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
elif args.algo == 'ppo':
advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-5)
for e in range(args.ppo_epoch):
if args.recurrent_policy:
data_generator = rollouts.recurrent_generator(
advantages, args.num_mini_batch)
else:
data_generator = rollouts.feed_forward_generator(
advantages, args.num_mini_batch)
for sample in data_generator:
observations_batch, states_batch, actions_batch, \
return_batch, masks_batch, old_action_log_probs_batch, \
adv_targ = sample
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy, states = actor_critic.evaluate_actions(
Variable(observations_batch), Variable(states_batch),
Variable(masks_batch), Variable(actions_batch))
adv_targ = Variable(adv_targ)
ratio = torch.exp(action_log_probs -
Variable(old_action_log_probs_batch))
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - args.clip_param,
1.0 + args.clip_param) * adv_targ
action_loss = -torch.min(
surr1,
surr2).mean() # PPO's pessimistic surrogate (L^CLIP)
value_loss = (Variable(return_batch) -
values).pow(2).mean()
optimizer.zero_grad()
(value_loss + action_loss -
dist_entropy * args.entropy_coef).backward()
nn.utils.clip_grad_norm(actor_critic.parameters(),
args.max_grad_norm)
optimizer.step()
rollouts.after_update()
if j % args.save_interval == 0 and args.save_dir != "":
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
hasattr(envs, 'ob_rms') and envs.ob_rms or None
]
torch.save(
save_model,
os.path.join(save_path, args.env_name + args.save_tag + ".pt"))
total_test_reward_list = []
step_test_list = []
for _ in range(args.num_tests):
test_obs = test_envs.reset()
update_current_obs(test_obs, test=True)
rollouts_test.observations[0].copy_(current_obs_test)
step_test = 0
total_test_reward = 0
while step_test < args.num_steps_test:
value_test, action_test, action_log_prob_test, states_test = actor_critic.act(
Variable(rollouts_test.observations[step_test],
volatile=True),
Variable(rollouts_test.states[step_test],
volatile=True),
Variable(rollouts_test.masks[step_test],
volatile=True))
cpu_actions_test = action_test.data.squeeze(
1).cpu().numpy()
# Observation, reward and next obs
obs_test, reward_test, done_test, info_test = test_envs.step(
cpu_actions_test)
# masks here doesn't really matter, but still
masks_test = torch.FloatTensor(
[[0.0] if done_test_ else [1.0]
for done_test_ in done_test])
# Maxime: clip the reward within [0,1] for more reliable training
# This code deals poorly with large reward values
reward_test = np.clip(reward_test, a_min=0,
a_max=None) / 400
total_test_reward += reward_test[0]
reward_test = torch.from_numpy(
np.expand_dims(np.stack(reward_test), 1)).float()
update_current_obs(obs_test)
rollouts_test.insert(step_test, current_obs_test, states_test.data, action_test.data, action_log_prob_test.data,\
value_test.data, reward_test, masks_test)
step_test += 1
if done_test:
break
#rollouts_test.reset() # Need to reinitialise with .cuda(); don't forget
total_test_reward_list.append(total_test_reward)
step_test_list.append(step_test)
append_to(tr.test_reward, tr,
sum(total_test_reward_list) / args.num_tests)
append_to(tr.test_episode_len, tr,
sum(step_test_list) / args.num_tests)
logger.log_scalar_rl(
"test_reward", tr.test_reward[0], args.sliding_wsize,
[tr.episodes_done, tr.global_steps_done, tr.iterations_done])
logger.log_scalar_rl(
"test_episode_len", tr.test_episode_len[0], args.sliding_wsize,
[tr.episodes_done, tr.global_steps_done, tr.iterations_done])
# Saving all the MyContainer variables
tr.save(
os.path.join(log_path, args.env_name + args.save_tag + ".p"))
if j % args.log_interval == 0:
reward_avg = 0.99 * reward_avg + 0.01 * final_rewards.mean()
end = time.time()
tr.global_steps_done = (j +
1) * args.num_processes * args.num_steps
print(
"Updates {}, num timesteps {}, FPS {}, running avg reward {:.3f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}"
.format(j, tr.global_steps_done,
int(tr.global_steps_done / (end - start)), reward_avg,
dist_entropy.data[0], value_loss.data[0],
action_loss.data[0]))
append_to(tr.pg_loss, tr, action_loss.data[0])
append_to(tr.val_loss, tr, value_loss.data[0])
append_to(tr.entropy_loss, tr, dist_entropy.data[0])
append_to(tr.train_reward_avg, tr, reward_avg)
logger.log_scalar_rl(
"train_pg_loss", tr.pg_loss[0], args.sliding_wsize,
[tr.episodes_done, tr.global_steps_done, tr.iterations_done])
logger.log_scalar_rl(
"train_val_loss", tr.val_loss[0], args.sliding_wsize,
[tr.episodes_done, tr.global_steps_done, tr.iterations_done])
logger.log_scalar_rl(
"train_entropy_loss", tr.entropy_loss[0], args.sliding_wsize,
[tr.episodes_done, tr.global_steps_done, tr.iterations_done])
logger.log_scalar_rl(
"train_reward_avg", tr.train_reward_avg[0], args.sliding_wsize,
[tr.episodes_done, tr.global_steps_done, tr.iterations_done])
"""
print("Updates {}, num timesteps {}, FPS {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}".
format(
j,
total_num_steps,
int(total_num_steps / (end - start)),
final_rewards.mean(),
final_rewards.median(),
final_rewards.min(),
final_rewards.max(), dist_entropy.data[0],
value_loss.data[0], action_loss.data[0])
)
"""
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, args.env_name,
args.algo)
except IOError:
pass