Annonated Algorithm Visualization

View code on GitHub

策略梯度（Policy Gradient，PG）算法的 PyTorch 实现。

策略梯度（也常指代 REINFORCE 算法）是学习策略的一类经典方法。
每个

(s_t，a_t)

将用于计算相应的对数概率，然后概率被反向传播计算得到梯度，梯度会乘以一个权重值，这个权重值通常是这局游戏中的累计收益。
最终的目标函数形式化定义为:

- \frac 1 N \sum_{n=1}^{N} log(\pi(a^n|s^n)) G_t^n

本文档主要包括:
- PG error 的实现。
- 主函数（测试函数）

概述
策略梯度（Policy Gradient，PG）算法的 PyTorch 实现。


x
 
1
from collections import namedtuple
2
import torch
3
4
pg_data = namedtuple('pg_data', ['logit', 'action', 'return_'])
5
pg_loss = namedtuple('pg_loss', ['policy_loss', 'entropy_loss'])
6
7
8
def pg_error(data: namedtuple) -> namedtuple:

对数据 data 进行解包: $<\pi(a|s), a, G_t>$


xxxxxxxxxx
9
 
9
    logit, action, return_ = data

根据 logit 构建策略分布，然后得到对应动作的概率的对数值。


xxxxxxxxxx
11
 
10
    dist = torch.distributions.categorical.Categorical(logits=logit)
11
    log_prob = dist.log_prob(action)

策略损失函数: $- \frac 1 N \sum_{n=1}^{N} log(\pi(a^n|s^n)) G_t^n$


xxxxxxxxxx
12
 
12
    policy_loss = -(log_prob * return_).mean()

熵奖赏（bonus）损失函数： $\frac 1 N \sum_{n=1}^{N} \sum_{a^n}\pi(a^n|s^n) log(\pi(a^n|s^n))$
注意：最终的损失是 policy_loss - entropy_weight * entropy_loss .


xxxxxxxxxx
13
 
13
    entropy_loss = dist.entropy().mean()

返回最终的各项损失函数：包含策略损失和熵损失。


xxxxxxxxxx
16
 
14
    return pg_loss(policy_loss, entropy_loss)
15
16

概述
策略梯度算法的测试函数，包括前向和反向传播测试。


xxxxxxxxxx
17
 
17
def test_pg():

设置相关参数：batch size=4, action=32


xxxxxxxxxx
18
 
18
    B, N = 4, 32

从随机分布中生成测试数据：logit, action, return_.


xxxxxxxxxx
21
 
19
    logit = torch.randn(B, N).requires_grad_(True)
20
    action = torch.randint(0, N, size=(B, ))
21
    return_ = torch.randn(B) * 2

计算 PG error。


xxxxxxxxxx
23
 
22
    data = pg_data(logit, action, return_)
23
    loss = pg_error(data)

测试 loss 是否是可微分的，是否能正确产生梯度


xxxxxxxxxx
30
 
24
    assert all([l.shape == tuple() for l in loss])
25
    assert logit.grad is None
26
    total_loss = sum(loss)
27
    total_loss.backward()
28
    assert isinstance(logit.grad, torch.Tensor)
29
30

如果读者关于本文档有任何问题和建议，可以在 GitHub 提 issue 或是直接发邮件给我们 (opendilab@pjlab.org.cn) 。