PPO算法拆解

最新推荐文章于 2025-12-21 15:46:06 发布

原创最新推荐文章于 2025-12-21 15:46:06 发布 · 1k 阅读

31 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #python #强化学习

强化学习专栏收录该内容

2 篇文章

订阅专栏

在强化学习legged_gym中rsl_rl所实现的PPO算法源码如下：

# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import torch
import torch.nn as nn
import torch.optim as optim

from rsl_rl.modules import ActorCritic
from rsl_rl.storage import RolloutStorage

class PPO:
    actor_critic: ActorCritic
    def __init__(self,
                 actor_critic,
                 num_learning_epochs=1,
                 num_mini_batches=1,
                 clip_param=0.2,
                 gamma=0.998,
                 lam=0.95,
                 value_loss_coef=1.0,
                 entropy_coef=0.0,
                 learning_rate=1e-3,
                 max_grad_norm=1.0,
                 use_clipped_value_loss=True,
                 schedule="fixed",
                 desired_kl=0.01,
                 device='cpu',
                 ):

        self.device = device

        self.desired_kl = desired_kl
        self.schedule = schedule
        self.learning_rate = learning_rate

        # PPO components
        self.actor_critic = actor_critic
        self.actor_critic.to(self.device)
        self.storage = None # initialized later
        self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
        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

    def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
        self.storage = RolloutStorage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device)

    def test_mode(self):
        self.actor_critic.test()
    
    def train_mode(self):
        self.actor_critic.train()

    def act(self, obs, critic_obs):
        if self.actor_critic.is_recurrent:
            self.transition.hidden_states = self.actor_critic.get_hidden_states()
        # Compute the actions and values
        self.transition.actions = self.actor_critic.act(obs).detach()
        self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
        self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach()
        self.transition.action_mean = self.actor_critic.action_mean.detach()
        self.transition.action_sigma = self.actor_critic.action_std.detach()
        # need to record obs and critic_obs before env.step()
        self.transition.observations = obs
        self.transition.critic_observations = critic_obs
        return self.transition.actions
    
    def process_env_step(self, rewards, dones, infos):
        self.transition.rewards = rewards.clone()
        self.transition.dones = dones
        # Bootstrapping on time outs
        if 'time_outs' in infos:
            self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1)

        # Record the transition
        self.storage.add_transitions(self.transition)
        self.transition.clear()
        self.actor_critic.reset(dones)
    
    def compute_returns(self, last_critic_obs):
        last_values= self.actor_critic.evaluate(last_critic_obs).detach()
        self.storage.compute_returns(last_values, self.gamma, self.lam)

    def update(self):
        mean_value_loss = 0
        mean_surrogate_loss = 0
        if self.actor_critic.is_recurrent:
            generator = self.storage.reccurent_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)
        for obs_batch, critic_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:


                self.actor_critic.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0])
                actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
                value_batch = self.actor_critic.evaluate(critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1])
                mu_batch = self.actor_critic.action_mean
                sigma_batch = self.actor_critic.action_std
                entropy_batch = self.actor_critic.entropy

                # KL
                if self.desired_kl != None and self.schedule == 'adaptive':
                    with torch.inference_mode():
                        kl = torch.sum(
                            torch.log(sigma_batch / old_sigma_batch + 1.e-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)

                        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)
                        
                        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()

                # Gradient step
                self.optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
                self.optimizer.step()

                mean_value_loss += value_loss.item()
                mean_surrogate_loss += surrogate_loss.item()

        num_updates = self.num_learning_epochs * self.num_mini_batches
        mean_value_loss /= num_updates
        mean_surrogate_loss /= num_updates
        self.storage.clear()

        return mean_value_loss, mean_surrogate_loss

1. 场景比喻

假设你经营着一家披萨店，店里有一位新来实习的学徒，你想教会他如何做出好吃的披萨。你可以把学徒想象成一个「神经网络」，你就是想要训练他的「大脑」，让他学会每次做披萨时都能做出正确的决策（怎么揉面、放多少酱料、何时烤制等）。

但是你不想一下子就把他灌输所有的知识，你希望通过一种“试错”（互动）方式，让学徒自己从实际操作里学到什么决策才是好的、什么决策是坏的。

在这个过程中，你会：

让学徒去做一段时间披萨（相当于“收集经验”）。
看学徒表现（把“做披萨”的过程数据存下来）。
给学徒打分或奖励（如果披萨好吃，就给高分；不好吃，就给低分）。
在一段时间之后把这些数据拿出来统一进行一次“复盘”，并根据得失进行调整（更新学徒的做法）。

PPO（Proximal Policy Optimization）就是这么一个试错 + 打分 + 调整的过程。代码中，主要有以下角色：

学徒的大脑（actor_critic）：

actor（演员）：决定下一步行动时该做什么（比如：在披萨上放多少酱料）。
critic（评论员）：决定学徒当前动作做得好不好，用价值来衡量（比如：预计这样做出来的披萨有多好吃）。

记事本或账本（RolloutStorage）：

记录学徒每一步的操作、获得的分数、以及一些当时的环境信息、学徒的心态等等。

PPO 这个管理者（PPO类本身）：

负责告诉学徒什么时候去做披萨，什么时候来复盘学习。
负责管理学徒的日常训练计划，比如学习率、奖励计算、经验数据存储等。
训练中的“奖惩原则”或“学习规则”也都在这里。

下面我们就从头到尾讲一下，这段代码是如何在“店里”训练“学徒”的。

2. init：初始化工具

当我们新开店时，会先做一些准备工作。比如，买工具、准备原材料、制定培训计划等等。对应的代码在 init 里完成了这些事：

actor_critic：学徒的大脑。本质是一个神经网络，里边包含了“做决定”和“估算好坏”两部分。
一系列超参数：
- clip_param: PPO 对行动调整幅度的限制。就像你告诉学徒：每次改进可以，但是不要一次性改太激进。
- gamma: 相当于给奖励打折扣的一个因子。比如说，后面才拿到的好处价值会稍微折扣一下。
- lam: 用来计算平滑的优势函数（GAE），帮助学徒知道具体哪一步做得好或不好。
- value_loss_coef, entropy_coef：在总的学习目标里，你要告诉学徒更重视哪部分，比如更重视“预测准确”还是更重视“保持灵活”。
- learning_rate: 学徒学习的速度。学徒如果一次学太多会消化不良，但学太慢又效率太低，需要找个合适的值。
optimizer：训练学徒的方式，这里选择了Adam，就像是你给学徒设定了“连续改进”的学习方式。
storage 和 transition：相当于一本“流水账本”和一个“临时记录纸片”，把学徒每个动作和相应的结果记下来，再定期汇总进行系统学习。

3. init_storage(…)：开本新账

你开了一本新的“流水账本”（RolloutStorage）用来记录学徒每一步做的事、得到的结果。就像开了一个新笔记本，专门记录学徒试做披萨的过程。

4. act(…)：让学徒操作并记录

当顾客来了，学徒开始做披萨了，这里就要让学徒根据所见的环境“观察”（obs、critic_obs）来决定“行动”（动作），并把一些关键信息记到“临时纸片”上。

决定行动（actor）：学徒看到饼皮大小、顾客口味等信息后，决定接下来要怎么做（放多少料、怎么搅拌）。
评价好坏（critic）：学徒内心也会有个估计，“我觉得这一步做出来的披萨会如何？”
记录下来（transition）：把“我的动作”，“我估计的评分”，“当前状态”都先写在一张纸片上。
等真正送到顾客或烤炉那一步后，学徒会拿到实际的结果，这时再把纸片上写的东西正式抄进账本。

5. process_env_step(…)：真正获取反馈并写入账本

就像披萨烤好了，顾客吃完了，你收到了评分（rewards），知道这次表现好不好，还有没有把披萨糊掉（dones 表示这一回合结束没有）。

把奖励记上去：这次顾客给了几星？这就是代码里 self.transition.rewards = rewards.clone()。
记录回合是否结束：如果这块披萨已完成，就要标记 dones，并重置学徒的状态（相当于一单结束重新来）。
处理超时：如果中间某些情况导致“时限到”，还会有一些额外的计算，把学徒自己估计的价值加回奖励（time_outs 的部分）。
正式记账：最后一行 self.storage.add_transitions(self.transition) 就像把那张“临时纸片”正式贴到我们的账本里。
清空纸片：学徒准备下一单做披萨了。

6. compute_returns(…)：结账时算总分

等学徒连着做完了若干个披萨，或者一个回合结束了，就需要来做一个小结：这段时间里学徒做得怎么样？要不要给一些鼓励或指导？

这里的 compute_returns 会把最终的回报（return）算出来，结合过去的估计，算出更精确的“优势函数”（优势：到底哪个操作比平庸的操作高出多少分）。
这些数据会留在账本里，以便稍后的整体学习。

7. update(…)：复盘学习，改进手法

这就是学徒真正集体反思和学习的阶段。
类似于：让学徒把前面做每个披萨的经历都翻出来，好好对比“我当初决定的做法”和“实际得到的结果”之间的差异，然后根据理论（PPO 策略）做改进：

拿出之前所有的记账：这里会把刚才存进账本的所有记录分成若干小段（mini-batch）。
对每个小段数据，学徒先计算现在的想法：
- 学徒看看“如果再来一次，我可能会怎么做（新的动作分布）？”
- 把这个新的想法跟之前实际做的动作对比，并计算变化率（ratio）。
防止学徒“一次改进过猛”：
- PPO 有一个 clip_param（剪切范围），像是你对学徒说：“慢慢改，不要一次把配方改得面目全非。”
- 因此在计算“损失函数”时，会把变化过大的地方剪一下。
价值函数的改进：
- 学徒还要学会更准确地“预估”自己每个动作带来的结果。这会涉及一个专门的价值损失（value loss）。
- 如果开启了“剪切价值”，也会对这部分做类似的限制，防止一口气把学徒内心的评分标准改得太离谱。
加入熵增强探索：
- 让学徒保持对各种可能性的探索，不要太快陷入“老是放同样的料”的习惯。我们给一点额外奖励鼓励多元做法，叫 entropy_coef。
合并成一个总的损失：
- 学徒这次的整体学习目标，就是“在保证做披萨决策改得恰当”的同时，也要让“评价系统更准”。
- 然后再在这个损失上进行一次梯度下降，更新学徒的脑袋（神经网络参数）。
自适应学习率（可选）：
- 如果发现学徒跟之前想法的差异太大（KL 散度过大），就降低学习率；如果差异过小，就提高学习率。就像如果学徒老是变动太大，你劝他先沉住气，别急；如果学徒一直不变动，你希望他大胆一些。

复盘完之后，学徒又变聪明了一些。代码里记录了最后的 mean_value_loss, mean_surrogate_loss 作为“这次复盘学习”的指标，下次你可以继续追踪。

8. 总结

像学徒制作披萨一样：学徒（Actor-Critic 网络）在不同的状态（环境观测）做出不同决策（动作），得到不同反馈（奖励）。
像记录流水账：每一步操作都记录在 storage 中，以便后续反思与学习。
像定期复盘：等学徒完成一段时间练习，就进行一次PPO更新（update），在这过程中分析旧做法与新做法的差异，用“剪切”等安全措施避免学徒一下子完全推翻自己之前的经验。
持续循环：多次“做批量披萨→记录→复盘更新”，学徒的“脑子”慢慢就学会了做出更好的披萨，动作也更合理，分数越来越高。

前半部分在做“采集经验”和“记账”
后半部分在做“学习更新和改进策略”
用生活经验理解，就是一直在“做→记→评→学→再做→再评”的螺旋式改进过程。