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