commit f3d2d1a85fca1ca31d4fce728a92a7e396cc2968 Author: Hongru Date: Sat Apr 4 15:51:08 2026 +0800 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7d1c5e5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,42 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +*.egg-info/ +dist/ +build/ + +# Virtual environments +venv/ +.venv/ +env/ +.env/ + +# PyTorch model checkpoints +*.pth + +# Images +step_response.png + +# Jupyter +.ipynb_checkpoints/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo + +# Logs +*.log +logs/ + +# TensorBoard +runs/ +events.out.tfevents.* + +# OS +.DS_Store +Thumbs.db diff --git a/algorithms/sac.py b/algorithms/sac.py new file mode 100644 index 0000000..1b349b5 --- /dev/null +++ b/algorithms/sac.py @@ -0,0 +1,110 @@ +import torch +import torch.nn.functional as F +import torch.optim as optim +from models.networks import Actor, Critic +import copy + +class SAC(object): + def __init__(self, state_dim, action_dim, max_action, config): + """ + 初始化 Soft Actor-Critic 算法 + """ + # 设备配置 (CPU 或 GPU) + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # 从 config 字典中加载超参数 + self.gamma = config.get('gamma', 0.99) # 折扣因子 + self.tau = config.get('tau', 0.005) # 目标网络软更新系数 (论文公式 9 下方) + self.alpha = config.get('alpha', 0.2) # 熵的温度参数 (控制探索的随机性) + + # 1. 实例化策略网络 (Actor) + self.actor = Actor(state_dim, action_dim, max_action).to(self.device) + self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=config.get('lr', 3e-4)) + + # 2. 实例化价值网络 (Critic) - 内部已经包含了 Q1 和 Q2 + self.critic = Critic(state_dim, action_dim).to(self.device) + self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=config.get('lr', 3e-4)) + + # 3. 实例化目标价值网络 (Target Critic) + # 用 copy.deepcopy 完美复制一份初始参数，并冻结其梯度计算 + self.critic_target = copy.deepcopy(self.critic) + for param in self.critic_target.parameters(): + param.requires_grad = False + + def select_action(self, state, evaluate=False): + """ + 与环境交互时使用的动作选择函数 + """ + # 将输入的状态转换为 PyTorch Tensor + state = torch.FloatTensor(state).unsqueeze(0).to(self.device) + + # 论文技巧：评估（测试）时使用均值动作，训练时使用采样动作 + with torch.no_grad(): + if evaluate: + _, _, action = self.actor.sample(state) # 第三个返回值是均值 + else: + action, _, _ = self.actor.sample(state) # 第一个返回值是加了噪声的采样值 + + # 转换回 numpy 数组，送给 Gym 环境执行 + return action.cpu().data.numpy().flatten() + + def update(self, replay_buffer, batch_size): + """ + 算法的核心心跳：从经验池采样并更新神经网络参数 + """ + # 从 Replay Buffer 中随机抽取一个 Batch 的数据 + state, action, reward, next_state, not_done = replay_buffer.sample(batch_size) + + # ================================================================= # + # 1. 更新 Critic # + # ================================================================= # + with torch.no_grad(): + # 拿到下一个状态的动作和其对应的对数概率 (用于计算熵) + next_action, next_log_prob, _ = self.actor.sample(next_state) + + # 使用目标网络计算下一个状态的 Q 值 (Q1 和 Q2) + target_Q1, target_Q2 = self.critic_target(next_state, next_action) + + # 【核心对抗高估】：取两个 Q 值的最小值 + target_Q = torch.min(target_Q1, target_Q2) + + # 【软贝尔曼方程】：目标 Q 值 = 奖励 + gamma * (目标 Q - alpha * 熵) + # 注意这里加上了 -self.alpha * next_log_prob，这就是论文中“把熵当做奖励”的体现 + target_Q = reward + not_done * self.gamma * (target_Q - self.alpha * next_log_prob) + + # 获取当前状态和动作对应的 Q 值预测 + current_Q1, current_Q2 = self.critic(state, action) + + # 计算 Critic 的损失 (均方误差 MSE) + critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q) + + # 优化 Critic 网络 + self.critic_optimizer.zero_grad() + critic_loss.backward() + self.critic_optimizer.step() + + # ================================================================= # + # 2. 更新 Actor # + # ================================================================= # + # 让当前 Actor 对**当前状态**重新采样一个动作 (注意：不能用 Buffer 里的旧动作) + pi_action, log_prob, _ = self.actor.sample(state) + + # 拿到更新后的 Critic 对这个新动作的打分 + q1_pi, q2_pi = self.critic(state, pi_action) + min_q_pi = torch.min(q1_pi, q2_pi) + + # 计算 Actor 的损失：最小化 (alpha * log_prob - min_Q) + # 等价于最大化 (min_Q - alpha * log_prob) -> 既要 Q 值大，又要熵大（分布广） + actor_loss = (self.alpha * log_prob - min_q_pi).mean() + + # 优化 Actor 网络 + self.actor_optimizer.zero_grad() + actor_loss.backward() + self.actor_optimizer.step() + + # ================================================================= # + # 3. 软更新 Target Critic # + # ================================================================= # + # 使用 EMA (指数移动平均) 缓慢将前线 Critic 的参数移交给 Target Critic + for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): + target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) diff --git a/configs/pendulum_config.yaml b/configs/pendulum_config.yaml new file mode 100644 index 0000000..c627dce --- /dev/null +++ b/configs/pendulum_config.yaml @@ -0,0 +1,22 @@ +# ========================================== +# Soft Actor-Critic (SAC) - Pendulum-v1 配置 +# ========================================== + +# --- 算法核心参数 --- +gamma: 0.99 # 折扣因子 (越接近1越看重长期收益) +tau: 0.005 # 目标网络软更新系数 (EMA平滑系数，越小越稳定) +alpha: 0.2 # 熵温度系数 (控制探索力度，Pendulum中0.2比较合适，如果是复杂环境可能需要调整或自动学习) +lr: 0.0003 # 学习率 (Actor 和 Critic 保持一致，3e-4 是 Adam 优化器的万金油) + +# --- 经验回放池参数 --- +buffer_size: 1000000 # 回放池最大容量 (100万条) +batch_size: 256 # 每次梯度更新抽样的 batch 大小 + +# --- 训练循环控制 --- +max_episodes: 200 # 总共训练多少个回合 (Episode) +max_steps: 200 # 每个回合最多走多少步 (Gym Pendulum 默认 200 步截断) +start_steps: 10000 # 纯随机动作探索的步数 (用来快速填补经验池的高多样性数据) + +# --- 扩展参数 (为了后续画图和保存模型备用) --- +eval_freq: 10 # 每隔多少个回合评估一次策略 +save_model: true # 是否保存最终模型权重 diff --git a/envs/__init__.py b/envs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/eval.py b/eval.py new file mode 100644 index 0000000..69e4387 --- /dev/null +++ b/eval.py @@ -0,0 +1,116 @@ +# ============================================================================== +# Author: Hongru Liu +# Affiliation: School of Power and Energy, Northwestern Polytechnical University +# Version: 1.0 +# Contact: hongruliu@mail.nwpu.edu.cn +# ============================================================================== + +import gymnasium as gym +import torch +import numpy as np +import matplotlib.pyplot as plt +import yaml +import os +from algorithms.sac import SAC + +def evaluate_and_plot(model_path): + """ + 加载训练好的模型，在环境中运行一个回合，并绘制类似阶跃响应的状态轨迹图 + """ + # 1. 读取配置和初始化环境 + config_path = os.path.join(os.path.dirname(__file__), 'configs', 'pendulum_config.yaml') + with open(config_path, "r", encoding="utf-8") as f: + config = yaml.safe_load(f) + + env = gym.make('Pendulum-v1') + state_dim = env.observation_space.shape[0] + action_dim = env.action_space.shape[0] + max_action = float(env.action_space.high[0]) + + # 2. 实例化算法并加载权重 + agent = SAC(state_dim, action_dim, max_action, config) + + if os.path.exists(model_path): + # 仅加载 Actor 的权重，因为测试阶段不需要 Critic 参与评估 + agent.actor.load_state_dict(torch.load(model_path, map_location=agent.device)) + print(f"[*] Successfully loaded model weights from: {model_path}") + else: + print(f"[!] Model file not found: {model_path}") + return + + # 3. 运行单个 Episode，收集状态和动作数据 + state, _ = env.reset() + + # 用于记录作图的数据列表 + angles = [] + velocities = [] + actions = [] + time_steps = [] + + episode_reward = 0 + for step in range(config['max_steps']): + # 【重点】：设置 evaluate=True，让网络输出确定的均值动作，关闭随机探索 + action = agent.select_action(state, evaluate=True) + + # 记录当前步的数据 + # Pendulum 的状态是 [cos(theta), sin(theta), theta_dot] + cos_theta, sin_theta, theta_dot = state + # 将 cos 和 sin 转换为实际角度 (弧度)，范围 [-pi, pi] + angle = np.arctan2(sin_theta, cos_theta) + + angles.append(angle) + velocities.append(theta_dot) + actions.append(action[0]) + time_steps.append(step) + + # 与环境交互 + next_state, reward, terminated, truncated, _ = env.step(action) + state = next_state + episode_reward += reward + + if terminated or truncated: + break + + print(f"[*] Evaluation completed. Total Reward: {episode_reward:.2f}") + + # 4. 绘制状态轨迹图 (类阶跃响应) + fig, axs = plt.subplots(3, 1, figsize=(10, 10), facecolor='white', sharex=True) + + # 子图 1: 角度响应 (Pendulum Angle) + axs[0].set_facecolor('white') + axs[0].plot(time_steps, angles, linewidth=2.5, color='#d62728', label='Angle (rad)') + axs[0].axhline(y=0, color='black', linestyle='--', linewidth=1.5, alpha=0.5) # 目标线 + axs[0].set_ylabel('Angle [rad]', fontsize=14, fontweight='bold') + axs[0].set_title('Pendulum Step Response Simulation', fontsize=16, fontweight='bold') + axs[0].tick_params(axis='y', labelsize=12) + axs[0].grid(True, linestyle='--', alpha=0.7) + axs[0].legend(fontsize=12, loc='upper right') + + # 子图 2: 角速度响应 (Angular Velocity) + axs[1].set_facecolor('white') + axs[1].plot(time_steps, velocities, linewidth=2.5, color='#2ca02c', label='Velocity (rad/s)') + axs[1].axhline(y=0, color='black', linestyle='--', linewidth=1.5, alpha=0.5) # 目标线 + axs[1].set_ylabel('Velocity [rad/s]', fontsize=14, fontweight='bold') + axs[1].tick_params(axis='y', labelsize=12) + axs[1].grid(True, linestyle='--', alpha=0.7) + axs[1].legend(fontsize=12, loc='upper right') + + # 子图 3: 控制输入 (Control Torque) + axs[2].set_facecolor('white') + axs[2].plot(time_steps, actions, linewidth=2.5, color='#9467bd', label='Torque (N.m)') + axs[2].set_xlabel('Time Step', fontsize=14, fontweight='bold') + axs[2].set_ylabel('Control Input', fontsize=14, fontweight='bold') + axs[2].tick_params(axis='both', which='major', labelsize=12) + axs[2].grid(True, linestyle='--', alpha=0.7) + axs[2].legend(fontsize=12, loc='upper right') + + plt.tight_layout() + save_path = "step_response.png" + plt.savefig(save_path, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none') + plt.close() + print(f"[*] Step response plot saved to: {save_path}") + +if __name__ == "__main__": + # 替换为你实际训练出来的最后一次模型文件名称 + # 例如：sac_actor_pendulum_ep200.pth + evaluate_and_plot("sac_actor_pendulum_ep200.pth") diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/networks.py b/models/networks.py new file mode 100644 index 0000000..eedfc20 --- /dev/null +++ b/models/networks.py @@ -0,0 +1,98 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributions import Normal + +# 确保代码可以在 GPU 上跑（如果有的话） +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +class Critic(nn.Module): + def __init__(self, state_dim, action_dim): + super(Critic, self).__init__() + + # Q1 网络架构 + self.l1 = nn.Linear(state_dim + action_dim, 256) + self.l2 = nn.Linear(256, 256) + self.l3 = nn.Linear(256, 1) + + # Q2 网络架构（和 Q1 完全一样，但参数是独立初始化的） + self.l4 = nn.Linear(state_dim + action_dim, 256) + self.l5 = nn.Linear(256, 256) + self.l6 = nn.Linear(256, 1) + + def forward(self, state, action): + # 把状态和动作拼接在一起，作为 Q 网络的输入 + sa = torch.cat([state, action], 1) + + # Q1 的前向传播 + q1 = F.relu(self.l1(sa)) + q1 = F.relu(self.l2(q1)) + q1 = self.l3(q1) + + # Q2 的前向传播 + q2 = F.relu(self.l4(sa)) + q2 = F.relu(self.l5(q2)) + q2 = self.l6(q2) + + # 训练时，我们需要同时返回两个 Q 值来算误差 + return q1, q2 + + +# 定义标准差的上下界，防止网络输出极端值导致计算崩溃（NaN） +LOG_SIG_MAX = 2 +LOG_SIG_MIN = -20 + +class Actor(nn.Module): + def __init__(self, state_dim, action_dim, max_action): + super(Actor, self).__init__() + + # 共享特征提取层 + self.l1 = nn.Linear(state_dim, 256) + self.l2 = nn.Linear(256, 256) + + # 均值输出层 + self.mean_linear = nn.Linear(256, action_dim) + # 对数标准差输出层（预测 log_std 比直接预测 std 更好优化） + self.log_std_linear = nn.Linear(256, action_dim) + + # 动作的最大物理边界（比如 Pendulum 的力矩最大是 2.0） + self.max_action = max_action + + def forward(self, state): + x = F.relu(self.l1(state)) + x = F.relu(self.l2(x)) + + mean = self.mean_linear(x) + log_std = self.log_std_linear(x) + + # 限制 log_std 的范围，防止数值不稳定 + log_std = torch.clamp(log_std, min=LOG_SIG_MIN, max=LOG_SIG_MAX) + return mean, log_std + + def sample(self, state): + mean, log_std = self.forward(state) + std = log_std.exp() + + # 构造一个高斯分布 + normal = Normal(mean, std) + + # normal.rsample() 内部执行的就是 a = mean + std * epsilon (其中 epsilon 是标准正态噪声) + # 这就是公式 (11) 的代码实现！用 rsample 才能让梯度传导回网络。 + x_t = normal.rsample() + + # 把动作压缩到 [-1, 1] 区间（这就是论文附录 C 里的 tanh 压扁函数） + y_t = torch.tanh(x_t) + + # 映射到真实的物理动作区间，比如 [-2.0, 2.0] + action = y_t * self.max_action + + # 计算这个动作的对数概率 log(pi(a|s))，用于后面算熵 + # 这行公式对应论文附录 C 的公式 (21)，是应用 tanh 后的概率修正 + log_prob = normal.log_prob(x_t) + log_prob -= torch.log(self.max_action * (1 - y_t.pow(2)) + 1e-6) + log_prob = log_prob.sum(1, keepdim=True) + + # mean 经过 tanh 就是测试时用的确定性动作 + mean = torch.tanh(mean) * self.max_action + + return action, log_prob, mean diff --git a/train.py b/train.py new file mode 100644 index 0000000..e0af75c --- /dev/null +++ b/train.py @@ -0,0 +1,86 @@ +import gymnasium as gym +import numpy as np +import torch +from algorithms.sac import SAC +from utils.replay_buffer import ReplayBuffer +import yaml +import os + +# 读取 YAML 配置文件 +config_path = os.path.join(os.path.dirname(__file__), 'configs', 'pendulum_config.yaml') +with open(config_path, "r", encoding="utf-8") as f: + config = yaml.safe_load(f) + +def main(): + # --------------------------------------------------------- + # 2. 实例化环境与获取维度信息 + # --------------------------------------------------------- + env = gym.make('Pendulum-v1') + + # 动态获取环境的维度信息，确保算法通用性 + state_dim = env.observation_space.shape[0] + action_dim = env.action_space.shape[0] + max_action = float(env.action_space.high[0]) + + print(f"环境已加载: 状态维度 {state_dim}, 动作维度 {action_dim}, 最大动作界限 {max_action}") + + # --------------------------------------------------------- + # 3. 实例化 SAC 代理与经验回放池 + # --------------------------------------------------------- + agent = SAC(state_dim, action_dim, max_action, config) + replay_buffer = ReplayBuffer(state_dim, action_dim, max_size=config['buffer_size']) + + total_steps = 0 # 记录与环境交互的总步数 + + # --------------------------------------------------------- + # 4. 训练大循环 + # --------------------------------------------------------- + for episode in range(config['max_episodes']): + # 重置环境，获取初始状态 (Gymnasium 返回 state 和 info) + state, _ = env.reset() + episode_reward = 0 + + for step in range(config['max_steps']): + # --- a. 动作选择策略 --- + # 强化学习工程技巧：在训练初期使用纯随机动作，收集高多样性的初始数据 + if total_steps < config['start_steps']: + action = env.action_space.sample() + else: + # 预热结束后，交由 SAC 的策略网络进行带噪声的采样 + action = agent.select_action(state, evaluate=False) + + # --- b. 与环境交互 --- + next_state, reward, terminated, truncated, _ = env.step(action) + + # 判断回合是否真正结束 (超时截断 truncated 不算做环境动力学意义上的 done) + done = float(terminated) + + # --- c. 存入经验池 --- + replay_buffer.add(state, action, reward, next_state, done) + + state = next_state + episode_reward += reward + total_steps += 1 + + # --- d. 核心学习逻辑 --- + # 只有当经验池里的数据量足够凑齐一个 Batch 时，才开始更新网络 + if replay_buffer.size > config['batch_size']: + agent.update(replay_buffer, config['batch_size']) + + # 如果提前倒地或撞毁，结束当前回合 + if terminated or truncated: + break + + # 打印当前回合的训练结果 + print(f"Episode: {episode+1:03d} | Total Steps: {total_steps:06d} | Reward: {episode_reward:.2f}") + + # --- 阶段性保存模型 (可选) --- + if (episode + 1) % 50 == 0: + torch.save(agent.actor.state_dict(), f"sac_actor_pendulum_ep{episode+1}.pth") + print(f"[*] 已保存第 {episode+1} 回合的模型权重。") + + env.close() + print("训练结束！") + +if __name__ == "__main__": + main() diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..1b349b5 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,110 @@ +import torch +import torch.nn.functional as F +import torch.optim as optim +from models.networks import Actor, Critic +import copy + +class SAC(object): + def __init__(self, state_dim, action_dim, max_action, config): + """ + 初始化 Soft Actor-Critic 算法 + """ + # 设备配置 (CPU 或 GPU) + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # 从 config 字典中加载超参数 + self.gamma = config.get('gamma', 0.99) # 折扣因子 + self.tau = config.get('tau', 0.005) # 目标网络软更新系数 (论文公式 9 下方) + self.alpha = config.get('alpha', 0.2) # 熵的温度参数 (控制探索的随机性) + + # 1. 实例化策略网络 (Actor) + self.actor = Actor(state_dim, action_dim, max_action).to(self.device) + self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=config.get('lr', 3e-4)) + + # 2. 实例化价值网络 (Critic) - 内部已经包含了 Q1 和 Q2 + self.critic = Critic(state_dim, action_dim).to(self.device) + self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=config.get('lr', 3e-4)) + + # 3. 实例化目标价值网络 (Target Critic) + # 用 copy.deepcopy 完美复制一份初始参数，并冻结其梯度计算 + self.critic_target = copy.deepcopy(self.critic) + for param in self.critic_target.parameters(): + param.requires_grad = False + + def select_action(self, state, evaluate=False): + """ + 与环境交互时使用的动作选择函数 + """ + # 将输入的状态转换为 PyTorch Tensor + state = torch.FloatTensor(state).unsqueeze(0).to(self.device) + + # 论文技巧：评估（测试）时使用均值动作，训练时使用采样动作 + with torch.no_grad(): + if evaluate: + _, _, action = self.actor.sample(state) # 第三个返回值是均值 + else: + action, _, _ = self.actor.sample(state) # 第一个返回值是加了噪声的采样值 + + # 转换回 numpy 数组，送给 Gym 环境执行 + return action.cpu().data.numpy().flatten() + + def update(self, replay_buffer, batch_size): + """ + 算法的核心心跳：从经验池采样并更新神经网络参数 + """ + # 从 Replay Buffer 中随机抽取一个 Batch 的数据 + state, action, reward, next_state, not_done = replay_buffer.sample(batch_size) + + # ================================================================= # + # 1. 更新 Critic # + # ================================================================= # + with torch.no_grad(): + # 拿到下一个状态的动作和其对应的对数概率 (用于计算熵) + next_action, next_log_prob, _ = self.actor.sample(next_state) + + # 使用目标网络计算下一个状态的 Q 值 (Q1 和 Q2) + target_Q1, target_Q2 = self.critic_target(next_state, next_action) + + # 【核心对抗高估】：取两个 Q 值的最小值 + target_Q = torch.min(target_Q1, target_Q2) + + # 【软贝尔曼方程】：目标 Q 值 = 奖励 + gamma * (目标 Q - alpha * 熵) + # 注意这里加上了 -self.alpha * next_log_prob，这就是论文中“把熵当做奖励”的体现 + target_Q = reward + not_done * self.gamma * (target_Q - self.alpha * next_log_prob) + + # 获取当前状态和动作对应的 Q 值预测 + current_Q1, current_Q2 = self.critic(state, action) + + # 计算 Critic 的损失 (均方误差 MSE) + critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q) + + # 优化 Critic 网络 + self.critic_optimizer.zero_grad() + critic_loss.backward() + self.critic_optimizer.step() + + # ================================================================= # + # 2. 更新 Actor # + # ================================================================= # + # 让当前 Actor 对**当前状态**重新采样一个动作 (注意：不能用 Buffer 里的旧动作) + pi_action, log_prob, _ = self.actor.sample(state) + + # 拿到更新后的 Critic 对这个新动作的打分 + q1_pi, q2_pi = self.critic(state, pi_action) + min_q_pi = torch.min(q1_pi, q2_pi) + + # 计算 Actor 的损失：最小化 (alpha * log_prob - min_Q) + # 等价于最大化 (min_Q - alpha * log_prob) -> 既要 Q 值大，又要熵大（分布广） + actor_loss = (self.alpha * log_prob - min_q_pi).mean() + + # 优化 Actor 网络 + self.actor_optimizer.zero_grad() + actor_loss.backward() + self.actor_optimizer.step() + + # ================================================================= # + # 3. 软更新 Target Critic # + # ================================================================= # + # 使用 EMA (指数移动平均) 缓慢将前线 Critic 的参数移交给 Target Critic + for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): + target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) diff --git a/utils/logger.py b/utils/logger.py new file mode 100644 index 0000000..d5f7b29 --- /dev/null +++ b/utils/logger.py @@ -0,0 +1,62 @@ +# ============================================================================== +# Author: Hongru Liu +# Affiliation: School of Power and Energy, Northwestern Polytechnical University +# Version: 1.0 +# Contact: hongruliu@mail.nwpu.edu.cn +# ============================================================================== + +import matplotlib.pyplot as plt +import numpy as np +import os + +class Logger(object): + def __init__(self): + """ + 初始化日志记录器，用于暂存训练过程中的各项指标 + """ + self.episode_rewards = [] + + def record(self, reward): + """ + 记录每个 Episode 的总奖励 + """ + self.episode_rewards.append(reward) + + def plot_learning_curve(self, save_dir="."): + """ + 绘制并保存学习曲线 (Learning Curve) + """ + # 确保保存目录存在 + os.makedirs(save_dir, exist_ok=True) + + # 创建画布，严格设置白色背景 + fig, ax = plt.subplots(figsize=(10, 6), facecolor='white') + ax.set_facecolor('white') + + # 绘制奖励曲线，使用加粗线条以满足论文发表的视觉要求 + ax.plot(self.episode_rewards, linewidth=2.5, color='#1f77b4', label='Episode Reward') + + # 计算并绘制 10 个 Episode 的滑动平均线，让趋势更清晰 + if len(self.episode_rewards) >= 10: + moving_avg = np.convolve(self.episode_rewards, np.ones(10)/10, mode='valid') + ax.plot(range(9, len(self.episode_rewards)), moving_avg, + linewidth=2.5, color='#ff7f0e', label='10-Episode Moving Average') + + # 设置全英文的坐标轴标签和图例，调整字体大小 + ax.set_xlabel('Episodes', fontsize=14, fontweight='bold') + ax.set_ylabel('Total Reward', fontsize=14, fontweight='bold') + ax.set_title('Training Learning Curve (Pendulum-v1)', fontsize=16, fontweight='bold') + + # 设置刻度字体大小 + ax.tick_params(axis='both', which='major', labelsize=12) + + # 增加网格线并设置图例 + ax.grid(True, linestyle='--', alpha=0.7) + ax.legend(fontsize=12, loc='lower right') + + # 紧凑布局并保存 + plt.tight_layout() + save_path = os.path.join(save_dir, "learning_curve.png") + plt.savefig(save_path, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none') + plt.close() + print(f"[*] Learning curve saved to: {save_path}") diff --git a/utils/replay_buffer.py b/utils/replay_buffer.py new file mode 100644 index 0000000..1fe80d7 --- /dev/null +++ b/utils/replay_buffer.py @@ -0,0 +1,59 @@ +import numpy as np +import torch + +class ReplayBuffer(object): + def __init__(self, state_dim, action_dim, max_size=int(1e6)): + """ + 初始化经验回放池 + 使用预分配的 Numpy 数组来提升存储和采样效率 + """ + self.max_size = max_size + self.ptr = 0 # 当前写入的指针位置 + self.size = 0 # 当前池子里的有效数据量 + + # 预先分配内存，避免动态扩张带来性能开销 + self.state = np.zeros((max_size, state_dim), dtype=np.float32) + self.action = np.zeros((max_size, action_dim), dtype=np.float32) + self.reward = np.zeros((max_size, 1), dtype=np.float32) + self.next_state = np.zeros((max_size, state_dim), dtype=np.float32) + + # 记录该状态是否是回合的结束 (1.0 表示结束，0.0 表示未结束) + self.not_done = np.zeros((max_size, 1), dtype=np.float32) + + # 自动检测 GPU + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def add(self, state, action, reward, next_state, done): + """ + 向回放池中添加一条新的转移数据 (Transition) + """ + # 将数据写入指针当前所在的位置 + self.state[self.ptr] = state + self.action[self.ptr] = action + self.reward[self.ptr] = reward + self.next_state[self.ptr] = next_state + + # 我们存储 1 - done，这样在贝尔曼方程更新时直接相乘即可： + # Q_target = r + gamma * V * not_done + self.not_done[self.ptr] = 1. - done + + # 移动指针，如果达到了最大容量，就回到开头覆盖最老的数据（环形结构） + self.ptr = (self.ptr + 1) % self.max_size + # 更新当前有效数据量 + self.size = min(self.size + 1, self.max_size) + + def sample(self, batch_size): + """ + 随机采样一个 batch 的数据，并直接转换为 PyTorch Tensor 放到 GPU/CPU 上 + """ + # 在 0 到当前有效数据量之间，随机生成 batch_size 个索引 + ind = np.random.randint(0, self.size, size=batch_size) + + # 提取数据并转为 Tensor + return ( + torch.FloatTensor(self.state[ind]).to(self.device), + torch.FloatTensor(self.action[ind]).to(self.device), + torch.FloatTensor(self.reward[ind]).to(self.device), + torch.FloatTensor(self.next_state[ind]).to(self.device), + torch.FloatTensor(self.not_done[ind]).to(self.device) + )