first commit
This commit is contained in:
+42
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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 # 是否保存最终模型权重
|
||||||
@@ -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")
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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)
|
||||||
@@ -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}")
|
||||||
@@ -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)
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user