概述
Advantage Actor-Critic (A2C) 算法的 PyTorch 版实现。 Related Link
from collections import namedtuple
import torch
import torch.nn.functional as F
a2c_data = namedtuple('a2c_data', ['logit', 'action', 'value', 'adv', 'return_', 'weight'])
a2c_loss = namedtuple('a2c_loss', ['policy_loss', 'value_loss', 'entropy_loss'])
def a2c_error(data: namedtuple) -> namedtuple:
对数据 data 进行解包: $$<\pi(a|s), a, V(s), A^{\pi}(s, a), G_t, w>$$
logit, action, value, adv, return_, weight = data
准备默认的权重(weight)。
if weight is None:
weight = torch.ones_like(value)
根据 logit 构建策略分布,然后得到对应动作的概率的对数值。
dist = torch.distributions.categorical.Categorical(logits=logit)
logp = dist.log_prob(action)
策略的损失函数: $$- \frac 1 N \sum_{n=1}^{N} log(\pi(a^n|s^n)) A^{\pi}(s^n, a^n)$$
policy_loss = -(logp * adv * weight).mean()
值函数的损失函数: $$\frac 1 N \sum_{n=1}^{N} (G_t^n - V(s^n))^2$$
value_loss = (F.mse_loss(return_, value, reduction='none') * weight).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 + value_weight * value_loss - entropy_weight * entropy_loss .
entropy_loss = (dist.entropy() * weight).mean()
Return the concrete loss items.
返回最终的各项损失函数:包含策略损失,值损失和熵损失。
return a2c_loss(policy_loss, value_loss, entropy_loss)
概述
A2C 算法的测试函数,包括前向和反向传播测试
def test_a2c():
设置相关参数:batch size=4, action=32
B, N = 4, 32
从随机分布中生成测试数据: logit, action, value, adv, return_.
logit = torch.randn(B, N).requires_grad_(True)
action = torch.randint(0, N, size=(B, ))
value = torch.randn(B).requires_grad_(True)
adv = torch.rand(B)
return_ = torch.randn(B) * 2
data = a2c_data(logit, action, value, adv, return_, None)
计算 A2C error
loss = a2c_error(data)
测试 loss 是否是可微分的,是否能正确产生梯度
assert logit.grad is None
assert value.grad is None
total_loss = sum(loss)
total_loss.backward()
assert isinstance(logit.grad, torch.Tensor)
assert isinstance(value.grad, torch.Tensor)
如果读者关于本文档有任何问题和建议,可以在 GitHub 提 issue 或是直接发邮件给我们 (opendilab@pjlab.org.cn) 。