10-actor-critic算法

1
class ActorCritic:
2
    # self.critic = ValueNet(state_dim, hidden_dim).to(device)  # 价值网络
3
    ...
4
    def update(self, transition_dict):
5
        states = torch.tensor(transition_dict['states'],
6
                              dtype=torch.float).to(self.device)
7
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(
8
            self.device)
9
        rewards = torch.tensor(transition_dict['rewards'],
10
                               dtype=torch.float).view(-1, 1).to(self.device)
11
        next_states = torch.tensor(transition_dict['next_states'],
12
                                   dtype=torch.float).to(self.device)
13
        dones = torch.tensor(transition_dict['dones'],
14
                             dtype=torch.float).view(-1, 1).to(self.device)
15

40 collapsed lines
16
        # 时序差分目标
17
        td_target = rewards + self.gamma * self.critic(next_states) * (1 -
18
                                                                       dones)
19
        # 时序差分差. 即：当前动作带来的额外奖励期望（在 critic 精准的情况下）
20
        td_delta = td_target - self.critic(states)
21
        log_probs = torch.log(self.actor(states).gather(1, actions))
22
        actor_loss = torch.mean(-log_probs * td_delta.detach())
23
        # 均方误差损失函数，这里直接 detach() 来实现类似 Double DQN 的效果... （不演了是吧）
24
        critic_loss = torch.mean(
25
            F.mse_loss(self.critic(states), td_target.detach()))
26
        self.actor_optimizer.zero_grad()
27
        self.critic_optimizer.zero_grad()
28
        actor_loss.backward()  # 计算策略网络的梯度
29
        critic_loss.backward()  # 计算价值网络的梯度
30
        self.actor_optimizer.step()  # 更新策略网络的参数
31
        self.critic_optimizer.step()  # 更新价值网络的参数
32

33
class PolicyNet(torch.nn.Module):
34
    def __init__(self, state_dim, hidden_dim, action_dim):
35
        super(PolicyNet, self).__init__()
36
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
37
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)
38

39
    # 输入状态 states: [batch_size, state_dim]
40
    # 输出动作概率分布
41
    def forward(self, x):
42
        x = F.relu(self.fc1(x))
43
        return F.softmax(self.fc2(x), dim=1)
44

45
class ValueNet(torch.nn.Module):
46
    def __init__(self, state_dim, hidden_dim):
47
        super(ValueNet, self).__init__()
48
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
49
        self.fc2 = torch.nn.Linear(hidden_dim, 1)
50

51
    # 输入状态 states: [batch_size, state_dim]
52
    # 输出状态价值 V(s): [batch_size, 1]
53
    def forward(self, x):
54
        x = F.relu(self.fc1(x))
55
        return self.fc2(x)