近端策略优化算法PPO的核心概念和PyTorch实现详解

 !pip install swig gymnasium torch tensordict pyvirtualdisplay  
 !pip install "gymnasium[box2d]"

 import random  
 import torch  
 import numpy as np  
 seed = 777  
 random.seed(seed)  
 np.random.seed(seed)  
 torch.manual_seed(seed)  
 torch.backends.cudnn.deterministic = True

 import gymnasium as gym  
from gymnasium.wrappers import RecordVideo  

# 创建将运行10个模拟的向量化环境
env_name = "LunarLander-v3"  
num_envs = 10  
envs = gym.make_vec(env_name, num_envs=num_envs, vectorization_mode="sync")  

# 创建我们的评估录制环境，用于当我们
# 想要测试我们的智能体在训练的各个阶段表现如何时使用
env = gym.make(env_name, render_mode="rgb_array")  
trigger = lambda t: True  
recording_output_directory = "./checkpoint_videos"  
 recording_env = RecordVideo(env, recording_output_directory, episode_trigger=trigger)

 from tensordict import TensorDict  
import torch  

# 关于这个类的快速说明
# 它假设训练环境在终端状态时自动重置，例如向量化环境
# 并且它假设评估环境不会自动重置
# 我还将num_steps_per_rollout固定为收集器初始化时的值，但这可以在
# get rollout函数中参数化
class PPORolloutCollector:  
    def __init__(self, agent, envs, num_steps_per_rollout, device, eval_env=None):  
        self.agent = agent  
        self.envs = envs  
        self.num_envs = envs.num_envs  
        self.num_steps_per_rollout = num_steps_per_rollout  
        self.device = device  
        self.eval_env = eval_env  

        self.obs_shape = envs.single_observation_space.shape  
        self.action_shape = envs.single_action_space.shape  
        self.initial_buffer_shape = (num_steps_per_rollout, envs.num_envs)  

        # 继续重置环境并存储观察
        # 并将next_done设置为false(我们假设环境不能以终端状态开始)
        obs, _ = envs.reset()  
        self.next_obs = torch.Tensor(obs).to(device)  
        self.next_done = torch.zeros(num_envs).to(device)  

    # 创建一个空缓冲区，将保存观察、来自智能体的动作、该动作的对数概率
    # 当前观察的评论家估计、从环境中获得的采取动作的实际奖励，以及
    # 动作是否导致终端状态
    # 我们故意选择不为每个观察记录"下一状态"，这基本上会使
    # 缓冲区的大小翻倍，由于我们按顺序收集观察并且在这里不打乱
    # 任何下游操作都可以只使用数组中的下一个值作为下一状态
    def _create_buffer(self):  
        return TensorDict({  
            "obs": torch.zeros(self.initial_buffer_shape + self.obs_shape).to(self.device),  
            "actions": torch.zeros(self.initial_buffer_shape + self.action_shape).to(self.device),  
            "log_probs": torch.zeros(self.initial_buffer_shape).to(self.device),  
            "rewards": torch.zeros(self.initial_buffer_shape).to(self.device),  
            "dones": torch.zeros(self.initial_buffer_shape).to(self.device),  
            "critic_values": torch.zeros(self.initial_buffer_shape).to(self.device),  
        })  

    # 将收集num_steps_per_rollout个观察对我们的训练环境的函数
    def get_next_rollout(self):  
        buffer = self._create_buffer()  

        # 获取最后记录的观察以及该观察是否为终端
        next_obs = self.next_obs  
        next_done = self.next_done  

        # 收集rollout
        for t in range(self.num_steps_per_rollout):  
            # 记录当前观察和终端状态
            buffer["obs"][t] = next_obs  
            buffer["dones"][t] = next_done  

            # 查询智能体下一个动作、该动作的对数概率和评论家估计
            with torch.no_grad():  
                action, log_prob, entropy = self.agent.get_actor_values(next_obs)  
                critic_value = self.agent.get_critic_value(next_obs)  

            # 记录值
            buffer["actions"][t] = action  
            buffer["log_probs"][t] = log_prob  
            buffer["critic_values"][t] = critic_value.flatten()  

            # 执行动作
            next_obs, reward, terminations, truncations, infos = envs.step(action.cpu().numpy())  

            # 形状化并存储奖励
            reward = torch.tensor(reward).to(self.device).view(-1)  
            buffer["rewards"][t] = reward  

            # 一些环境会终止(意味着智能体处于最终状态)，
            # 其他会截断(例如达到时间限制但不在终端状态)
            # 这些是重要的区别，但对我们的目的来说意味着同样的事情，模拟结束了
            # 所以如果任一为真且模拟重置，我们将next done设置为true
            next_done = np.logical_or(terminations, truncations)  
            # 为下一轮存储下一个obs和next done
            next_obs, next_done = torch.Tensor(next_obs).to(self.device), torch.Tensor(next_done).to(self.device)  

        # 在缓冲区中存储下一个obs和next done
        # 这是为了在稍后计算优势时处理边缘情况
        buffer['next_obs'] = next_obs  
        buffer['next_done'] = next_done  
        # 我们还需要这个下一状态的评论家估计，我们将使用它来引导最终状态的奖励
        # 当我们计算gae时
        with torch.no_grad():  
            buffer['next_value'] = self.agent.get_critic_value(next_obs).reshape(1, -1)  
        self.next_obs = next_obs  
        self.next_done = next_done  
        return buffer  

      
    # 用于在我们的环境上评估智能体，将运行整个模拟直到终止
    # 与上面非常相似，唯一的区别是我们将手动检查环境是否终止
    # 然后结束循环
    # 将返回一个普通的python字典，包含奖励、熵、奖励平均值、我们智能体的平均熵，以及每次运行的总奖励
    def run_eval_rollout(self, num_episodes: int = 5):  
        assert self.eval_env is not None, "No eval_env provided."  

        rewards_per_timestep = []  
        entropies_per_timestep = []  
        final_rewards = []  
        total_entropies = []  

        for _ in range(num_episodes):  
            obs, _ = self.eval_env.reset()  
            obs = torch.tensor(obs, device=self.device).unsqueeze(0)  
            done = False  

            episode_rewards = []  
            episode_entropies = []  

            while not done:  
                with torch.no_grad():  
                    action, _, entropy = self.agent.get_actor_values(obs)  
                    action = action.squeeze()  
                obs_np, reward, term, trunc, _ = self.eval_env.step(action.cpu().numpy())  
                done = term or trunc  

                obs = torch.tensor(obs_np, device=self.device).unsqueeze(0)  
                episode_rewards.append(float(reward))  
                episode_entropies.append(entropy.item())  

            rewards_per_timestep.append(episode_rewards)  
            entropies_per_timestep.append(episode_entropies)  
            final_rewards.append(sum(episode_rewards))  
            total_entropies.append(sum(episode_entropies) / len(episode_entropies))  

        return {  
            "rewards_per_timestep": rewards_per_timestep,  
            "average_reward_per_run": sum(final_rewards) / len(final_rewards),  
            "average_entropy_per_run": sum(total_entropies) / len(total_entropies),  
            "entropies_per_timestep": entropies_per_timestep,  
            "final_rewards": final_rewards,  
         }

 import torch.nn as nn  
from torch.distributions.categorical import Categorical  

# 一个初始化辅助函数。在强化学习问题中，正交初始化
# 网络的权重可能会有帮助，这意味着每层的输出尽可能
# 不相关。这可以通过帮助梯度更好地流动和避免可能
# 从朴素初始化中产生的相关特征问题来改善训练稳定性和效率。
# 可以把这想象成确保网络不会以交叉的线路开始。
# 这个函数取自CleanRL的PPO实现。
def layer_init(layer, std=np.sqrt(2), bias_const=0.0):  
    torch.nn.init.orthogonal_(layer.weight, std)  
    torch.nn.init.constant_(layer.bias, bias_const)  
    return layer  

      
class Agent(nn.Module):  
    def __init__(self, envs):  
        super().__init__()  
        # 我们的输入数组有多大？
        # 注意如果我们做的是图像观察之类的
        # 我们需要重新设计这个，但lunar lander只提供
        # 代表世界状态的数字数组
        input_shape = np.array(envs.single_observation_space.shape).prod()  

        # 我们的智能体将有多少个动作？
        action_shape = envs.single_action_space.n  

        # 创建一个3层评论家，它将预测
        # 我们的智能体在给定观察下预期收到的总奖励
        self.critic = nn.Sequential(  
            layer_init(nn.Linear(input_shape, 64)),  
            nn.Tanh(),  
            layer_init(nn.Linear(64, 64)),  
            nn.Tanh(),  
            layer_init(nn.Linear(64,1), std=1.0)  
        )  

        # 创建一个3层演员，它将输出一个概率分布
        # 表示在给定观察下最好采取哪个动作的概率
        self.actor = nn.Sequential(  
            layer_init(nn.Linear(input_shape, 64)),  
            nn.Tanh(),  
            layer_init(nn.Linear(64, 64)),  
            nn.Tanh(),  
            layer_init(nn.Linear(64, action_shape), std=1.0)  
        )  

      
    def save(self, path):  
        torch.save(self.state_dict(), path)  

    def load(self, path):  
        self.load_state_dict(torch.load(path))  

    def get_critic_value(self, x):  
        return self.critic(x)  

    def get_actor_values(self, x, action=None):  
        logits = self.actor(x)  
        probs = Categorical(logits=logits)  
        if action is None:  
            action = probs.sample()  
         return action, probs.log_prob(action), probs.entropy()