57 lines
2.1 KiB
Python
57 lines
2.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.distributions import Normal
|
|
|
|
class ValueNetwork(nn.Module):
|
|
"""
|
|
状态值函数网络 (Critic)，用于估计状态的内在价值 V(s)
|
|
增加了隐藏层容量，以确保在多轮迭代中能够更准确地拟合优势函数
|
|
"""
|
|
def __init__(self, state_dim, hidden_dim=128):
|
|
super(ValueNetwork, self).__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(state_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, 1)
|
|
)
|
|
|
|
def forward(self, state):
|
|
# 返回对当前状态的价值评估
|
|
return self.net(state)
|
|
|
|
|
|
class PolicyNetwork(nn.Module):
|
|
"""
|
|
策略网络 (Actor)，输出高斯分布的均值和标准差，适用于连续动作空间
|
|
"""
|
|
def __init__(self, state_dim, action_dim, action_bound, hidden_dim=128):
|
|
super(PolicyNetwork, self).__init__()
|
|
self.action_bound = action_bound # 环境允许的最大物理动作幅度
|
|
|
|
self.net = nn.Sequential(
|
|
nn.Linear(state_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, action_dim),
|
|
nn.Tanh() # 关键修复：强制均值输出在 [-1, 1] 之间，防止动作空间爆炸
|
|
)
|
|
# 将初始对数标准差设为 -0.5 (对应的标准差约为 0.6)
|
|
# 较小的初始方差有助于防止初期探索步子迈得太大导致系统崩溃
|
|
self.action_log_std = nn.Parameter(torch.full((1, action_dim), -0.5))
|
|
|
|
def forward(self, state):
|
|
# 计算动作均值并将其缩放到实际的物理边界内
|
|
action_mean = self.net(state) * self.action_bound
|
|
# 限制标准差的范围以提高数值稳定性
|
|
action_std = torch.exp(self.action_log_std.expand_as(action_mean))
|
|
return action_mean, action_std
|
|
|
|
def evaluate(self, state):
|
|
# 评估当前状态，返回构建好的高斯动作分布
|
|
mean, std = self.forward(state)
|
|
dist = Normal(mean, std)
|
|
return dist
|