Annonated Algorithm Visualization

View code on GitHub

PyTorch implementation of Advantage Actor-Critic (A2C)
REINFORCE method usually suffers from high variance for gradient estimation and Actor-Critic method can only get a biased gradient estimation.
To combine these two methods, A2C uses a baseline function for normalization. By subtracting the baseline function to the total return, the variance for gradient estimation is reduced.
In practical, the baseline function is set to be the value function. The final target function is formulated as:

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

Also in this way, the estimation is guaranteed to be unbiased.
Supplementary material for explaining why baseline function can reduce variance: Related Link
This document mainly includes:
- Implementation of A2C error.
- Main function (test function)

Overview
Implementation of A2C (Advantage Actor-Critic) Related Link


x
 
1
from collections import namedtuple
2
import torch
3
import torch.nn.functional as F
4
5
a2c_data = namedtuple('a2c_data', ['logit', 'action', 'value', 'adv', 'return_', 'weight'])
6
a2c_loss = namedtuple('a2c_loss', ['policy_loss', 'value_loss', 'entropy_loss'])
7
8
9
def a2c_error(data: namedtuple) -> namedtuple:

Unpack data: $<\pi(a|s), a, V(s), A^{\pi}(s, a), G_t, w>$


xxxxxxxxxx
10
 
10
    logit, action, value, adv, return_, weight = data

Prepare weight for default cases.


xxxxxxxxxx
12
 
11
    if weight is None:
12
        weight = torch.ones_like(value)

Prepare policy distribution from logit and get log propability.


xxxxxxxxxx
14
 
13
    dist = torch.distributions.categorical.Categorical(logits=logit)
14
    logp = dist.log_prob(action)

Policy loss: $- \frac 1 N \sum_{n=1}^{N} log(\pi(a^n|s^n)) A^{\pi}(s^n, a^n)$


xxxxxxxxxx
15
 
15
    policy_loss = -(logp * adv * weight).mean()

Value loss: $\frac 1 N \sum_{n=1}^{N} (G_t^n - V(s^n))^2$


xxxxxxxxxx
16
 
16
    value_loss = (F.mse_loss(return_, value, reduction='none') * weight).mean()

Entropy bonus: $\frac 1 N \sum_{n=1}^{N} \sum_{a^n}\pi(a^n|s^n) log(\pi(a^n|s^n))$
P.S. the final loss is policy_loss + value_weight * value_loss - entropy_weight * entropy_loss .


xxxxxxxxxx
17
 
17
    entropy_loss = (dist.entropy() * weight).mean()

Return the concrete loss items.


xxxxxxxxxx
20
 
18
    return a2c_loss(policy_loss, value_loss, entropy_loss)
19
20

Overview
Test function of A2C, for both forward and backward operations.


xxxxxxxxxx
21
 
21
def test_a2c():

batch size=4, action=32


xxxxxxxxxx
22
 
22
    B, N = 4, 32

Generate logit, action, value, adv, return_.


xxxxxxxxxx
28
 
23
    logit = torch.randn(B, N).requires_grad_(True)
24
    action = torch.randint(0, N, size=(B, ))
25
    value = torch.randn(B).requires_grad_(True)
26
    adv = torch.rand(B)
27
    return_ = torch.randn(B) * 2
28
    data = a2c_data(logit, action, value, adv, return_, None)

Compute A2C error.


xxxxxxxxxx
29
 
29
    loss = a2c_error(data)

Assert the loss is differentiable.


xxxxxxxxxx
37
 
30
    assert logit.grad is None
31
    assert value.grad is None
32
    total_loss = sum(loss)
33
    total_loss.backward()
34
    assert isinstance(logit.grad, torch.Tensor)
35
    assert isinstance(value.grad, torch.Tensor)
36
37

If you have any questions or advices about this documation, you can raise issues in GitHub (https://github.com/opendilab/PPOxFamily) or email us (opendilab@pjlab.org.cn).