[学习笔记]PPO_penalty：trpo的工程化版本

在trpo出现后，强化学习迎来了一波热潮。

但是，trpo确实是有点复杂，而且有一些缺陷，比如：

为了保护策略的稳定性，trpo使用硬约束，只要策略超出接受范围，就直接裁断，这让策略的利用率大大降低

还有，二阶优化+线搜索确实是很难算啊。虽然从原本的二阶矩阵求逆已经简化为求解线性方程组，但是也不免让人浮想联翩：如果能够避免二阶优化问题，那么效果会不会要更好呢？

这就是ppo所解决的问题了。

那么，接下来就来顺一下trpo是怎么变成ppo的吧

首先，trpo的创新是什么，是将策略的相似程度给量化了，也就是KL散度

但是，与此同时，引入KL散度会带来一个问题：这会导致原有策略空间的崩溃

于是，trpo引入了fisher矩阵，来重新构建策略空间的度量

甚至需要引入二阶优化等一系列复杂问题。

虽然这很稳定，但是这一点也不酷。

所谓要用神经网络来打败神经网络

一个朴素的想法就是：全都丢给神经网络和优化器！

毕竟，从这个学科出现以来，我们一直都是这么做的

把特征提取也塞到神经网络里变成CNN，把各种东西全都塞到神经网络里

（其实这样说也不准确，ppo_penalty是把“不能接受”这件事写到奖励函数中作为惩罚）

当计算完了kl散度后，将其作为奖励函数的一部分，再加一个系数β。

如果惩罚过大，则减小β，惩罚过小，则加大β。

这样，复杂的二阶优化也就变成了神经网络自己就能处理的一阶优化问题。

没有什么复杂的数学，没有什么艰难的推导，这很工程。

（插一点私货：我一直认为，某个方法的简洁程度在某种意义上表达了这种方法的上限。trpo固然稳定，但是其过于复杂的推导与应用似乎确实代表了上限不很高。而大模型的爆发也正是从ppo的提出开始的）

之后，ppo-penalty还提出了一个训练策略：数据复用

就是说：在一个epoch中，将一次策略进行好几次训练。

这样就提高了数据的利用效率。

可能有小伙伴会和我有一样的疑惑：那这个与经验回放是不是有什么关系？

是的，这样说吧，经验回放是在一个epoch中使用多个轮次策略，打破了时间的约束，但策略的稳定性无法保证

数据复用是在一个epoch中把这一轮产生的数据进行多次复用，就像看错题那样。

这个提升也是建立在一阶优化问题简单，而且约束要更宽松的前提下的。

所以，ppo_penalty的思想是很巧妙的，贴一下代码，下一节是更巧妙地ppo_clip，也是ppo名声大噪的部分。

# PPO-Penalty: 通过KL散度惩罚实现策略更新，动态调整惩罚系数以近似满足KL约束
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gymnasium as gym
import time
from torch.distributions import Categorical

# Hyper Parameters
BATCH_SIZE = 4000
LR = 0.01  # 价值网络学习率
LR_POLICY = 0.0003  # 策略网络学习率（PPO通常使用较小的学习率）
GAMMA = 0.99
LAMBDA = 0.95
TARGET_KL = 0.01  # 目标KL散度（对应TRPO的MAX_KL）
KL_PENALTY_COEF = 1.0  # KL惩罚系数初始值
ADAPTIVE_KL_TARGET = 1.5  # 自适应KL调整的倍数阈值
EPOCHS = 10  # 每次更新时对同一批数据迭代次数
EPS = 1e-8

env = gym.make('CartPole-v1', render_mode='human')
env = env.unwrapped
N_ACTIONS = env.action_space.n
N_STATES = env.observation_space.shape[0]
ENV_A_SHAPE = 0 if isinstance(
    env.action_space.sample(), int) else env.action_space.sample().shape


class PolicyNet(nn.Module):
    def __init__(self):
        super(PolicyNet, self).__init__()
        self.fc1 = nn.Linear(N_STATES, 50)
        self.fc1.weight.data.normal_(0, 0.1)
        self.out = nn.Linear(50, N_ACTIONS)
        self.out.weight.data.normal_(0, 0.1)

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        action_logits = self.out(x)
        return F.softmax(action_logits, dim=-1)


class ValueNet(nn.Module):
    def __init__(self):
        super(ValueNet, self).__init__()
        self.fc1 = nn.Linear(N_STATES, 50)
        self.fc1.weight.data.normal_(0, 0.1)
        self.out = nn.Linear(50, 1)
        self.out.weight.data.normal_(0, 0.1)

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        state_value = self.out(x)
        return state_value


class PPO_Penalty:
    def __init__(self):
        self.policy_net = PolicyNet()
        self.value_net = ValueNet()
        self.policy_optimizer = torch.optim.Adam(
            self.policy_net.parameters(), lr=LR_POLICY)
        self.value_optimizer = torch.optim.Adam(
            self.value_net.parameters(), lr=LR)
        self.kl_penalty_coef = KL_PENALTY_COEF  # 动态调整的KL惩罚系数

    def choose_action(self, x):  # 与AC一致，得出概率后采样
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        probs = self.policy_net(x)
        m = Categorical(probs)
        action = m.sample()
        return action.item(), m.log_prob(action)

    def compute_advantages(self, rewards, values, dones, next_value):
        advantages = np.zeros(len(rewards))
        last_advantage = 0

        for t in reversed(range(len(rewards))):
            if t == len(rewards) - 1:
                next_val = next_value
            else:
                next_val = values[t + 1] * (1 - dones[t])

            delta = rewards[t] + GAMMA * next_val - values[t]
            advantages[t] = delta + GAMMA * LAMBDA * \
                (1 - dones[t]) * last_advantage
            last_advantage = advantages[t]

        return advantages

    def update_value_net(self, states, returns):
        states = torch.FloatTensor(states)
        returns = torch.FloatTensor(returns).unsqueeze(1)

        for _ in range(10):
            values = self.value_net(states)
            value_loss = F.mse_loss(values, returns)

            self.value_optimizer.zero_grad()
            value_loss.backward()
            self.value_optimizer.step()

    def update_policy(self, states, actions, log_probs_old, advantages):
        states = torch.FloatTensor(states)
        actions = torch.LongTensor(actions)
        log_probs_old = torch.FloatTensor(log_probs_old)
        advantages = torch.FloatTensor(advantages)

        # 标准化优势
        advantages = (advantages - advantages.mean()) / \
            (advantages.std() + EPS)

        # 获取旧策略的概率分布（用于KL计算）
        with torch.no_grad():
            old_probs = self.policy_net(states)
            old_log_probs = torch.log(old_probs.gather(
                1, actions.unsqueeze(1))).squeeze()
            old_probs_clone = old_probs.clone()

        total_policy_loss = 0.0
        total_kl = 0.0

        # 对同一批数据进行多轮优化
        for _ in range(EPOCHS):
            # 计算新策略
            probs = self.policy_net(states)
            log_probs = torch.log(probs.gather(
                1, actions.unsqueeze(1))).squeeze()

            # 计算比率 r(θ)
            ratio = torch.exp(log_probs - log_probs_old)

            # 计算KL散度 (平均每个样本)
            kl = (old_probs_clone * (torch.log(old_probs_clone + EPS) -
                  torch.log(probs + EPS))).sum(dim=1).mean()

            # PPO-Penalty 损失函数: L = -E[ r(θ) * A ] + β * KL
            policy_loss = -(ratio * advantages).mean() + \
                self.kl_penalty_coef * kl

            # 更新策略网络
            self.policy_optimizer.zero_grad()
            policy_loss.backward()
            torch.nn.utils.clip_grad_norm_(
                self.policy_net.parameters(), max_norm=0.5)  # 梯度裁剪
            self.policy_optimizer.step()

            total_policy_loss += policy_loss.item()
            total_kl += kl.item()

        # 计算平均KL并动态调整惩罚系数
        avg_kl = total_kl / EPOCHS
        if avg_kl < TARGET_KL / ADAPTIVE_KL_TARGET:
            self.kl_penalty_coef *= 0.5  # KL太小，减小惩罚
        elif avg_kl > TARGET_KL * ADAPTIVE_KL_TARGET:
            self.kl_penalty_coef *= 2.0  # KL太大，增大惩罚

        return avg_kl


ppo = PPO_Penalty()

print('\nTraining with PPO-Penalty...')
time.sleep(2)

for i_episode in range(400):
    states, actions, log_probs, rewards, dones, values = [], [], [], [], [], []

    s, info = env.reset()
    ep_r = 0

    while True:
        env.render()
        a, log_prob = ppo.choose_action(s)

        s_, r, terminated, truncated, info = env.step(a)
        done = terminated or truncated

        # modify the reward (same as before)
        x, x_dot, theta, theta_dot = s_
        r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8
        r2 = (env.theta_threshold_radians - abs(theta)) / \
            env.theta_threshold_radians - 0.5
        r = r1 + r2

        # store transition
        states.append(s)
        actions.append(a)
        log_probs.append(log_prob.item())
        rewards.append(r)
        dones.append(done)

        # value estimate
        with torch.no_grad():
            state_tensor = torch.FloatTensor(s).unsqueeze(0)
            value = ppo.value_net(state_tensor).item()
            values.append(value)

        ep_r += r

        if done:
            # 计算最后一个状态的value
            with torch.no_grad():
                if terminated:
                    next_value = 0.0
                else:
                    next_value = ppo.value_net(
                        torch.FloatTensor(s_).unsqueeze(0)).item()

            # 计算优势函数
            advantages = ppo.compute_advantages(
                rewards, values, dones, next_value)

            # 计算回报
            returns = advantages + values

            # 更新价值网络
            ppo.update_value_net(states, returns)

            # 更新策略网络
            avg_kl = ppo.update_policy(states, actions, log_probs, advantages)

            print('Ep: ', i_episode, '| Ep_r: ', round(ep_r, 2), '| Avg_KL: ', round(
                avg_kl, 6), '| Beta: ', round(ppo.kl_penalty_coef, 4))
            break

        s = s_

env.close()