# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations import torch import torch.nn as nn import torch.optim as optim from itertools import chain from rsl_rl.modules import ActorCritic from rsl_rl.modules.rnd import RandomNetworkDistillation from rsl_rl.storage import RolloutStorage from rsl_rl.utils import string_to_callable class PPO: """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" policy: ActorCritic """The actor critic module.""" def __init__( self, policy, num_learning_epochs=5, num_mini_batches=4, clip_param=0.2, gamma=0.99, lam=0.95, value_loss_coef=1.0, entropy_coef=0.01, learning_rate=0.001, max_grad_norm=1.0, use_clipped_value_loss=True, schedule="adaptive", desired_kl=0.01, device="cpu", normalize_advantage_per_mini_batch=False, # RND parameters rnd_cfg: dict | None = None, # Symmetry parameters symmetry_cfg: dict | None = None, # Distributed training parameters multi_gpu_cfg: dict | None = None, ): # device-related parameters self.device = device self.is_multi_gpu = multi_gpu_cfg is not None # Multi-GPU parameters if multi_gpu_cfg is not None: self.gpu_global_rank = multi_gpu_cfg["global_rank"] self.gpu_world_size = multi_gpu_cfg["world_size"] else: self.gpu_global_rank = 0 self.gpu_world_size = 1 # RND components if rnd_cfg is not None: # Extract parameters used in ppo rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) # Create RND module self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) # Create RND optimizer params = self.rnd.predictor.parameters() self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) else: self.rnd = None self.rnd_optimizer = None # Symmetry components if symmetry_cfg is not None: # Check if symmetry is enabled use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] # Print that we are not using symmetry if not use_symmetry: print("Symmetry not used for learning. We will use it for logging instead.") # If function is a string then resolve it to a function if isinstance(symmetry_cfg["data_augmentation_func"], str): symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) # Check valid configuration if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): raise ValueError( "Data augmentation enabled but the function is not callable:" f" {symmetry_cfg['data_augmentation_func']}" ) # Store symmetry configuration self.symmetry = symmetry_cfg else: self.symmetry = None # PPO components self.policy = policy self.policy.to(self.device) # Create optimizer self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) # Create rollout storage self.storage: RolloutStorage = None # type: ignore self.transition = RolloutStorage.Transition() # PPO parameters self.clip_param = clip_param self.num_learning_epochs = num_learning_epochs self.num_mini_batches = num_mini_batches self.value_loss_coef = value_loss_coef self.entropy_coef = entropy_coef self.gamma = gamma self.lam = lam self.max_grad_norm = max_grad_norm self.use_clipped_value_loss = use_clipped_value_loss self.desired_kl = desired_kl self.schedule = schedule self.learning_rate = learning_rate self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): # create rollout storage self.storage = RolloutStorage( training_type, num_envs, num_transitions_per_env, obs, actions_shape, self.device, ) def act(self, obs): if self.policy.is_recurrent: self.transition.hidden_states = self.policy.get_hidden_states() # compute the actions and values self.transition.actions = self.policy.act(obs).detach() self.transition.values = self.policy.evaluate(obs).detach() self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() self.transition.action_mean = self.policy.action_mean.detach() self.transition.action_sigma = self.policy.action_std.detach() # need to record obs before env.step() self.transition.observations = obs return self.transition.actions def process_env_step(self, obs, rewards, dones, extras): # update the normalizers self.policy.update_normalization(obs) if self.rnd: self.rnd.update_normalization(obs) # Record the rewards and dones # Note: we clone here because later on we bootstrap the rewards based on timeouts self.transition.rewards = rewards.clone() self.transition.dones = dones # Compute the intrinsic rewards and add to extrinsic rewards if self.rnd: # Compute the intrinsic rewards self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) # Add intrinsic rewards to extrinsic rewards self.transition.rewards += self.intrinsic_rewards # Bootstrapping on time outs if "time_outs" in extras: self.transition.rewards += self.gamma * torch.squeeze( self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 ) # record the transition self.storage.add_transitions(self.transition) self.transition.clear() self.policy.reset(dones) def compute_returns(self, obs): # compute value for the last step last_values = self.policy.evaluate(obs).detach() self.storage.compute_returns( last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch ) def update(self): # noqa: C901 mean_value_loss = 0 mean_surrogate_loss = 0 mean_entropy = 0 # -- RND loss if self.rnd: mean_rnd_loss = 0 else: mean_rnd_loss = None # -- Symmetry loss if self.symmetry: mean_symmetry_loss = 0 else: mean_symmetry_loss = None # generator for mini batches if self.policy.is_recurrent: generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) else: generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) # iterate over batches for ( obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, hid_states_batch, masks_batch, ) in generator: # number of augmentations per sample # we start with 1 and increase it if we use symmetry augmentation num_aug = 1 # original batch size # we assume policy group is always there and needs augmentation original_batch_size = obs_batch.batch_size[0] # check if we should normalize advantages per mini batch if self.normalize_advantage_per_mini_batch: with torch.no_grad(): advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) # Perform symmetric augmentation if self.symmetry and self.symmetry["use_data_augmentation"]: # augmentation using symmetry data_augmentation_func = self.symmetry["data_augmentation_func"] # returned shape: [batch_size * num_aug, ...] obs_batch, actions_batch = data_augmentation_func( obs=obs_batch, actions=actions_batch, env=self.symmetry["_env"], ) # compute number of augmentations per sample # we assume policy group is always there and needs augmentation num_aug = int(obs_batch.batch_size[0] / original_batch_size) # repeat the rest of the batch # -- actor old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) # -- critic target_values_batch = target_values_batch.repeat(num_aug, 1) advantages_batch = advantages_batch.repeat(num_aug, 1) returns_batch = returns_batch.repeat(num_aug, 1) # Recompute actions log prob and entropy for current batch of transitions # Note: we need to do this because we updated the policy with the new parameters # -- actor self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) # -- critic value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) # -- entropy # we only keep the entropy of the first augmentation (the original one) mu_batch = self.policy.action_mean[:original_batch_size] sigma_batch = self.policy.action_std[:original_batch_size] entropy_batch = self.policy.entropy[:original_batch_size] # KL if self.desired_kl is not None and self.schedule == "adaptive": with torch.inference_mode(): kl = torch.sum( torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1, ) kl_mean = torch.mean(kl) # Reduce the KL divergence across all GPUs if self.is_multi_gpu: torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) kl_mean /= self.gpu_world_size # Update the learning rate # Perform this adaptation only on the main process # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, # then the learning rate should be the same across all GPUs. if self.gpu_global_rank == 0: if kl_mean > self.desired_kl * 2.0: self.learning_rate = max(1e-5, self.learning_rate / 1.5) elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: self.learning_rate = min(1e-2, self.learning_rate * 1.5) # Update the learning rate for all GPUs if self.is_multi_gpu: lr_tensor = torch.tensor(self.learning_rate, device=self.device) torch.distributed.broadcast(lr_tensor, src=0) self.learning_rate = lr_tensor.item() # Update the learning rate for all parameter groups for param_group in self.optimizer.param_groups: param_group["lr"] = self.learning_rate # Surrogate loss ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) surrogate = -torch.squeeze(advantages_batch) * ratio surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( ratio, 1.0 - self.clip_param, 1.0 + self.clip_param ) surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() # Value function loss if self.use_clipped_value_loss: value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( -self.clip_param, self.clip_param ) value_losses = (value_batch - returns_batch).pow(2) value_losses_clipped = (value_clipped - returns_batch).pow(2) value_loss = torch.max(value_losses, value_losses_clipped).mean() else: value_loss = (returns_batch - value_batch).pow(2).mean() loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() # Symmetry loss if self.symmetry: # obtain the symmetric actions # if we did augmentation before then we don't need to augment again if not self.symmetry["use_data_augmentation"]: data_augmentation_func = self.symmetry["data_augmentation_func"] obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) # compute number of augmentations per sample num_aug = int(obs_batch.shape[0] / original_batch_size) # actions predicted by the actor for symmetrically-augmented observations mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) # compute the symmetrically augmented actions # note: we are assuming the first augmentation is the original one. # We do not use the action_batch from earlier since that action was sampled from the distribution. # However, the symmetry loss is computed using the mean of the distribution. action_mean_orig = mean_actions_batch[:original_batch_size] _, actions_mean_symm_batch = data_augmentation_func( obs=None, actions=action_mean_orig, env=self.symmetry["_env"] ) # compute the loss (we skip the first augmentation as it is the original one) mse_loss = torch.nn.MSELoss() symmetry_loss = mse_loss( mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] ) # add the loss to the total loss if self.symmetry["use_mirror_loss"]: loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss else: symmetry_loss = symmetry_loss.detach() # Random Network Distillation loss # TODO: Move this processing to inside RND module. if self.rnd: # extract the rnd_state # TODO: Check if we still need torch no grad. It is just an affine transformation. with torch.no_grad(): rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) # predict the embedding and the target predicted_embedding = self.rnd.predictor(rnd_state_batch) target_embedding = self.rnd.target(rnd_state_batch).detach() # compute the loss as the mean squared error mseloss = torch.nn.MSELoss() rnd_loss = mseloss(predicted_embedding, target_embedding) # Compute the gradients # -- For PPO self.optimizer.zero_grad() loss.backward() # -- For RND if self.rnd: self.rnd_optimizer.zero_grad() # type: ignore rnd_loss.backward() # Collect gradients from all GPUs if self.is_multi_gpu: self.reduce_parameters() # Apply the gradients # -- For PPO nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.optimizer.step() # -- For RND if self.rnd_optimizer: self.rnd_optimizer.step() # Store the losses mean_value_loss += value_loss.item() mean_surrogate_loss += surrogate_loss.item() mean_entropy += entropy_batch.mean().item() # -- RND loss if mean_rnd_loss is not None: mean_rnd_loss += rnd_loss.item() # -- Symmetry loss if mean_symmetry_loss is not None: mean_symmetry_loss += symmetry_loss.item() # -- For PPO num_updates = self.num_learning_epochs * self.num_mini_batches mean_value_loss /= num_updates mean_surrogate_loss /= num_updates mean_entropy /= num_updates # -- For RND if mean_rnd_loss is not None: mean_rnd_loss /= num_updates # -- For Symmetry if mean_symmetry_loss is not None: mean_symmetry_loss /= num_updates # -- Clear the storage self.storage.clear() # construct the loss dictionary loss_dict = { "value_function": mean_value_loss, "surrogate": mean_surrogate_loss, "entropy": mean_entropy, } if self.rnd: loss_dict["rnd"] = mean_rnd_loss if self.symmetry: loss_dict["symmetry"] = mean_symmetry_loss return loss_dict """ Helper functions """ def broadcast_parameters(self): """Broadcast model parameters to all GPUs.""" # obtain the model parameters on current GPU model_params = [self.policy.state_dict()] if self.rnd: model_params.append(self.rnd.predictor.state_dict()) # broadcast the model parameters torch.distributed.broadcast_object_list(model_params, src=0) # load the model parameters on all GPUs from source GPU self.policy.load_state_dict(model_params[0]) if self.rnd: self.rnd.predictor.load_state_dict(model_params[1]) def reduce_parameters(self): """Collect gradients from all GPUs and average them. This function is called after the backward pass to synchronize the gradients across all GPUs. """ # Create a tensor to store the gradients grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] if self.rnd: grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] all_grads = torch.cat(grads) # Average the gradients across all GPUs torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) all_grads /= self.gpu_world_size # Get all parameters all_params = self.policy.parameters() if self.rnd: all_params = chain(all_params, self.rnd.parameters()) # Update the gradients for all parameters with the reduced gradients offset = 0 for param in all_params: if param.grad is not None: numel = param.numel() # copy data back from shared buffer param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) # update the offset for the next parameter offset += numel