def update(self, rollouts: RolloutStorage) -> Tuple[float, float, float]:
advantages = self.get_advantages(rollouts)
value_loss_epoch = 0.0
action_loss_epoch = 0.0
dist_entropy_epoch = 0.0
for _e in range(self.ppo_epoch):
profiling_wrapper.range_push("PPO.update epoch")
data_generator = rollouts.recurrent_generator(
advantages, self.num_mini_batch)
for batch in data_generator:
(
values,
action_log_probs,
dist_entropy,
_,
) = self.actor_critic.evaluate_actions(
batch["observations"],
batch["recurrent_hidden_states"],
batch["prev_actions"],
batch["masks"],
batch["actions"],
)
ratio = torch.exp(action_log_probs - batch["action_log_probs"])
surr1 = ratio * batch["advantages"]
surr2 = (torch.clamp(ratio, 1.0 - self.clip_param, 1.0 +
self.clip_param) * batch["advantages"])
action_loss = -(torch.min(surr1, surr2).mean())
if self.use_clipped_value_loss:
value_pred_clipped = batch["value_preds"] + (
values - batch["value_preds"]).clamp(
-self.clip_param, self.clip_param)
value_losses = (values - batch["returns"]).pow(2)
value_losses_clipped = (value_pred_clipped -
batch["returns"]).pow(2)
value_loss = 0.5 * torch.max(value_losses,
value_losses_clipped)
else:
value_loss = 0.5 * (batch["returns"] - values).pow(2)
value_loss = value_loss.mean()
dist_entropy = dist_entropy.mean()
self.optimizer.zero_grad()
total_loss = (value_loss * self.value_loss_coef + action_loss -
dist_entropy * self.entropy_coef)
self.before_backward(total_loss)
total_loss.backward()
self.after_backward(total_loss)
self.before_step()
self.optimizer.step()
self.after_step()
value_loss_epoch += value_loss.item()
action_loss_epoch += action_loss.item()
dist_entropy_epoch += dist_entropy.item()
profiling_wrapper.range_pop() # PPO.update epoch
num_updates = self.ppo_epoch * self.num_mini_batch
value_loss_epoch /= num_updates
action_loss_epoch /= num_updates
dist_entropy_epoch /= num_updates
return value_loss_epoch, action_loss_epoch, dist_entropy_epoch
def _worker_fn(world_rank: int, world_size: int, port: int,
unused_params: bool):
device = (torch.device("cuda")
if torch.cuda.is_available() else torch.device("cpu"))
tcp_store = distrib.TCPStore( # type: ignore
"127.0.0.1", port, world_size, world_rank == 0)
distrib.init_process_group("gloo",
store=tcp_store,
rank=world_rank,
world_size=world_size)
config = get_config("habitat_baselines/config/test/ppo_pointnav_test.yaml")
obs_space = gym.spaces.Dict({
IntegratedPointGoalGPSAndCompassSensor.cls_uuid:
gym.spaces.Box(
low=np.finfo(np.float32).min,
high=np.finfo(np.float32).max,
shape=(2, ),
dtype=np.float32,
)
})
action_space = ActionSpace({"move": EmptySpace()})
actor_critic = PointNavBaselinePolicy.from_config(config, obs_space,
action_space)
# This use adds some arbitrary parameters that aren't part of the computation
# graph, so they will mess up DDP if they aren't correctly ignored by it
if unused_params:
actor_critic.unused = nn.Linear(64, 64)
actor_critic.to(device=device)
ppo_cfg = config.RL.PPO
agent = DDPPO(
actor_critic=actor_critic,
clip_param=ppo_cfg.clip_param,
ppo_epoch=ppo_cfg.ppo_epoch,
num_mini_batch=ppo_cfg.num_mini_batch,
value_loss_coef=ppo_cfg.value_loss_coef,
entropy_coef=ppo_cfg.entropy_coef,
lr=ppo_cfg.lr,
eps=ppo_cfg.eps,
max_grad_norm=ppo_cfg.max_grad_norm,
use_normalized_advantage=ppo_cfg.use_normalized_advantage,
)
agent.init_distributed()
rollouts = RolloutStorage(
ppo_cfg.num_steps,
2,
obs_space,
action_space,
ppo_cfg.hidden_size,
num_recurrent_layers=actor_critic.net.num_recurrent_layers,
is_double_buffered=False,
)
rollouts.to(device)
for k, v in rollouts.buffers["observations"].items():
rollouts.buffers["observations"][k] = torch.randn_like(v)
# Add two steps so batching works
rollouts.advance_rollout()
rollouts.advance_rollout()
# Get a single batch
batch = next(rollouts.recurrent_generator(rollouts.buffers["returns"], 1))
# Call eval actions through the internal wrapper that is used in
# agent.update
value, action_log_probs, dist_entropy, _ = agent._evaluate_actions(
batch["observations"],
batch["recurrent_hidden_states"],
batch["prev_actions"],
batch["masks"],
batch["actions"],
)
# Backprop on things
(value.mean() + action_log_probs.mean() + dist_entropy.mean()).backward()
# Make sure all ranks have very similar parameters
for param in actor_critic.parameters():
if param.grad is not None:
grads = [param.grad.detach().clone() for _ in range(world_size)]
distrib.all_gather(grads, grads[world_rank])
for i in range(world_size):
assert torch.isclose(grads[i], grads[world_rank]).all()
def update(self, rollouts: RolloutStorage) -> Tuple[float, float, float]:
advantages = self.get_advantages(rollouts)
value_loss_epoch = 0.0
action_loss_epoch = 0.0
dist_entropy_epoch = 0.0
for _e in range(self.ppo_epoch):
data_generator = rollouts.recurrent_generator(
advantages, self.num_mini_batch)
for sample in data_generator:
(
obs_batch,
recurrent_hidden_states_batch,
actions_batch,
prev_actions_batch,
value_preds_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,
_,
) = self.actor_critic.evaluate_actions(
obs_batch,
recurrent_hidden_states_batch,
prev_actions_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.clip_param,
1.0 + self.clip_param) * adv_targ)
action_loss = -torch.min(surr1, surr2).mean()
if self.use_clipped_value_loss:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(
-self.clip_param, self.clip_param)
value_losses = (values - return_batch).pow(2)
value_losses_clipped = (value_pred_clipped -
return_batch).pow(2)
value_loss = (
0.5 *
torch.max(value_losses, value_losses_clipped).mean())
else:
value_loss = 0.5 * (return_batch - values).pow(2).mean()
self.optimizer.zero_grad()
total_loss = (value_loss * self.value_loss_coef + action_loss -
dist_entropy * self.entropy_coef)
self.before_backward(total_loss)
total_loss.backward()
self.after_backward(total_loss)
self.before_step()
self.optimizer.step()
self.after_step()
value_loss_epoch += value_loss.item()
action_loss_epoch += action_loss.item()
dist_entropy_epoch += dist_entropy.item()
num_updates = self.ppo_epoch * self.num_mini_batch
value_loss_epoch /= num_updates
action_loss_epoch /= num_updates
dist_entropy_epoch /= num_updates
return value_loss_epoch, action_loss_epoch, dist_entropy_epoch