基于“蘑菇书”的强化学习知识点（十五）：第五章的代码：PPO.ipynb及其涉及的其他代码的更新以及注解（gym版本 ＞= 0.26）

摘要

本系列知识点讲解基于蘑菇书EasyRL中的内容进行详细的疑难点分析！具体内容请阅读蘑菇书EasyRL！

---

对应蘑菇书附书代码——PPO.ipynb

---

# -*- coding: utf-8 -*-


import torch.nn as nn
import torch.nn.functional as F # 提供激活函数（如 F.relu、F.softmax）等函数，这些函数是无状态的。


class ActorSoftmax(nn.Module):
    """
    演员网络用于输出在当前状态下采取各个动作的概率分布，常用于策略梯度方法中。
    该类继承自 nn.Module，确保网络参数自动注册，支持自动求导。
    """
    def __init__(self, input_dim, output_dim, hidden_dim=256):
        """
        解释：
        - ActorSoftmax 继承自 nn.Module，因此在构造函数中调用 super(…) 初始化父类。
        参数：
        - input_dim：输入特征数（状态的维度）；
        - output_dim：输出的维度，通常对应动作数量；
        - hidden_dim：隐藏层神经元个数，默认设置为 256。
        """
        super(ActorSoftmax, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)
    def forward(self,x):
        """x 是输入状态，形状通常为 (batch_size, input_dim)。"""
        '''
        self.fc1(x) 计算线性变换，将输入从 (batch_size, input_dim) 映射为 (batch_size, hidden_dim)。
        F.relu 激活函数将所有负值置零，增加非线性。 
        举例：
        - 假设 x 为一个 1×4 张量 [0.5, -0.2, 0.1, 0.3]，
          经过 fc1 后得到 1×256 张量，其中某一分量可能为 0.8，经过 ReLU 则保持为 0.8；负值则变为 0。
        '''
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        '''
        F.softmax(…, dim=1) 对每个样本的输出向量进行 softmax，使得输出每一行的和为 1，
        表示各动作的概率分布。 
        举例：
        - 如果 fc3 输出为 [0.2, 1.8, 0.5]，
          则 softmax 后可能得到概率向量约为 [exp(0.2)/Z, exp(1.8)/Z, exp(0.5)/Z]，
          计算后例如为 [0.17, 0.74, 0.09]（Z 为归一化常数）。
        '''
        probs = F.softmax(self.fc3(x),dim=1)
        '''probs 是一个形状为 (batch_size, output_dim) 的概率分布张量。'''
        return probs
    
    
class Critic(nn.Module):
    """
    评论家网络用于估计给定状态的价值函数，即返回一个标量值。
    它同样继承自 nn.Module，结构与演员类似，但输出层只有一个节点。
    """
    def __init__(self,input_dim,output_dim,hidden_dim=256):
        """
        作用：
        - 初始化评论家网络，要求输出维度 output_dim 必须为 1，因为评论家只估计状态价值。
        参数：
        - input_dim：状态特征数；
        - output_dim：输出维度，必须为 1；
        - hidden_dim：隐藏层神经元个数，默认 256。
        """
        super(Critic,self).__init__()
        assert output_dim == 1, "critic must output a single value!"
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)
    def forward(self,x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        value = self.fc3(x)
        return value
    
    
import random
from collections import deque


class ReplayBufferQue:
    '''DQN的经验回放池，每次采样batch_size个样本'''
    """
    作用：
    - 定义一个名为 ReplayBufferQue 的类，用于存储和采样经验数据。
    - 注释中说明其主要用途是 DQN 中的经验回放池，每次从中采样一定数量（batch_size）的样本。
    """
    def __init__(self, capacity: int) -> None:
        """参数 capacity 表示回放池的容量，也就是最多存储的转移数量。"""
        '''self.capacity = capacity 
        - 将传入的容量保存为类的一个属性。'''
        self.capacity = capacity
        '''self.buffer = deque(maxlen=self.capacity)
        - 使用 collections.deque 构造一个双端队列，并设置最大长度为 capacity。
        - 当缓冲区达到最大容量时，新加入的样本会自动淘汰最早存入的样本。
        例如：
        - 若 capacity 为 10000，则最多保留 10000 个经验数据，超过后最旧的数据将被丢弃。'''
        self.buffer = deque(maxlen=self.capacity)
    def push(self,transitions):
        '''_summary_
        Args:
            trainsitions (tuple): _description_
        '''
        """
        作用：
        - 将一条经验（通常为一个 tuple，例如 (state, action, reward, next_state, done)）添加到缓冲区中。
        举例：
        - 假设 transitions = (s, a, r, s', done)，
          调用 push(transitions) 后，该 tuple 被追加到 self.buffer 的末尾。
        """
        self.buffer.append(transitions)
    def sample(self, batch_size: int, sequential: bool = False):
        """
        作用：
        - 从缓冲区中采样一批经验数据，并以“解压”形式返回。
        """
        '''
        首先检查请求的 batch_size 是否超过当前存储的样本数，
        如果超过，则将 batch_size 调整为当前缓冲区大小。
        '''
        if batch_size > len(self.buffer):
            batch_size = len(self.buffer)
        if sequential: # sequential sampling
            """顺序采用"""
            '''
            随机生成一个起始索引 rand，使得从 rand 到 rand+batch_size 的区间在缓冲区内；
            使用列表推导式取出这一连续区间的样本，最后用 zip(*batch) 将样本“解压”，
            即将各个元素分别组合成一个元组（例如所有状态组成一个 tuple、所有动作组成一个 tuple……）。
            '''
            '''假设缓冲区有 5 个样本，batch_size=3，
            则 rand 的取值范围为 0 到 2（例如随机选到 rand=1）。'''
            rand = random.randint(0, len(self.buffer) - batch_size)
            '''例子：如果缓冲区存储的样本依次编号为 0,1,2,3,4，且 rand=1，batch_size=3，
            则抽取样本 1、2、3。'''
            batch = [self.buffer[i] for i in range(rand, rand + batch_size)]
            return zip(*batch)
        else:
            """随机采样"""
            '''
            使用 random.sample 从 self.buffer 中随机采样 batch_size 个不同的样本；
            同样用 zip(*batch) 解压返回。'''
            '''注意：random.sample 返回的是不重复的样本，顺序可能不同于存储顺序。'''
            batch = random.sample(self.buffer, batch_size)
            return zip(*batch)
    def clear(self):
        """
        作用：
        - 清空缓冲区中的所有经验数据。
        """
        self.buffer.clear()
    def __len__(self):
        """
        作用：
        - 定义 len(ReplayBufferQue) 的行为，返回缓冲区中当前存储的样本数。
        """
        return len(self.buffer)


class PGReplay(ReplayBufferQue):
    '''PG的经验回放池，每次采样所有样本，因此只需要继承ReplayBufferQue，重写sample方法即可
    '''
    """
    继承说明：
    - PGReplay 类继承自 ReplayBufferQue，但用于策略梯度（Policy Gradient）的经验回放。
    - 在策略梯度方法中，通常在一次回合结束后，我们希望使用整个回合的经验来更新策略，
      而不是随机抽取部分数据。
    """
    def __init__(self):
        """
        作用：
        - 重写 init 方法，不设置最大容量（即没有容量限制），直接创建一个空的 deque 存储经验。
        - 这与 ReplayBufferQue 不同，后者设置了最大长度（capacity）。
        """
        self.buffer = deque()
    def sample(self):
        ''' sample all the transitions
        '''
        """
        作用：
        - 重写 sample 方法，使得每次采样时返回缓冲区中所有的经验数据，而不是随机抽取部分样本。
        """
        batch = list(self.buffer)
        return zip(*batch)    
    
    
import torch
from torch.distributions import Categorical
import numpy as np


class Agent:
    """
    定义 Agent 类，用于构造策略梯度智能体。
    该类不是继承自其他特定父类，只是一个普通的 Python 类，但它内部使用了 PyTorch 模块构建神经网络和优化器。
    """
    def __init__(self,cfg) -> None:
        '''作用：
        保存折扣因子 γ，用于后续计算累计折扣回报。例如，若 γ=0.9，则后续奖励乘以 0.9 的幂。'''
        self.gamma = cfg.gamma
        '''作用：
        根据配置 cfg['device'] 创建一个 torch.device 对象，用于指定运行设备（如 "cpu" 或 "cuda"）。'''
        self.device = torch.device(cfg.device) 
        '''
        作用：
        - 创建演员网络（策略网络），其输入维度为状态数（cfg.n_states），输出维度为动作数（cfg.n_actions），
          隐藏层维度由 cfg.actor_hidden_dim 指定。
        说明：
        - ActorSoftmax 是一个全连接网络，输出经过 softmax 后得到每个动作的概率分布。
          最后用 .to(self.device) 将网络移动到指定设备。
        '''
        self.actor = ActorSoftmax(cfg.n_states,cfg.n_actions, hidden_dim = cfg.actor_hidden_dim).to(self.device)
        '''
        作用：
        - 创建评论家网络（价值网络），输入维度为状态数（cfg.n_states），
          输出维度为 1（即单个标量），隐藏层维度由 cfg.critic_hidden_dim 指定。
        说明：
        - Critic 网络用于估计当前状态的价值，即 V(s)。
        '''
        self.critic = Critic(cfg.n_states,1,hidden_dim=cfg.critic_hidden_dim).to(self.device)
        '''作用：
        使用 Adam 优化器分别为演员网络和评论家网络设置优化器，
        学习率分别由 cfg.actor_lr 和 cfg.critic_lr 指定。
        例如，若 actor_lr=0.001，则演员网络的参数每次更新步长由此控制。'''
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=cfg.actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=cfg.critic_lr)
        '''作用：
        将传入的经验存储器 memory 保存到智能体中，用于后续更新时采样整个回合的经验。'''
        self.memory = PGReplay()
        '''
        k_epochs：在每次更新中对收集到的经验进行多次（K 次）梯度下降更新。
        eps_clip：PPO 中用来限制新旧策略比值变化的剪切参数。
        entropy_coef：熵正则项系数，用于鼓励策略的探索，即使得策略分布保持一定的随机性。
        '''
        self.k_epochs = cfg.k_epochs # update policy for K epochs
        self.eps_clip = cfg.eps_clip # clip parameter for PPO
        self.entropy_coef = cfg.entropy_coef # entropy coefficient
        '''
        sample_count：用于统计智能体采样动作的总次数，当达到一定频率（update_freq）时触发策略更新。
        update_freq：指定每隔多少步更新一次策略。
        '''
        self.sample_count = 0
        self.update_freq = cfg.update_freq

    def sample_action(self,state):
        """动作采样"""
        '''作用：
        每次采样动作时先将 sample_count 加 1，用于判断是否达到更新频率。'''
        self.sample_count += 1
        '''作用：
        - 将输入状态（通常为 NumPy 数组）转换为 torch 张量，并将其放入指定设备。
        -- unsqueeze(dim=0) 为状态添加一个 batch 维度，保证输入网络时形状为 (1, n_states)。'''\

        if isinstance(state, tuple):
            state = state[0]
        else:
            state = state

        state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
        '''作用：
        - 将状态传入演员网络，得到动作概率分布 probs。
        -- 假设演员网络输出为一个张量，形状 (1, n_actions)，例如 [0.7, 0.2, 0.1] 表示 3 个动作的概率。'''
        probs = self.actor(state)
        '''作用：
        - 根据输出概率构造一个 Categorical 分布对象，用于采样离散动作。
        -- Categorical 分布接受一个概率向量，并支持 sample() 方法。'''
        dist = Categorical(probs)
        '''作用：
        - 根据分布采样一个动作。
        -- 例如，如果概率向量为 [0.7, 0.2, 0.1]，那么大概率采样到动作 0。'''
        action = dist.sample()
        '''作用：
        - 计算所采样动作的对数概率 log(pi(a|s))，并 detach() 掉梯度，
          保存到 self.log_probs 供后续更新使用。
        -- 例如，如果动作 0 的概率为 0.7，则 log_prob 约为 log(0.7) ≈ -0.357。
        '''
        self.log_probs = dist.log_prob(action).detach()
        '''作用：
        - 将采样得到的动作转换为标量整数返回。
        -- detach() 移除梯度计算，
           cpu() 将数据放到 CPU 上，
           numpy() 转为 NumPy 数组，
           item() 提取标量值。'''
        return action.detach().cpu().numpy().item()
    @torch.no_grad() # 使用 @torch.no_grad() 装饰器，表示在该方法中不需要计算梯度，适用于测试或评估阶段。
    def predict_action(self,state):
        '''作用：
        - 将输入状态（通常为 NumPy 数组）转换为 torch 张量，并将其放入指定设备。
        -- unsqueeze(dim=0) 为状态添加一个 batch 维度，保证输入网络时形状为 (1, n_states)。'''
        
        if isinstance(state, tuple):
            state = state[0]
        else:
            state = state
        
        state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
        '''作用：
        - 将状态传入演员网络，得到动作概率分布 probs。
        -- 假设演员网络输出为一个张量，形状 (1, n_actions)，例如 [0.7, 0.2, 0.1] 表示 3 个动作的概率。'''
        probs = self.actor(state)
        '''作用：
        - 根据输出概率构造一个 Categorical 分布对象，用于采样离散动作。
        -- Categorical 分布接受一个概率向量，并支持 sample() 方法。'''
        dist = Categorical(probs)
        '''作用：
        - 根据分布采样一个动作。
        -- 例如，如果概率向量为 [0.7, 0.2, 0.1]，那么大概率采样到动作 0。'''
        action = dist.sample()
        '''作用：
        - 将采样得到的动作转换为标量整数返回。
        -- detach() 移除梯度计算，
           cpu() 将数据放到 CPU 上，
           numpy() 转为 NumPy 数组，
           item() 提取标量值。'''
        return action.detach().cpu().numpy().item()
    def update(self):
        # update policy every n steps
        """
        演员－评论家智能体中策略更新的核心部分，
        它结合了蒙特卡罗回报估计和 PPO 的剪切目标（clipped surrogate objective）。
        """
        '''作用：
        - 只有当采样步数（self.sample_count）是更新频率（self.update_freq）的整数倍时，才进行策略更新。
        数值例子：
        - 如果 update_freq 设置为 10，当前 sample_count 为 15，
          则 15 % 10 = 5 ≠ 0，因此直接返回，不执行更新；
          只有当 sample_count 为 10、20、30…时才会更新。
        '''
        if self.sample_count % self.update_freq != 0:
            return
        # print("update policy")
        '''作用：
        - 从内存中采样整个回合的经验数据，分别得到状态、动作、旧对数概率、奖励和 done 标志。
        -- 将状态和动作数据转换为 NumPy 数组，再转换为 torch 张量，以便后续计算。
        -- 旧的对数概率也转换为张量。        
        从经验存储器中采样一整个回合的数据。假设我们采样得到：
        - old_states = [s1, s2, s3]
        - old_actions = [1, 0, 1]
        - old_log_probs = [-0.35, -0.80, -0.40]
        - old_rewards = [1, 0, 2]
        - old_dones = [False, False, True]
        '''
        old_states, old_actions, old_log_probs, old_rewards, old_dones = self.memory.sample()
        
        # convert to tensor
        # 过滤掉字典部分
        old_states = [s[0] if isinstance(s, tuple) else s for s in old_states]
        # 转换为 numpy 数组
        old_states = np.array(old_states, dtype=np.float32)
        old_states = torch.tensor(np.array(old_states), device=self.device, dtype=torch.float32)
        old_actions = torch.tensor(np.array(old_actions), device=self.device, dtype=torch.float32)
        old_log_probs = torch.tensor(old_log_probs, device=self.device, dtype=torch.float32)
        # monte carlo estimate of state rewards
        '''蒙特卡罗估计累计折扣回报
        作用：
        - 计算从每个时刻到回合结束的累计折扣回报 G。
        - G_t={r_{t}}+{γr _{t+1}}+{γ^2r_{t+2}}+ ⋯
        - 如果 dt=True, Gt=rt, 否则 G_t=r_t+γG_{t+1}
        具体过程：
        - 遍历奖励和 done 标志的倒序序列。
        - 如果 done 为 True（表示回合终止），将 discounted_sum 重置为 0；
          否则累计折扣回报为 reward + γ × discounted_sum。
        - 使用 returns.insert(0, discounted_sum) 在列表头部插入计算结果，
          从而保证最终返回的 returns 顺序与原始序列一致。
        举例：
        - 假设 old_rewards = [1, 0, 2]，old_dones = [False, False, True]，γ=0.9。
        - 倒序遍历：
        -- 最后一步：reward=2, done=True，设 discounted_sum=0 → discounted_sum=2+0=2，插入返回序列；
        -- 第二步：reward=0, done=False，discounted_sum=2*0.9+0=1.8，插入到最前面；
        -- 第一步：reward=1, done=False，discounted_sum=1+1.8*0.9=1+1.62=2.62；
        -- 返回的 returns = [2.62, 1.8, 2].
        '''
        returns = []
        discounted_sum = 0
        for reward, done in zip(reversed(old_rewards), reversed(old_dones)):
            if done:
                discounted_sum = 0
            discounted_sum = reward + (self.gamma * discounted_sum)
            returns.insert(0, discounted_sum)
        # Normalizing the rewards:
        '''作用：
        - 将计算好的 returns 转换为张量，并归一化，使其均值为 0，标准差为 1。
        - 这样做可以使梯度更新更加稳定。'''
        returns = torch.tensor(returns, device=self.device, dtype=torch.float32)
        returns = (returns - returns.mean()) / (returns.std() + 1e-5) # 1e-5 to avoid division by zero
        for _ in range(self.k_epochs):
            """重复进行策略更新多次（例如 k_epochs = 4），以充分利用同一批数据。"""
            # compute advantage
            '''
            作用：
            - 使用评论家网络对所有旧状态进行前向传播，得到状态价值估计V(s_t) values。
            数值例子：
            - 假设 critic(old_states) 返回 tensor([0.5, 0.0, -0.5])（形状 (3,)），
              表示对三个状态的价值估计。
            '''
            values = self.critic(old_states) # detach to avoid backprop through the critic
            '''
            作用：
            - 计算优势值 advantage = normalized returns - critic 估计,
              即归一化的累计回报与评论家估计之间的差异：At=~Gt−V(st)
              （detach() 防止梯度传到 critic 网络）。
            数值例子：
            - 之前 normalized returns = [1.37, -0.97, -0.40]，values.detach() = [0.5, 0.0, -0.5]
              则 advantage = [1.37 - 0.5, -0.97 - 0.0, -0.40 - (-0.5)] = [0.87, -0.97, 0.10].
            '''
            advantage = returns - values.detach()
            # get action probabilities
            '''
            作用：
            - 将旧状态批量传入演员网络，得到当前策略下的动作概率分布 probs。πθ(a|s)
            - 构造 Categorical 分布对象 dist，这个分布用于计算 log 概率和熵等信息。
            '''
            probs = self.actor(old_states)
            dist = Categorical(probs)
            # get new action probabilities
            '''
            作用：
            - 计算当前策略下，对于之前采样的旧动作（old_actions）的对数概率，得到 new_probs。
            -- 注意：这里 new_probs 实际上是 log(pi(a|s)) 的新计算值。
            数值例子：
            - 假设演员网络对于一个状态输出概率 [0.7, 0.2, 0.1]，
              而旧动作为 1，则 new_probs = log(0.2) ≈ -1.6094。
            '''
            new_probs = dist.log_prob(old_actions)
            # compute ratio (pi_theta / pi_theta__old):
            '''
            这里通过对数概率差再取指数实现。
            - 这个 ratio 是 PPO 算法中重要的部分，用于衡量当前策略与旧策略的变化。
            '''
            ratio = torch.exp(new_probs - old_log_probs) # old_log_probs must be detached
            # compute surrogate loss
            '''
            作用：
            - surr1 表示原始的策略梯度目标：ratio × advantage。
            - surr2 使用 torch.clamp 限制 ratio 的范围在 [1 - eps_clip, 1 + eps_clip] 内，
              再乘以 advantage。
            - 最终目标使用两者的最小值，以防止策略更新幅度过大，从而实现 PPO 中的剪切策略目标。
            '''
            surr1 = ratio * advantage
            surr2 = torch.clamp(ratio, 1 - self.eps_clip, 1 + self.eps_clip) * advantage
            # compute actor loss
            '''
            作用：
            - 计算演员网络的损失。
            - 首先取 surr1 和 surr2 的逐元素最小值，再取平均，前面加负号，使得目标变为最大化该值。
            - 同时加入一个熵项（dist.entropy() 的均值），乘以一个系数 entropy_coef，
              用于鼓励策略的多样性和探索。
            - 例如，如果最小的 surr 值为 0.2，则 actor_loss = -0.2 + entropy_coef * entropy；
              熵项越大，loss 越低，鼓励策略不至于过早收敛为确定性策略。
            '''
            actor_loss = -torch.min(surr1, surr2).mean() + self.entropy_coef * dist.entropy().mean()
            # compute critic loss
            '''
            作用：
            - 计算评论家网络的均方误差损失（MSE loss），
              即 (returns - values)² 的平均值，用于使评论家估计更准确。
            '''
            critic_loss = (returns - values).pow(2).mean()
            # take gradient step
            '''
            清零演员和评论家的梯度；
            分别对 actor_loss 和 critic_loss 进行反向传播（backward()），计算梯度；
            使用各自的优化器（这里是 Adam 或 RMSprop，之前在构造函数中初始化）更新网络参数。
            '''
            self.actor_optimizer.zero_grad()
            self.critic_optimizer.zero_grad()
            actor_loss.backward()
            critic_loss.backward()
            self.actor_optimizer.step()
            self.critic_optimizer.step()
        '''
        作用：
        - 更新完成后，清空存储在 memory 中的经验数据，为下一次策略更新准备。
        '''
        self.memory.clear()    
    
    
    
    

# 定义训练    
import copy


def train(cfg, env, agent):
    ''' 训练
    '''
    print("开始训练！")
    rewards = []        # 用于记录每个回合的累计奖励。
    steps = []          # 用于记录每回合的步数（即每回合执行了多少步）。
    best_ep_reward = 0  # 记录目前为止评估时得到的最高平均回合奖励。
    output_agent = None # 用于存储表现最佳的智能体，通过深拷贝（copy.deepcopy）保存当前模型。
    for i_ep in range(cfg.train_eps):
        """
        作用：
        - 循环 cfg.train_eps 次，每次代表一个训练回合。
          例如，如果 cfg.train_eps 为 400，则总共训练 400 个回合。
        """
        ep_reward = 0        # 累计当前回合的奖励，初始设为 0。
        ep_step = 0          # 记录当前回合执行的步数。
        state = env.reset()  # 重置环境，获得初始状态。
                             # 例如，若环境是一个赛道或网格环境，
                             # 可能返回一个状态向量，如 [y, x, v_y, v_x]。
        for _ in range(cfg.max_steps):
            ep_step += 1
            '''根据当前状态，通过智能体的 sample_action() 方法选择一个动作。
            若 state 为某状态，演员网络输出概率 [0.7, 0.2, 0.1]，
            则大概率选择动作 0（对应概率 0.7）。'''
            action = agent.sample_action(state)  # 选择动作
            next_state, reward, done, truncated, info = env.step(action)
            # next_state, reward, done, *_ = env.step(action)  # 更新环境，返回transition
            '''将当前步的经验 (state, action, log_probs, reward, done) 
            存入智能体的经验存储器 memory。'''
            agent.memory.push((state, action,agent.log_probs,reward,done))  # 保存transition
            state = next_state  # 更新下一个状态
            agent.update()  # 更新智能体
            ep_reward += reward  # 累加奖励
            if done:
                break
        '''
        每隔 cfg.eval_per_episode 个回合进行一次策略评估。
        例如，如果 cfg.eval_per_episode 为 10，则每 10 个回合进行评估。
        '''
        if (i_ep+1)%cfg.eval_per_episode == 0:
            '''初始化 sum_eval_reward 记录所有评估回合奖励的总和。'''
            sum_eval_reward = 0
            '''循环 cfg.eval_eps 个评估回合（例如 5 个回合）。'''
            for _ in range(cfg.eval_eps):
                '''重置环境得到初始状态。'''
                eval_ep_reward = 0
                state = env.reset()
                for _ in range(cfg.max_steps):
                    '''
                    在最多 cfg.max_steps 步内，
                    使用 predict_action() 选择动作（通常为确定性选择，即取最大概率动作），
                    与环境交互，累计评估回合奖励。
                    '''
                    action = agent.predict_action(state)  # 选择动作
                    next_state, reward, done, truncated, info = env.step(action)
                    # next_state, reward, done, _ = env.step(action)  # 更新环境，返回transition
                    state = next_state  # 更新下一个状态
                    eval_ep_reward += reward  # 累加奖励
                    if done:
                        break
                sum_eval_reward += eval_ep_reward
            '''计算本次评估的平均奖励 mean_eval_reward'''
            mean_eval_reward = sum_eval_reward/cfg.eval_eps
            '''
            如果 mean_eval_reward 大于或等于当前记录的最佳回合奖励 best_ep_reward，
            则更新 best_ep_reward，并使用 copy.deepcopy(agent) 保存当前智能体模型到 output_agent。
            '''
            if mean_eval_reward >= best_ep_reward:
                best_ep_reward = mean_eval_reward
                output_agent = copy.deepcopy(agent)
                print(f"回合：{i_ep+1}/{cfg.train_eps}，奖励：{ep_reward:.2f}，评估奖励：{mean_eval_reward:.2f}，最佳评估奖励：{best_ep_reward:.2f}，更新模型！")
            else:
                print(f"回合：{i_ep+1}/{cfg.train_eps}，奖励：{ep_reward:.2f}，评估奖励：{mean_eval_reward:.2f}，最佳评估奖励：{best_ep_reward:.2f}")
        '''
        将当前回合的步数和累计奖励分别存入 steps 和 rewards 列表中。
        '''
        steps.append(ep_step)
        rewards.append(ep_reward)
    print("完成训练！")
    env.close()
    return output_agent,{'rewards':rewards}


def test(cfg, env, agent):
    """测试函数用于在训练完成后评估智能体表现，不更新模型，只记录回合奖励。"""
    print("开始测试！")
    '''初始化 rewards 和 steps 列表来记录测试回合的奖励和步数。'''
    rewards = []  # 记录所有回合的奖励
    steps = []
    for i_ep in range(cfg.test_eps):
        ep_reward = 0  # 记录一回合内的奖励
        ep_step = 0
        state = env.reset()  # 重置环境，返回初始状态
        for _ in range(cfg.max_steps):
            ep_step+=1
            action = agent.predict_action(state)  # 选择动作
            
            next_state, reward, done, truncated, info = env.step(action)
            # next_state, reward, done, _ = env.step(action)  # 更新环境，返回transition
            state = next_state  # 更新下一个状态
            ep_reward += reward  # 累加奖励
            if done:
                break
        steps.append(ep_step)
        rewards.append(ep_reward)
        print(f"回合：{i_ep+1}/{cfg.test_eps}，奖励：{ep_reward:.2f}")
    print("完成测试")
    env.close()
    return {'rewards':rewards}    
    

# 定义环境    
import gym
import os
import numpy as np
def all_seed(env,seed = 1):
    ''' 万能的seed函数
    '''
    if seed == 0:
        return
    
    
    # env.seed(seed) # env config
    # 如果 FrozenLakeEnv 没有 seed 方法，则为其添加一个
    if not hasattr(env, 'seed'):
        def seed_fn(self, seed=None):
            env.reset(seed=seed)
            return [seed]
        env.seed = seed_fn.__get__(env, type(env))
    
    '''固定 NumPy 中的随机数生成器的种子。'''
    np.random.seed(seed)
    '''固定 Python 内置 random 模块的随机种子。'''
    random.seed(seed)
    '''分别固定 PyTorch 在 CPU 和 GPU 上的随机数生成器种子。'''
    torch.manual_seed(seed) # config for CPU
    torch.cuda.manual_seed(seed) # config for GPU
    '''设置环境变量 PYTHONHASHSEED，使 Python 内部哈希函数输出确定。'''
    os.environ['PYTHONHASHSEED'] = str(seed) # config for python scripts
    # config for cudnn
    torch.backends.cudnn.deterministic = True # 强制使用确定性算法；
    torch.backends.cudnn.benchmark = False    # 关闭动态寻找最优算法；
    torch.backends.cudnn.enabled = False      # 禁用 CuDNN，确保计算完全可控。
def env_agent_config(cfg):
    '''
    根据 cfg.env_name（例如 "CartPole-v0" 或 "Racetrack-v0"）调用 Gym 的 make() 方法创建环境实例。
    '''
    env = gym.make(cfg.env_name) # 创建环境
    '''
    调用上面定义的 all_seed 函数，为环境及全局随机数生成器设置种子。
    '''
    all_seed(env,seed=cfg.seed)
    '''从环境的 observation_space 中提取状态空间的维度。'''
    n_states = env.observation_space.shape[0]
    '''从环境的 action_space 中获取动作数量。'''
    n_actions = env.action_space.n
    print(f"状态空间维度：{n_states}，动作空间维度：{n_actions}")
    # 更新n_states和n_actions到cfg参数中
    setattr(cfg, 'n_states', n_states)
    setattr(cfg, 'n_actions', n_actions) 
    agent = Agent(cfg)
    return env,agent    
    
    
# 设置参数    
import matplotlib.pyplot as plt
import seaborn as sns
class Config:
    def __init__(self) -> None:
        self.env_name = "CartPole-v1" # 环境名字
        self.new_step_api = False # 是否用gym的新api
        self.algo_name = "PPO" # 算法名字
        self.mode = "train" # train or test
        self.seed = 1 # 随机种子
        self.device = "cuda" # device to use
        self.train_eps = 200 # 训练的回合数
        self.test_eps = 20 # 测试的回合数
        self.max_steps = 200 # 每个回合的最大步数
        self.eval_eps = 5 # 评估的回合数
        self.eval_per_episode = 10 # 评估的频率

        self.gamma = 0.99 # 折扣因子
        self.k_epochs = 4 # 更新策略网络的次数
        self.actor_lr = 0.0003 # actor网络的学习率
        self.critic_lr = 0.0003 # critic网络的学习率
        self.eps_clip = 0.2 # epsilon-clip
        self.entropy_coef = 0.01 # entropy的系数
        self.update_freq = 100 # 更新频率
        self.actor_hidden_dim = 256 # actor网络的隐藏层维度
        self.critic_hidden_dim = 256 # critic网络的隐藏层维度

def smooth(data, weight=0.9):  
    '''用于平滑曲线，类似于Tensorboard中的smooth曲线
    对输入数据序列进行指数平滑，使得数据曲线更加平滑，减少噪声。
    '''
    last = data[0] 
    smoothed = []
    for point in data:
        smoothed_val = last * weight + (1 - weight) * point  # 计算平滑值
        smoothed.append(smoothed_val)                    
        last = smoothed_val                                
    return smoothed

def plot_rewards(rewards,cfg, tag='train'):
    ''' 画图
    '''
    sns.set()
    plt.figure()  # 创建一个图形实例，方便同时多画几个图
    plt.title(f"{tag}ing curve on {cfg.device} of {cfg.algo_name} for {cfg.env_name}")
    plt.xlabel('epsiodes')
    plt.plot(rewards, label='rewards')
    plt.plot(smooth(rewards), label='smoothed')
    plt.legend()    
    plt.show()
    
    
# 获取参数
cfg = Config() 
# 训练
env, agent = env_agent_config(cfg)
best_agent,res_dic = train(cfg, env, agent)
 
plot_rewards(res_dic['rewards'], cfg, tag="train")  
# 测试
res_dic = test(cfg, env, best_agent)
plot_rewards(res_dic['rewards'], cfg, tag="test")  # 画出结果