请写一段连续动作空间的强化学习代码

在强化学习中，处理连续动作空间通常涉及使用基于梯度的策略优化算法，如Policy Gradient。这里是一个简化版的Python代码示例，使用了PyTorch库，假设我们有一个环境模型`Env`，它返回状态、奖励和是否终止：

import torch
from torch.distributions import Normal

class Policy(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_size):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_size)
        self.fc_mu = nn.Linear(hidden_size, action_dim)
        self.fc_std = nn.Linear(hidden_size, action_dim)

    # 这里假设action_dim是一维连续值
    def forward(self, state):
        x = F.relu(self.fc1(state))
        mu = self.fc_mu(x)
        std = torch.exp(self.fc_std(x))  # 对标准差取exp确保正数
        return Normal(mu, std)  # 创建一个正态分布，表示我们的策略

def select_action(state, policy_net, device):
    state = torch.FloatTensor(state).to(device)
    distribution = policy_net(state)
    action = distribution.sample()  # 根据策略采样行动
    log_prob = distribution.log_prob(action)  # 记录log概率以便计算优势函数
    return action.item(), log_prob.item()

# 强化学习训练循环
policy_net = Policy(state_dim, action_dim, hidden_size).to(device)
optimizer = torch.optim.Adam(policy_net.parameters())

for episode in range(num_episodes):
    state = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        action, log_prob = select_action(state, policy_net, device)
        next_state, reward, done, _ = env.step(action)
        
        # 更新网络并存储经验
        # (在这里添加更新步骤，比如利用REINFORCE算法）
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_reward += reward
        state = next_state