99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
|
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
|