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 实现。

from collections import namedtuple
import torch

pg_data = namedtuple('pg_data', ['logit', 'action', 'return_'])
pg_loss = namedtuple('pg_loss', ['policy_loss', 'entropy_loss'])


def pg_error(data: namedtuple) -> namedtuple:

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

    logit, action, return_ = data

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

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

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

    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 .

    entropy_loss = dist.entropy().mean()

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

    return pg_loss(policy_loss, entropy_loss)

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

def test_pg():

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

    B, N = 4, 32

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

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

计算 PG error。

    data = pg_data(logit, action, return_)
    loss = pg_error(data)

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

    assert all([l.shape == tuple() for l in loss])
    assert logit.grad is None
    total_loss = sum(loss)
    total_loss.backward()
    assert isinstance(logit.grad, torch.Tensor)

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