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
class PPOAgent(ResearchAgent):
"""The TensorForceAgent. Acts through the algorith, not here."""
def __init__(self, actor_critic, character=characters.Bomber, **kwargs):
self._actor_critic = actor_critic
super(PPOAgent, self).__init__(character, **kwargs)
def cuda(self):
self._actor_critic.cuda()
if hasattr(self, "_rollout"):
self._rollout.cuda()
@property
def model(self):
return self._actor_critic
@property
def optimizer(self):
return self._optimizer
def set_eval(self):
self._actor_critic.eval()
def set_train(self):
self._actor_critic.train()
def _rollout_data(self, step, num_agent, num_agent_end=None):
if num_agent_end is not None:
assert (num_agent_end > num_agent)
observations = Variable(
self._rollout.observations[step, num_agent:num_agent_end])
states = Variable(self._rollout.states[step,
num_agent:num_agent_end])
masks = Variable(self._rollout.masks[step,
num_agent:num_agent_end])
else:
observations = Variable(self._rollout.observations[step,
num_agent],
volatile=True)
states = Variable(self._rollout.states[step, num_agent],
volatile=True)
masks = Variable(self._rollout.masks[step, num_agent],
volatile=True)
return observations, states, masks
def actor_critic_act(self, step, num_agent=0, deterministic=False):
"""Uses the actor_critic to take action.
Args:
step: The int timestep that we are acting.
num_agent: Agent id that's running. Non-zero when agent has copies.
Returns:
See the actor_critic's act function in model.py.
"""
# NOTE: Training uses this --> it uses act(..., deterministic=False).
return self._actor_critic.act(*self.get_rollout_data(step, num_agent),
deterministic=deterministic)
def get_rollout_data(self, step, num_agent, num_agent_end=None):
return self._rollout_data(step, num_agent, num_agent_end)
def actor_critic_call(self, step, num_agent=0):
observations, states, masks = self._rollout_data(step, num_agent)
return self._actor_critic(observations, states, masks)[0].data
def _evaluate_actions(self, observations, states, masks, actions):
return self._actor_critic.evaluate_actions(observations, states, masks,
actions)
def _optimize(self,
value_loss,
action_loss,
dist_entropy,
entropy_coef,
value_loss_coef,
max_grad_norm,
kl_loss=None,
kl_factor=0,
only_value_loss=False,
add_nonlin=False):
self._optimizer.zero_grad()
# only update the value head (to be used when fine tuning a model
# trained with BC without a value predictor) -- only beginning of finetuning
if only_value_loss:
loss = value_loss * value_loss_coef
# stop the gradients from flowing through the
# parameters that are used to compute the actions (i.e. critic / policy head)
# and only backprop the value loss through the value head
#(i.e. parameters used exclusively to predict the value)
for p in self._actor_critic.parameters():
p.requires_grad = False
self._actor_critic.critic_linear.requires_grad = True
if add_nonlin:
self._actor_critic.fc_critic.requires_grad = True
loss.backward()
else:
loss = value_loss * value_loss_coef + action_loss \
- dist_entropy * entropy_coef
if kl_factor > 0 and not use_is:
loss += kl_factor * kl_loss
loss.backward()
nn.utils.clip_grad_norm(self._actor_critic.parameters(), max_grad_norm)
self._optimizer.step()
if hasattr(self, '_scheduler'):
self._scheduler.step(loss)
if only_value_loss:
for p in self._actor_critic.parameters():
p.requires_grad = True
def halve_lr(self):
for i, param_group in enumerate(self._optimizer.param_groups):
old_lr = float(param_group['lr'])
new_lr = max(old_lr * 0.5, 1e-7)
param_group['lr'] = new_lr
def compute_advantages(self, next_value_agents, use_gae, gamma, tau):
for num_agent, next_value in enumerate(next_value_agents):
self._rollout.compute_returns(next_value, use_gae, gamma, tau,
num_agent)
advantages = self._rollout.compute_advantages()
diff = (advantages - advantages.mean())
advantages = diff / (advantages.std() + 1e-5)
return advantages
def initialize(self, args, obs_shape, action_space,
num_training_per_episode, num_episodes, total_steps,
num_epoch, optimizer_state_dict, num_steps, uniform_v,
uniform_v_prior):
params = self._actor_critic.parameters()
self._optimizer = optim.Adam(params, lr=args.lr, eps=args.eps)
if optimizer_state_dict:
self._optimizer.load_state_dict(optimizer_state_dict)
if args.use_lr_scheduler:
self._scheduler = optim.lr_scheduler.ReduceLROnPlateau(
self._optimizer, mode='min', verbose=True)
self._rollout = RolloutStorage(num_steps, args.num_processes,
obs_shape, action_space,
self._actor_critic.state_size,
num_training_per_episode)
self.num_episodes = num_episodes
self.total_steps = total_steps
self.num_epoch = num_epoch
self.uniform_v = uniform_v
self.uniform_v_prior = uniform_v_prior
def update_rollouts(self, obs, timestep):
self._rollout.observations[timestep, :, :, :, :, :].copy_(obs)
def insert_rollouts(self,
step,
current_obs,
states,
action,
action_log_prob,
value,
reward,
mask,
action_log_prob_distr=None,
dagger_prob_distr=None,
expert_action_log_prob=None,
training_action_log_prob=None):
self._rollout.insert(step,
current_obs,
states,
action,
action_log_prob,
value,
reward,
mask,
action_log_prob_distr,
dagger_prob_distr,
expert_action_log_prob=None,
training_action_log_prob=None)
def ppo(self,
advantages,
num_mini_batch,
batch_size,
num_steps,
clip_param,
entropy_coef,
value_loss_coef,
max_grad_norm,
action_space,
anneal=False,
lr=1e-4,
eps=1e-5,
kl_factor=0,
only_value_loss=False,
add_nonlin=False,
use_is=False,
use_retrace=False,
lambda_retrace=1.0):
action_losses = []
value_losses = []
dist_entropies = []
kl_losses = []
kl_loss = None
total_losses = []
if hasattr(self._actor_critic, 'gru'):
data_generator = self._rollout.recurrent_generator(
advantages, num_mini_batch, batch_size, num_steps, kl_factor,
use_is)
else:
data_generator = self._rollout.feed_forward_generator(
advantages, num_mini_batch, batch_size, num_steps,
action_space, kl_factor, use_is)
for sample in data_generator:
observations_batch, states_batch, actions_batch, return_batch, \
masks_batch, old_action_log_probs_batch, adv_targ, \
action_log_probs_distr_batch, dagger_probs_distr_batch, \
expert_action_log_probs_batch, training_action_log_probs_batch \
= sample
# Reshape to do in a single forward pass for all steps
result = self._evaluate_actions(Variable(observations_batch),
Variable(states_batch),
Variable(masks_batch),
Variable(actions_batch))
values, action_log_probs, dist_entropy, states = result
adv_targ = Variable(adv_targ)
ratio = action_log_probs
ratio -= Variable(old_action_log_probs_batch)
ratio = torch.exp(ratio)
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - clip_param, 1.0 + clip_param)
surr2 *= adv_targ
action_loss = -torch.min(surr1, surr2).mean()
value_loss = (Variable(return_batch) - values) \
.pow(2).mean()
total_loss = value_loss * value_loss_coef + action_loss \
- dist_entropy * entropy_coef
if kl_factor > 0 and not use_is:
criterion = nn.KLDivLoss()
kl_loss = criterion(Variable(action_log_probs_distr_batch),
Variable(dagger_probs_distr_batch))
total_loss += kl_factor * kl_loss
self._optimize(value_loss, action_loss, dist_entropy, entropy_coef,
value_loss_coef, max_grad_norm, kl_loss, kl_factor,
only_value_loss, add_nonlin)
lr = self._optimizer.param_groups[0]['lr']
action_losses.append(action_loss.data[0])
value_losses.append(value_loss.data[0])
dist_entropies.append(dist_entropy.data[0])
if kl_factor > 0 and not use_is:
kl_losses.append(kl_loss.data[0])
total_losses.append(total_loss.data[0])
return action_losses, value_losses, dist_entropies, \
kl_losses, total_losses, lr
def copy_ex_model(self):
"""Creates a copy without the model. This is for operating with homogenous training."""
return PPOAgent(None,
self._character,
num_processes=self._num_processes)
def copy_with_model(self):
"""Creates a copy with the model. This is for operating with frozen backplay."""
return PPOAgent(self._actor_critic,
self._character,
num_processes=self._num_processes)
def after_epoch(self):
self._rollout.after_epoch()
def set_new_model(self, model, cuda=False):
self._actor_critic = model
if cuda:
self._actor_critic.cuda()
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]))
class PPOAgent(object):
def __init__(self,args):
self.args = args
self.device = torch.device('cuda') if args.cuda else torch.device('cpu')
dummy_env = gym.make(self.args.env_name)
self.actor = ACNet(dummy_env.action_space.n,args.feedforward)
del dummy_env
if args.load_dir is not None:
actorState = torch.load(args.load_dir,map_location=lambda storage, loc: storage)
if args.continue_training:
self.actor.load_state_dict(actorState)
print("Loaded pretrained model successfully")
if args.transfer:
self.actor.load_autoturn_model(actorState)
if args.cuda:
self.actor.cuda()
self.actor_optimizer = optim.Adam(self.actor.parameters(),lr=self.args.lr)
self.env_list = [make_env(self.args.env_name,self.args.seed,i) for i in range(self.args.num_processes)]
if self.args.num_processes > 1:
self.envs = gym_vecenv.SubprocVecEnv(self.env_list)
else:
self.envs = gym_vecenv.DummyVecEnv(self.env_list)
if len(self.envs.observation_space.shape) == 1:
self.envs = gym_vecenv.VecNormalize(self.envs)
self.obs_shape = self.envs.observation_space.shape
self.obs_shape = (self.obs_shape[0] * args.num_stack, *self.obs_shape[1:])
self.state_shape = 1 if args.feedforward else 256
self.rollouts = RolloutStorage(self.args.num_fwd_steps, self.args.num_processes, self.obs_shape, self.envs.action_space, self.state_shape)
self.num_updates = int(args.num_frames)//args.num_fwd_steps//args.num_processes
self.current_obs = torch.zeros(self.args.num_processes,*self.obs_shape)
self.writer = SummaryWriter(log_dir=self.args.save_dir)
self.fortress_threshold = 650
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(self.actor_optimizer,
mode='max',factor=0.2,patience=15,verbose=True,threshold=1e-3,
threshold_mode='rel')
#self.scheduler2 = torch.optim.lr_scheduler.MultiStepLR(self.actor_optimizer,milestones=[40,80],gamma=0.3)
def update_current_obs(self,obs):
shape_dim0 = self.envs.observation_space.shape[0]
obs = torch.from_numpy(obs).float()
if self.args.num_stack > 1:
self.current_obs[:, :-shape_dim0] = self.current_obs[:, shape_dim0:]
self.current_obs[:, -shape_dim0:] = obs
def train(self):
obs = self.envs.reset()
self.update_current_obs(obs)
self.rollouts.observations[0].copy_(self.current_obs)
episode_rewards = torch.zeros([self.args.num_processes,1])
final_rewards = torch.zeros([self.args.num_processes,1])
num_destruction = 0
if self.args.cuda:
self.current_obs = self.current_obs.cuda()
self.rollouts.cuda()
start = time.time()
num_episodes = 0
for iteration in range(self.num_updates):
for step in range(self.args.num_fwd_steps):
with torch.no_grad():
value,action,action_log_prob,states = self.actor.act(self.rollouts.observations[step],
self.rollouts.states[step],
self.rollouts.masks[step])
cpu_actions = action.squeeze(1).cpu().numpy()
obs,reward,done,info = self.envs.step(cpu_actions)
num_destruction += sum(info)
reward = torch.from_numpy(np.expand_dims(np.stack(reward),1)).float()
episode_rewards += reward
masks = torch.FloatTensor([[0.0] if i else [1.0] for i in done])
final_rewards*=masks
final_rewards += (1-masks)*episode_rewards
episode_rewards *= masks
if self.args.cuda:
masks = masks.cuda()
if self.current_obs.dim() == 4:
self.current_obs *= masks.unsqueeze(2).unsqueeze(2)
else:
self.current_obs *= masks
self.update_current_obs(obs)
self.rollouts.insert(step,self.current_obs,states,action,action_log_prob,value,reward,masks)
with torch.no_grad():
next_value = self.actor.get_value(self.rollouts.observations[-1],
self.rollouts.states[-1],
self.rollouts.masks[-1]).detach()
self.rollouts.compute_returns(next_value,True,self.args.gamma,self.args.tau)
if not self.args.a2c:
advantages = self.rollouts.returns[:-1] - self.rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean())/(advantages.std()+1e-5)
for i in range(self.args.ppo_epoch):
if self.args.feedforward:
data_generator = self.rollouts.feed_forward_generator(advantages,self.args.num_mini_batch)
else:
data_generator = self.rollouts.recurrent_generator(advantages,self.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
values,action_log_probs,dist_entropy,states = self.actor.evaluate_actions(observations_batch,
states_batch,
masks_batch,
actions_batch)
ratio = torch.exp(action_log_probs - old_action_log_probs_batch)
surr1 = ratio*adv_targ
surr2 = torch.clamp(ratio,1.0-self.args.clip_param,1.0+self.args.clip_param)*adv_targ
action_loss = -torch.min(surr1,surr2).mean()
value_loss = (values-return_batch).pow(2).mean()
self.actor_optimizer.zero_grad()
actorLoss = action_loss + self.args.value_loss_coeff*value_loss - self.args.entropy_coeff*dist_entropy
actorLoss.backward()
torch.nn.utils.clip_grad_norm_(self.actor.parameters(),self.args.max_grad_norm)
self.actor_optimizer.step()
else:
values,action_log_probs,dist_entropy,states = self.actor.evaluate_actions(self.rollouts.observations[:-1].view(-1,*self.obs_shape),
self.rollouts.states[0].view(-1,1 if self.args.feedforward else 256),
self.rollouts.masks[:-1].view(-1,1),
self.rollouts.actions.view(-1,1))
values = values.view(self.args.num_fwd_steps,self.args.num_processes,1)
action_log_probs = action_log_probs.view(self.args.num_fwd_steps,self.args.num_processes,1)
advantages = self.rollouts.returns[:-1] - values
value_loss = advantages.pow(2).mean()
action_loss = -(advantages.detach()*action_log_probs).mean()
self.actor_optimizer.zero_grad()
actorLoss = action_loss + self.args.value_loss_coeff*value_loss - self.args.entropy_coeff*dist_entropy
actorLoss.backward()
self.actor_optimizer.step()
self.rollouts.after_update()
if num_destruction>self.fortress_threshold:
torch.save(self.actor.state_dict(),self.args.save_dir+'/'+self.args.env_name+'_'+str(iteration)+'_ppo_actor.pth.tar')
self.fortress_threshold = num_destruction
if iteration%self.args.log_interval == 0:
end = time.time()
total_num_steps = (iteration+1)*self.args.num_processes*self.args.num_fwd_steps
num_destruction /= self.args.num_processes
print("Updates {}, num timesteps {}, FPS {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}, num fortress destroyed {:.2f}".
format(iteration, total_num_steps,
int(total_num_steps / (end - start)),
final_rewards.mean(),
final_rewards.median(),
final_rewards.min(),
final_rewards.max(), dist_entropy.item(),
value_loss.item(), action_loss.item(),num_destruction))
self.writer.add_scalar('data/rewardmean',final_rewards.mean(),total_num_steps)
self.writer.add_scalar('data/distentropy',dist_entropy.item(),total_num_steps)
self.writer.add_scalar('data/valueloss',value_loss.item(),total_num_steps)
self.writer.add_scalar('data/actionloss',action_loss.item(),total_num_steps)
self.writer.add_scalar('data/numdestruction',num_destruction,total_num_steps)
#self.scheduler.step(final_rewards.mean())i
#self.scheduler2.step()
num_destruction = 0
if iteration%self.args.save_interval==0:
torch.save(self.actor.state_dict(),self.args.save_dir+'/'+self.args.env_name+'_'+str(iteration)+'_ppo_actor.pth.tar')
self.writer.export_scalars_to_json(self.args.save_dir+'/'+self.args.env_name+"_all_scalars.json")
self.envs.close()
self.writer.close()
hx, cx = rollouts.states[-1].split(model.state_size, 1)
next_value = model.forward(
(Variable(rollouts.observations[-1],
volatile=True), (Variable(hx), Variable(cx)),
Variable(rollouts.masks[-1], volatile=True)))[0].data
rollouts.compute_returns(next_value, args.gamma, args.tau)
# After constructing points, now train points
advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
advantages = (advantages - advantages.mean()) / (advantages.std() +
1e-5)
for e in range(args.ppo_epoch):
data_generator = rollouts.recurrent_generator(
advantages, args.minibatch)
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
hx, cx = states_batch.split(model.state_size, 1)
values, logits, _, states = model.forward(
(Variable(observations_batch),
(Variable(hx), Variable(cx)), Variable(masks_batch)))
# Compute action probabilities and entropy distance
action_log_probs, dist_entropy = logprobs_and_entropy(
logits, Variable(actions_batch))
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():
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
