if self.training_type == "rl": if "critic" in extras["observations"]: self.privileged_obs_type = "critic" # actor-critic reinforcement learnig, e.g., PPO else: self.privileged_obs_type = None if self.training_type == "distillation": if "teacher" in extras["observations"]: self.privileged_obs_type = "teacher" # policy distillation else: self.privileged_obs_type = None
# resolve dimensions of privileged observations if self.privileged_obs_type is not None: num_privileged_obs = extras["observations"][self.privileged_obs_type].shape[1] else: num_privileged_obs = num_obs
if"rnd_cfg"inself.alg_cfg andself.alg_cfg["rnd_cfg"] isnotNone: # check if rnd gated state is present rnd_state = extras["observations"].get("rnd_state") if rnd_state isNone: raise ValueError("Observations for the key 'rnd_state' not found in infos['observations'].") # get dimension of rnd gated state num_rnd_state = rnd_state.shape[1] # add rnd gated state to config self.alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state # scale down the rnd weight with timestep (similar to how rewards are scaled down in legged_gym envs) self.alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt
# if using symmetry then pass the environment config object 对称 if"symmetry_cfg"inself.alg_cfg andself.alg_cfg["symmetry_cfg"] isnotNone: # this is used by the symmetry function for handling different observation terms self.alg_cfg["symmetry_cfg"]["_env"] = env
# number of augmentations per sample # we start with 1 and increase it if we use symmetry augmentation num_aug = 1 # original batch size 一个batch有多少数据 original_batch_size = obs_batch.shape[0]
# check if we should normalize advantages per mini batch 优势函数标准化 ifself.normalize_advantage_per_mini_batch: with torch.no_grad(): advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8)
# Perform symmetric augmentation 对称数据增强 ifself.symmetry andself.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"], obs_type="policy" ) critic_obs_batch, _ = data_augmentation_func( obs=critic_obs_batch, actions=None, env=self.symmetry["_env"], obs_type="critic" ) # compute number of augmentations per sample num_aug = int(obs_batch.shape[0] / original_batch_size) # repeat the rest of the batch 把其他的张量也重复 num_aug 次，保证每个增强样本都有对应的标签,且其他张量在镜像前后一致，所以直接复制就行 # 扩容前大家都是b行，扩容后刚刚说的obs_batch actions_batch critic_obs_batch 变成2B行， # -- 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(critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) # -- entropy # we only keep the entropy of the first augmentation (the original one) 只用原始样本计算新旧策略的 KL和entropy regularization # 注意这三个的更新步骤，他们在actor_critic.py里定义了@property，实际上返回的是self.distribution的相关计算结果， # 而self.distribution是在self.policy.act里面执行update_distribution的时候更新的 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]
# Reduce the KL divergence across all GPUs ifself.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. ifself.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.0and kl_mean > 0.0: self.learning_rate = min(1e-2, self.learning_rate * 1.5)
# 多 GPU 中，rank 0 是主进程。这里只让主进程执行学习率判断，避免多个进程分别修改学习率。然后修改完后同步 # Update the learning rate for all GPUs ifself.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 inself.optimizer.param_groups: param_group["lr"] = self.learning_rate
# Symmetry loss 镜像损失 ifself.symmetry: # obtain the symmetric actions # if we did augmentation before then we don't need to augment again 前面做过就不需要再做一次镜像（数据增强）了
ifnotself.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"], obs_type="policy" ) # 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 # act_inference() 通常输出 deterministic action，也就是 actor 的均值动作，不采样。 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"], obs_type="policy" )
# 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 ifself.symmetry["use_mirror_loss"]: loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss else: symmetry_loss = symmetry_loss.detach()
# Random Network Distillation loss 注意这里它使用独立的self.rnd_optimizer，所以没有加到ppo的loss里 # ifself.rnd: # 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)
surrogate loss、entropy loss 更新 actor value loss 更新 critic mirror loss 更新 actor AMP loss、gradient penalty 更新 discriminator RND loss 使用独立 optimizer 更新 RND predictor
motion_visualization/walk.txt [736, 52]，FrameDuration=0.033 s motion_amp_expert/walk.txt [735, 52]，FrameDuration=0.033 s motion_visualization/run.txt [345, 52]，FrameDuration=0.033 s motion_amp_expert/run.txt [344, 52]，FrameDuration=0.033 s
