33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# 连续动作策略网络 (Deterministic Actor)
|
|
class ContActor(nn.Module):
|
|
def __init__(self, state_dim, action_dim, max_action):
|
|
super(ContActor, self).__init__()
|
|
self.fc1 = nn.Linear(state_dim, 128)
|
|
self.fc2 = nn.Linear(128, action_dim)
|
|
self.max_action = max_action
|
|
|
|
def forward(self, state):
|
|
x = F.relu(self.fc1(state))
|
|
# 使用 tanh 将输出限制在 [-1, 1]，然后映射到物理实际的边界 (例如 [-2.0, 2.0])
|
|
action = torch.tanh(self.fc2(x)) * self.max_action
|
|
return action
|
|
|
|
# 连续动作 Q 值网络 (用于评价 (State, Action) 对)
|
|
class ContQCritic(nn.Module):
|
|
def __init__(self, state_dim, action_dim):
|
|
super(ContQCritic, self).__init__()
|
|
# 将状态和动作拼接后输入网络
|
|
self.fc1 = nn.Linear(state_dim + action_dim, 128)
|
|
self.fc2 = nn.Linear(128, 1)
|
|
|
|
def forward(self, state, action):
|
|
# 拼接维度
|
|
xu = torch.cat([state, action], 1)
|
|
x = F.relu(self.fc1(xu))
|
|
q_value = self.fc2(x)
|
|
return q_value
|