其他一样.

1
class PolicyNet(torch.nn.Module):
2
    def __init__(self, state_dim, hidden_dim, action_dim):
3
        super(PolicyNet, self).__init__()
4
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
5
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)
6

7
    # 输入当前状态，输出各个可选动作的概率分布, 例如 [0.3, 0.7]
8
    def forward(self, x):
9
        x = F.relu(self.fc1(x))
10
        return F.softmax(self.fc2(x), dim=1)
11

12
class REINFORCE:
13
    def __init__(self, state_dim, hidden_dim, action_dim, learning_rate, gamma,
14
                 device):
15
        self.policy_net = PolicyNet(state_dim, hidden_dim,
82 collapsed lines
16
                                    action_dim).to(device)
17
        ...
18

19
    def take_action(self, state):  # 根据动作概率分布随机采样
20
        state = torch.tensor([state], dtype=torch.float).to(self.device)
21
        probs = self.policy_net(state)
22
        action_dist = torch.distributions.Categorical(probs)
23
        action = action_dist.sample() # 随机选一个
24
        return action.item()
25

26
    def update(self, transition_dict):
27
        # 一个 list 即是开始到 done 所有步骤
28
        reward_list = transition_dict['rewards']
29
        state_list = transition_dict['states']
30
        action_list = transition_dict['actions']
31

32
        G = 0
33
        self.optimizer.zero_grad()
34
        for i in reversed(range(len(reward_list))):  # 从最后一步算起
35
            reward = reward_list[i]
36
            state = torch.tensor([state_list[i]],
37
                                 dtype=torch.float).to(self.device)
38
            action = torch.tensor([action_list[i]]).view(-1, 1).to(self.device)
39
            # 这里仅仅是提取rollout过程中采纳那个 action 的概率.
40
            # 这个概率也是最终 loss 中唯一的梯度来源. 其他 action 的可能回报在这里不产生影响.
41
            log_prob = torch.log(self.policy_net(state).gather(1, action))
42
            G = self.gamma * G + reward
43
            loss = -log_prob * G  # 每一步的损失函数
44
            loss.backward()  # 反向传播计算梯度
45
        self.optimizer.step()  # 梯度下降
46

47
# 其他一样.
48
learning_rate = 1e-3
49
num_episodes = 1000
50
hidden_dim = 128
51
gamma = 0.98
52
device = torch.device("cuda") if torch.cuda.is_available() else torch.device(
53
    "cpu")
54

55
env_name = "CartPole-v0"
56
env = gym.make(env_name)
57
env.seed(0)
58
torch.manual_seed(0)
59
state_dim = env.observation_space.shape[0]
60
action_dim = env.action_space.n
61
agent = REINFORCE(state_dim, hidden_dim, action_dim, learning_rate, gamma,
62
                  device)
63

64
return_list = []
65
for i in range(10):
66
    with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:
67
        for i_episode in range(int(num_episodes / 10)):
68
            episode_return = 0
69
            transition_dict = {
70
                'states': [],
71
                'actions': [],
72
                'next_states': [],
73
                'rewards': [],
74
                'dones': []
75
            }
76
            state = env.reset()
77
            done = False
78
            while not done:
79
                action = agent.take_action(state)
80
                nexth_state, reward, done, _ = env.step(action)
81
                transition_dict['states'].append(state)
82
                transition_dict['actions'].append(action)
83
                transition_dict['next_states'].append(next_state)
84
                transition_dict['rewards'].append(reward)
85
                transition_dict['dones'].append(done)
86
                state = next_state
87
                episode_return += reward
88
            return_list.append(episode_return)
89
            agent.update(transition_dict)
90
            if (i_episode + 1) % 10 == 0:
91
                pbar.set_postfix({
92
                    'episode':
93
                    '%d' % (num_episodes / 10 * i + i_episode + 1),
94
                    'return':
95
                    '%.3f' % np.mean(return_list[-10:])
96
                })
97
            pbar.update(1)