111 lines
5.4 KiB
Python
111 lines
5.4 KiB
Python
|
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)
|