Annonated Algorithm Visualization

View code on GitHub

PyTorch implementation of Policy Gradient (PG)
Policy gradient (also known as REINFORCE) is a classical method for learning a policy.
Each

(s_t,a_t)

will be used to compute corresponding log probability. Then the probability is back-propagated then and calculate gradient. The gradient will be multiplied by a weight value, which is the accumulated return in this game.
The final target function is formulated as:

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

This document mainly includes:
- Implementation of PG error.
- Main function (test function)

Overview
Implementation of PG (Policy Gradient)


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:

Unpack data: $<\pi(a|s), a, G_t>$


xxxxxxxxxx
9
 
9
    logit, action, return_ = data

Prepare policy distribution from logit and get log propability.


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

Policy loss: $- \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()

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 - entropy_weight * entropy_loss .


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

Return the concrete loss items.


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

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


xxxxxxxxxx
17
 
17
def test_pg():

batch size=4, action=32


xxxxxxxxxx
18
 
18
    B, N = 4, 32

Generate 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

Compute PG error.


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

Assert the loss is differentiable.


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

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).