From 7f9d7b2ee624b446af80144c4b6c1933edd163f1 Mon Sep 17 00:00:00 2001 From: Hongru Date: Wed, 25 Mar 2026 15:38:08 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=20TRPO=20=E7=AE=97=E6=B3=95?= =?UTF-8?q?=E5=AE=9E=E7=8E=B0=EF=BC=8C=E5=8C=85=E6=8B=AC=E6=A0=B8=E5=BF=83?= =?UTF-8?q?=E6=95=B0=E5=AD=A6=E5=BC=95=E6=93=8E=E3=80=81=E6=99=BA=E8=83=BD?= =?UTF-8?q?=E4=BD=93=E3=80=81=E7=BD=91=E7=BB=9C=E7=BB=93=E6=9E=84=E5=8F=8A?= =?UTF-8?q?=E8=AE=AD=E7=BB=83=E5=85=A5=E5=8F=A3=EF=BC=8C=E5=AE=8C=E5=96=84?= =?UTF-8?q?=E7=8E=AF=E5=A2=83=E4=BA=A4=E4=BA=92=E4=B8=8E=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E5=A4=84=E7=90=86=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 61 ++++++++-- TRPO/agent.py | 295 +++++++++++++++++++++++++++++++++++++++++++++++++ TRPO/main.py | 142 ++++++++++++++++++++++++ TRPO/models.py | 74 +++++++++++++ TRPO/utils.py | 99 +++++++++++++++++ 5 files changed, 660 insertions(+), 11 deletions(-) create mode 100644 TRPO/agent.py create mode 100644 TRPO/main.py create mode 100644 TRPO/models.py create mode 100644 TRPO/utils.py diff --git a/README.md b/README.md index fc93922..be2e6d7 100644 --- a/README.md +++ b/README.md @@ -14,21 +14,43 @@ RL-Study/ │ ├── SAC.ipynb # SAC (Soft Actor-Critic) 算法 │ └── *_training_results.png # 训练结果可视化 ├── RawBook/ # 原书资源 -└── RL_Algothrithms/ # 核心算法实现 - ├── agents/ # 智能体实现 - │ ├── a2c.py # A2C (Advantage Actor-Critic) - │ └── qac.py # QAC (Soft Actor-Critic / Q-Value Actor-Critic) - ├── networks.py # 神经网络定义 - ├── utils.py # 工具函数 - └── main.py # 训练入口 +├── RL_Algothrithms/ # 核心算法实现 +│ ├── agents/ # 智能体实现 +│ │ ├── a2c.py # A2C (Advantage Actor-Critic) +│ │ ├── qac.py # QAC (Q-Value Actor-Critic) +│ │ ├── off_pac.py # Off-PAC (Off-Policy Actor-Critic) +│ │ ├── ddpg.py # DDPG (Deep Deterministic Policy Gradient) +│ │ └── dpac.py # DPAC (Deterministic Policy Actor-Critic) +│ ├── networks.py # 离散动作空间网络 (Actor, QCritic, VCritic) +│ ├── networks_cont.py # 连续动作空间网络 (ContActor, ContQCritic) +│ ├── utils.py # 工具函数 +│ ├── main.py # 离散动作空间训练入口 (CartPole-v1) +│ ├── disp_main.py # 离散动作空间多算法对比 +│ └── cont_main.py # 连续动作空间训练入口 (Pendulum-v1) +└── TRPO/ # TRPO (Trust Region Policy Optimization) 独立实现 + ├── models.py # ActorNet (高斯策略), CriticNet (值函数) + ├── utils.py # RolloutBuffer, GAE, 共轭梯度, FVP + ├── agent.py # TRPO 智能体 + └── main.py # TRPO 训练入口 (Pendulum-v1) ``` ## 已实现算法 -| 算法 | 文件 | 说明 | -|------|------|------| -| A2C | [a2c.py](RL_Algothrithms/agents/a2c.py) | Advantage Actor-Critic，同步版本 | -| QAC | [qac.py](RL_Algothrithms/agents/qac.py) | Q-Value Actor-Critic，支持 GPU | +### 离散动作空间（CartPole-v1） + +| 算法 | 文件 | 说明 | +| ------- | ------------------------------------------------------ | ------------------------------------------------------ | +| A2C | [agents/a2c.py](RL_Algothrithms/agents/a2c.py) | Advantage Actor-Critic，同步版本，V-critic，带熵正则化 | +| QAC | [agents/qac.py](RL_Algothrithms/agents/qac.py) | Q-Value Actor-Critic，On-policy SARSA 风格，支持 GPU | +| Off-PAC | [agents/off_pac.py](RL_Algothrithms/agents/off_pac.py) | Off-Policy Actor-Critic，带重要性采样，epsilon 探索 | + +### 连续动作空间（Pendulum-v1） + +| 算法 | 文件 | 说明 | +| ---- | ------------------------------------------------ | -------------------------------------------------------------- | +| DDPG | [agents/ddpg.py](RL_Algothrithms/agents/ddpg.py) | Deep Deterministic Policy Gradient，离策略，带目标网络和软更新 | +| DPAC | [agents/dpac.py](RL_Algothrithms/agents/dpac.py) | Deterministic Policy Actor-Critic，在策略 | +| TRPO | [agent.py](TRPO/agent.py) | Trust Region Policy Optimization，共轭梯度法 + 线搜索 + GAE | ## 环境配置 @@ -38,10 +60,27 @@ pip install torch numpy matplotlib gymnasium ## 快速开始 +### 离散动作空间（CartPole-v1） + ```bash cd RL_Algothrithms python main.py --agent a2c # 训练 A2C python main.py --agent qac # 训练 QAC +python disp_main.py # 多算法对比训练 +``` + +### 连续动作空间（Pendulum-v1） + +```bash +cd RL_Algothrithms +python cont_main.py # 训练 DDPG 和 DPAC +``` + +### TRPO + +```bash +cd TRPO +python main.py # 训练 TRPO ``` ## 关于原书 diff --git a/TRPO/agent.py b/TRPO/agent.py new file mode 100644 index 0000000..4693b60 --- /dev/null +++ b/TRPO/agent.py @@ -0,0 +1,295 @@ +import torch +import torch.nn as nn +import numpy as np +from models import ActorNet, CriticNet +import torch.optim as optim +from torch.distributions import Normal + +# ========================================== +# 辅助工具函数：参数与向量的互相转换 +# ========================================== + +def get_flat_params_from(model): + """ + 把模型中散落在各个层的所有参数，按顺序拼接成一个巨大的一维 Tensor。 + 相当于数学推导中的参数向量 theta。 + """ + params = [] + for param in model.parameters(): + params.append(param.data.view(-1)) + return torch.cat(params) + +def set_flat_params_to(model, flat_params): + """ + 把计算好的新一维参数向量，按照对应尺寸还原、塞回神经网络的各个层中。 + 这是用来真正执行 theta_new 赋值的。 + """ + prev_ind = 0 + for param in model.parameters(): + flat_size = int(np.prod(param.size())) + # 截取对应长度的数据，并 reshape 回原始层的形状 + param.data.copy_(flat_params[prev_ind:prev_ind + flat_size].view(param.size())) + prev_ind += flat_size + +def get_flat_grad_from(loss, model): + """ + 对指定的 loss 求网络参数的梯度，并直接压平成一个一维 Tensor 返回。 + 相当于计算梯度向量 g。 + """ + # retain_graph=True 是因为我们后面算二阶导可能还会用到当前的计算图 + grads = torch.autograd.grad(loss, model.parameters(), retain_graph=True) + return torch.cat([grad.view(-1) for grad in grads]) + + +# ========================================== +# TRPO 核心数学引擎 +# ========================================== + +def conjugate_gradient(fvp_func, b, nsteps=10, residual_tol=1e-10): + """ + 共轭梯度法 (Conjugate Gradient, CG) + 用于近似求解线性方程组 Ax = b，在这里也就是求解 As = g。 + 它不需要矩阵 A，只需要一个能计算 A*v 的函数 (也就是下面的 fvp_func)。 + + 参数: + fvp_func: 传入一个向量 v，返回 Fisher矩阵乘该向量的结果 A*v + b: 目标向量 (在 TRPO 中就是策略梯度向量 g) + nsteps: 迭代次数 (通常 10 次就能逼近得很好) + """ + x = torch.zeros_like(b) # 初始解设为 0 + r = b.clone() # 初始残差 + p = b.clone() # 初始搜索方向 + rdotr = torch.dot(r, r) # 残差的内积 + + for i in range(nsteps): + # 计算 A * p (也就是 FVP) + Ap = fvp_func(p) + + # 计算步长 alpha + alpha = rdotr / (torch.dot(p, Ap) + 1e-8) + + # 更新解 x + x += alpha * p + + # 更新残差 r + r -= alpha * Ap + new_rdotr = torch.dot(r, r) + + # 如果残差已经足够小，提前退出 + if new_rdotr < residual_tol: + break + + # 计算方向更新系数 beta + beta = new_rdotr / rdotr + + # 更新搜索方向 p + p = r + beta * p + rdotr = new_rdotr + + return x + + +def fisher_vector_product(actor_net, states, vector, damping=0.1): + """ + 海森向量积 / 费雪信息矩阵-向量积 (FVP) + 这是 TRPO 的绝对核心黑科技：通过连续两次自动求导，计算 A * v，无需显式构造 A！ + """ + # 1. 用当前的策略网络计算出旧的均值和标准差 (停止梯度更新，作为基准点) + mean_old, std_old = actor_net(states) + mean_old = mean_old.detach() + std_old = std_old.detach() + + # 2. 重新进行一次前向传播，保留计算图 + mean, std = actor_net(states) + + # 3. 解析计算 KL 散度 (高斯分布的精确闭式解) + # 公式: log(std/std_old) + (std_old^2 + (mean_old - mean)^2) / (2 * std^2) - 0.5 + # 注意：TRPO 是对状态空间求期望，所以最后要求均值 (mean) + kl = torch.log(std / std_old) + (std_old.pow(2) + (mean_old - mean).pow(2)) / (2.0 * std.pow(2)) - 0.5 + kl = kl.sum(dim=1, keepdim=True).mean() + + # 4. 第一次求导：计算 KL 对网络参数的一阶梯度 (Jacobian) + # create_graph=True 极其关键：它让一阶梯度本身也成为计算图的一部分，为求二阶导做准备 + grads = torch.autograd.grad(kl, actor_net.parameters(), create_graph=True) + flat_grad_kl = torch.cat([grad.view(-1) for grad in grads]) + + # 5. 计算一阶梯度向量与传入向量 v 的内积 + # 这个点积的结果是一个标量 (Scalar) + kl_v = torch.dot(flat_grad_kl, vector) + + # 6. 第二次求导：对上面的点积标量再次求参数的梯度 + # 根据微积分法则，梯度的点积的梯度 = Hessian * v + grads_v = torch.autograd.grad(kl_v, actor_net.parameters()) + flat_grad_grad_kl = torch.cat([grad.contiguous().view(-1) for grad in grads_v]).detach() + + # 7. 加上阻尼项 (Damping) + # 给对角线加上一个小常数 (damping * vector)，确保 FIM 矩阵正定，提升 CG 求解的数值稳定性 + return flat_grad_grad_kl + vector * damping + + +class TRPOAgent: + def __init__(self, state_dim, action_dim, max_kl=0.01, cg_iters=10, cg_residual_tol=1e-10, cg_damping=0.1): + """ + 初始化 TRPO 智能体 + """ + self.actor = ActorNet(state_dim, action_dim) + self.critic = CriticNet(state_dim) + + # Critic 使用普通的 Adam 优化器即可 (学习率设为 1e-3) + self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=1e-3) + + # TRPO 超参数 + self.max_kl = max_kl # 也就是公式里的 delta (信任域边界) + self.cg_iters = cg_iters # 共轭梯度法的迭代次数 + self.cg_residual_tol = cg_residual_tol + self.cg_damping = cg_damping # 海森矩阵的阻尼系数 + + def compute_surrogate_obj(self, states, actions, old_log_probs, advantages): + """ + 计算替代目标函数 (Surrogate Objective) + 公式: L = E[ (pi_new / pi_old) * A ] + """ + # 计算当前策略下动作的对数概率 + mean, std = self.actor(states) + dist = Normal(mean, std) + # 注意: 如果动作是多维的，需要对各维度的 log_prob 求和 + new_log_probs = dist.log_prob(actions).sum(dim=1, keepdim=True) + + # 计算重要性采样比率 (Ratio): exp(log_new - log_old) = new / old + ratio = torch.exp(new_log_probs - old_log_probs) + + # 替代目标函数 (最大化目标，所以返回均值) + surrogate_obj = (ratio * advantages).mean() + return surrogate_obj + + def update(self, rollout_buffer, next_state, done): + """ + 执行一次完整的 TRPO 更新 (Actor 和 Critic) + 参数: + rollout_buffer: 收集好数据的经验池 + next_state: 轨迹结束时的下一个状态 (用于计算 GAE) + done: 轨迹是否结束的标志位 + """ + + # 在更新前，先用 Critic 算出 GAE 和 Returns + with torch.no_grad(): + # 如果环境已经 done（比如倒立摆摔倒或超时），那未来的预期价值就是 0 + # 否则，用 Critic 网络预测一下 next_state 的价值 + if done: + last_value = 0.0 + else: + next_state_tensor = torch.tensor(next_state, dtype=torch.float32) + last_value = self.critic(next_state_tensor).item() + + # 真正调用 buffer 的函数，计算出 numpy 格式的 returns 和 advantages + returns_np, advantages_np = rollout_buffer.compute_returns_and_advantages(last_value) + + # 转成深度学习需要的 Tensor，并对齐维度 [Batch, 1] + returns = torch.tensor(returns_np, dtype=torch.float32).view(-1, 1) + advantages = torch.tensor(advantages_np, dtype=torch.float32).view(-1, 1) + + # 1. 从经验池中获取并整理数据 + states, actions, old_log_probs = rollout_buffer.get_data() + + + # ========================================== + # 第一步：计算目标函数的一阶梯度 (g) + # ========================================== + surrogate_obj_old = self.compute_surrogate_obj(states, actions, old_log_probs, advantages) + # 注意 TRPO 是最大化目标，所以 loss 是负的 objective + loss_actor = -surrogate_obj_old + + # 使用我们在 Stage 3 写的辅助函数获取铺平的一阶梯度 g + g = get_flat_grad_from(loss_actor, self.actor).detach() + + # ========================================== + # 第二步：用共轭梯度法 (CG) 求自然梯度方向 (s) + # ========================================== + # 定义一个局部函数，把 states 封进去，专供 CG 调用 + def fvp_callable(v): + return fisher_vector_product(self.actor, states, v, self.cg_damping) + + # 解方程 As = g，得到搜索方向 step_dir (即公式里的 s) + step_dir = conjugate_gradient(fvp_callable, -g, self.cg_iters, self.cg_residual_tol) + + # ========================================== + # 第三步：计算理论最大步长 (beta) + # ========================================== + # s^T A s (通过 FVP 再算一次) + sAs = torch.dot(step_dir, fvp_callable(step_dir)) + + # 为了防止除以 0 的数值不稳定，加上 1e-8 + # 公式: beta = sqrt( 2 * delta / (s^T A s) ) + beta = torch.sqrt(2 * self.max_kl / (sAs + 1e-8)) + + # 完整的最大更新步长向量 + full_step = beta * step_dir + + # ========================================== + # 第四步：回溯线搜索 (Backtracking Line Search) + # ========================================== + # 获取当前 Actor 的初始参数 + old_params = get_flat_params_from(self.actor) + + success = False + fraction = 1.0 # 步长衰减系数的初始值 + + # 尝试 10 次缩小步长 + for i in range(10): + # 试探性地迈出一步: theta_new = theta_old + fraction * full_step + new_params = old_params + fraction * full_step + set_flat_params_to(self.actor, new_params) + + # 在新参数下，重新评估目标函数和 KL 散度 + with torch.no_grad(): + # 1. 检查目标函数是否真的提升了 + surrogate_obj_new = self.compute_surrogate_obj(states, actions, old_log_probs, advantages) + improvement = surrogate_obj_new - surrogate_obj_old + + # 2. 检查 KL 散度是否满足约束 (<= max_kl) + mean_new, std_new = self.actor(states) + mean_old, std_old = self.actor(states) # 注意这里应该用一开始记录的固定不变的 old 值，为了严谨，我们在最开始算一次并脱离计算图 + # (为了简化代码，更严谨的做法是在线搜索外先算好 old 分布参数传进来) + + # 我们在这里写一个快速的 KL 检查逻辑 + mean_old_frozen, std_old_frozen = self.actor(states) + set_flat_params_to(self.actor, old_params) # 临时切回老参数获取冻结的分布 + mean_old_frozen, std_old_frozen = self.actor(states) + mean_old_frozen, std_old_frozen = mean_old_frozen.detach(), std_old_frozen.detach() + + set_flat_params_to(self.actor, new_params) # 再切回新参数算 KL + mean_new, std_new = self.actor(states) + kl = torch.log(std_new / std_old_frozen) + (std_old_frozen.pow(2) + (mean_old_frozen - mean_new).pow(2)) / (2.0 * std_new.pow(2)) - 0.5 + kl_mean = kl.sum(dim=1).mean() + + # 判断双重安全条件：目标提升了，且 KL 没超标 + if improvement.item() > 0 and kl_mean.item() <= self.max_kl: + success = True + print(f"线搜索成功: 迭代次数 {i}, 提升值 {improvement.item():.4f}, KL 散度 {kl_mean.item():.4f}") + break + else: + # 如果失败了，步长减半，再试一次 + fraction *= 0.5 + + # 如果 10 次尝试都失败了，说明这里地形太差，安全起见我们不更新 Actor 了 + if not success: + print("线搜索失败，放弃此次 Actor 更新，保持原参数。") + set_flat_params_to(self.actor, old_params) + + # ========================================== + # 第五步：更新 Critic (价值网络) + # ========================================== + # Critic 的更新非常简单，就是标准的深度学习监督训练，目标是逼近真实的 Return + criterion = nn.MSELoss() + + # 多次迭代更新 Critic 以充分拟合价值 (通常设为 10 次) + for _ in range(10): + values = self.critic(states) + loss_critic = criterion(values, returns) + + self.critic_optimizer.zero_grad() + loss_critic.backward() + self.critic_optimizer.step() + + # 清空经验池，为下一次环境交互做准备 + rollout_buffer.clear() diff --git a/TRPO/main.py b/TRPO/main.py new file mode 100644 index 0000000..70d86fe --- /dev/null +++ b/TRPO/main.py @@ -0,0 +1,142 @@ +import gymnasium as gym +import torch +import numpy as np +import matplotlib.pyplot as plt # 新增：用于画图 +from agent import TRPOAgent +from utils import RolloutBuffer + +# 新增：用于计算滑动平均，让曲线更平滑 +def moving_average(data, window_size=5): + """计算简单滑动平均""" + if len(data) < window_size: + return data + return np.convolve(data, np.ones(window_size)/window_size, mode='valid') + +def plot_learning_curve(rewards, smoothed_rewards): + """绘制学习曲线""" + plt.figure(figsize=(10, 6)) + + # 绘制原始奖励（浅色） + plt.plot(rewards, color='blue', alpha=0.3, label='Raw Average Reward') + + # 绘制平滑后的奖励（深色，粗线） + # 注意：滑动平均后数据点会变少，需要调整 X 轴起始位置 + if len(smoothed_rewards) > 0: + x_ticks = np.arange(len(rewards) - len(smoothed_rewards), len(rewards)) + plt.plot(x_ticks, smoothed_rewards, color='red', linewidth=2, label='Smoothed Reward (MA-5)') + + plt.title('TRPO Training Performance on Pendulum-v1') + plt.xlabel('Iteration') + plt.ylabel('Average Reward') + plt.grid(True, linestyle='--', alpha=0.5) + plt.legend() + + # 保存图片 + plt.savefig('trpo_training_curve.png') + print("\n训练曲线图已保存为 'trpo_training_curve.png'") + + # 如果有 GUI 界面则显示 + try: + plt.show() + except Exception: + print("无法显示图形界面（可能是无头服务器），已跳过 plt.show()。") + +def main(): + # 1. 初始化环境 + #env_name = 'Pendulum-v1' + # env = gym.make('Pendulum-v1') + # 兼容 Gymnasium + try: + env = gym.make('Pendulum-v1', render_mode=None) + except Exception: + env = gym.make('Pendulum-v1') + + state_dim = env.observation_space.shape[0] # type: ignore + action_dim = env.action_space.shape[0] # type: ignore + max_action = float(env.action_space.high[0]) # type: ignore + min_action = float(env.action_space.low[0]) # type: ignore + + print(f"环境加载成功! 状态维度: {state_dim}, 动作维度: {action_dim}") + print(f"动作范围: [{min_action}, {max_action}]") + + # 2. 初始化智能体和经验池 + agent = TRPOAgent(state_dim, action_dim, max_kl=0.01) + buffer = RolloutBuffer() + + # 3. 设置训练超参数 + max_iterations = 300 # 训练迭代总轮数 + batch_size = 2000 # 每次更新收集的步数 + + # 新增：用于记录绘图数据 + history_rewards = [] + + # 4. 主训练循环 + for iteration in range(max_iterations): + state = env.reset() + if isinstance(state, tuple): state = state[0] + + episode_rewards = [] + ep_reward = 0 + steps = 0 + done = False + + while steps < batch_size: + state_tensor = torch.FloatTensor(state).unsqueeze(0) + + with torch.no_grad(): + action, log_prob = agent.actor.get_action(state_tensor) + value = agent.critic(state_tensor) + + action_np = action.squeeze(0).numpy() + clipped_action = np.clip(action_np, min_action, max_action) + + # env.step 兼容性 + step_result = env.step(clipped_action) + if len(step_result) == 5: + next_state, reward, terminated, truncated, _ = step_result + done = terminated or truncated + else: + next_state, reward, done, _ = step_result + + ep_reward += reward + buffer.add(state=state, action=action_np, reward=reward, + next_state=next_state, done=done, + log_prob=log_prob.item(), value=value.item()) + + state = next_state + steps += 1 + + if done: + episode_rewards.append(ep_reward) + state = env.reset() + if isinstance(state, tuple): state = state[0] + ep_reward = 0 + + print(f"正在更新参数 (Iteration {iteration + 1}/{max_iterations})...") + agent.update(buffer, state, done) + + # 5. 记录和打印日志 + if episode_rewards: + avg_reward = np.mean(episode_rewards) + history_rewards.append(avg_reward) # 新增：记录数据 + print(f"Iteration: {iteration + 1} | Average Reward: {avg_reward:.2f} | Max Reward: {np.max(episode_rewards):.2f}") + print("-" * 50) + else: + # 如果 batch_size 刚好结束时没有完成任何 episode， + # 为了画图不中断，我们沿用上一次的奖励（或者简单处理） + if history_rewards: + history_rewards.append(history_rewards[-1]) + else: + history_rewards.append(-2000) # 初始默认低分 + + print("训练结束！") + env.close() + + # 6. 新增：绘图 + print("正在生成训练曲线图...") + smoothed = moving_average(history_rewards, window_size=5) + plot_learning_curve(history_rewards, smoothed) + +if __name__ == '__main__': + # 确保安装了 matplotlib: pip install matplotlib + main() diff --git a/TRPO/models.py b/TRPO/models.py new file mode 100644 index 0000000..4d8541f --- /dev/null +++ b/TRPO/models.py @@ -0,0 +1,74 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributions import Normal + +class ActorNet(nn.Module): + """ + 策略网络 (Actor)：输入环境状态，输出连续动作的高斯分布参数 (均值和标准差) + """ + def __init__(self, state_dim, action_dim, hidden_dim=64): + super(ActorNet, self).__init__() + # 定义两层隐藏层，提取状态特征 + self.fc1 = nn.Linear(state_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, hidden_dim) + + # 输出层：计算动作的均值 (Mean) + self.mean_layer = nn.Linear(hidden_dim, action_dim) + + # 定义对数标准差 (Log Standard Deviation) 为可训练的独立参数 + # 注意：在 TRPO/PPO 中，通常让标准差独立于状态，这能让训练更稳定 + # 初始值设为 0，意味着初始标准差为 exp(0) = 1.0 + self.log_std = nn.Parameter(torch.zeros(1, action_dim)) + + def forward(self, state): + # 前向传播提取特征 + x = F.tanh(self.fc1(state)) + x = F.tanh(self.fc2(x)) + + # 计算均值 + mean = self.mean_layer(x) + + # 将对数标准差扩展到与 batch size 相同的维度 + log_std = self.log_std.expand_as(mean) + # 转化为标准差 + std = torch.exp(log_std) + + return mean, std + + def get_action(self, state): + """ + 根据当前状态采样动作，并返回对应的对数概率 (log probability) + """ + # 获取当前状态的均值和标准差 + mean, std = self.forward(state) + + # 构建正态分布 (高斯分布) + dist = Normal(mean, std) + + # 从分布中采样一个动作 + action = dist.sample() + + # 返回采样动作及其对数概率 (后续计算重要性采样权重时必须用到) + return action, dist.log_prob(action) + + +class CriticNet(nn.Module): + """ + 价值网络 (Critic)：输入环境状态，评估该状态的预期收益 (标量 V 值) + """ + def __init__(self, state_dim, hidden_dim=64): + super(CriticNet, self).__init__() + # 定义隐藏层 + self.fc1 = nn.Linear(state_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, hidden_dim) + + # 输出层：计算状态价值 V(s)，输出维度为 1 + self.value_layer = nn.Linear(hidden_dim, 1) + + def forward(self, state): + # 前向传播计算价值 + x = F.tanh(self.fc1(state)) + x = F.tanh(self.fc2(x)) + value = self.value_layer(x) + return value diff --git a/TRPO/utils.py b/TRPO/utils.py new file mode 100644 index 0000000..0377e91 --- /dev/null +++ b/TRPO/utils.py @@ -0,0 +1,99 @@ +import torch +import numpy as np + +class RolloutBuffer: + """ + 经验回放池：用于收集智能体与环境交互的轨迹数据， + 并在一个回合（或一个 Batch）结束后计算优势函数 GAE 和目标价值。 + """ + def __init__(self): + # 初始化存储列表 + self.states = [] + self.actions = [] + self.rewards = [] + self.next_states = [] + self.dones = [] + self.log_probs = [] + self.values = [] + + def add(self, state, action, reward, next_state, done, log_prob, value): + """ + 向池子中添加一步交互的数据 + """ + self.states.append(state) + self.actions.append(action) + self.rewards.append(reward) + self.next_states.append(next_state) + self.dones.append(done) + self.log_probs.append(log_prob) + self.values.append(value) + + def clear(self): + """ + 清空池子，准备收集下一批数据 + """ + self.states.clear() + self.actions.clear() + self.rewards.clear() + self.next_states.clear() + self.dones.clear() + self.log_probs.clear() + self.values.clear() + + def compute_returns_and_advantages(self, last_value, gamma=0.99, lam=0.95): + """ + 计算广义优势估计 (GAE) 和 目标价值 (Returns)。 + 这是 TRPO/PPO 最核心的数据处理步骤！ + + 参数: + last_value: 截断处（或回合结束时）的最后一个状态的 V 值。 + gamma: 折扣因子 (Discount factor)。 + lam: GAE 的平滑参数 (Lambda)，用于权衡偏差和方差。 + """ + # 将列表转换为 NumPy 数组，方便进行向量化运算 + rewards = np.array(self.rewards, dtype=np.float32) + values = np.array(self.values, dtype=np.float32) + dones = np.array(self.dones, dtype=np.float32) + + # 预分配数组空间 + advantages = np.zeros_like(rewards, dtype=np.float32) + last_gae_lam = 0 + + # 逆序遍历轨迹：从最后一步往前推算 + for t in reversed(range(len(rewards))): + if t == len(rewards) - 1: + # 如果是最后一步，next_value 就是传入的 last_value + next_non_terminal = 1.0 - dones[t] + next_value = last_value + else: + # 否则，next_value 就是下一步的 value + next_non_terminal = 1.0 - dones[t] + next_value = values[t + 1] + + # 计算 TD 误差 (Temporal Difference Error) + # delta = r_t + gamma * V(s_{t+1}) - V(s_t) + delta = rewards[t] + gamma * next_value * next_non_terminal - values[t] + + # 递推计算 GAE + # A_t = delta_t + gamma * lambda * A_{t+1} + advantages[t] = last_gae_lam = delta + gamma * lam * next_non_terminal * last_gae_lam + + # 目标价值 = 优势函数 + 状态价值 + returns = advantages + values + + # 优势函数标准化 (Advantage Normalization) + # 这是一个极度重要的工程 Trick，能大幅提升训练稳定性 + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) + + return returns, advantages + + def get_data(self): + """ + 将收集到的所有数据转换为 PyTorch Tensor，供后续网络训练使用 + """ + # 将 NumPy 数组转为 Tensor + state_tensor = torch.tensor(np.array(self.states), dtype=torch.float32) + action_tensor = torch.tensor(np.array(self.actions), dtype=torch.float32) + old_log_probs_tensor = torch.tensor(np.array(self.log_probs), dtype=torch.float32) + + return state_tensor, action_tensor, old_log_probs_tensor