Annonated Algorithm Visualization

View code on GitHub

PyTorch implementation of Proximal Policy Optimization (PPO)
PPO is one of the most common algorithms in reinforcement learning, which combines Actor-Critic methods and Trust Region Policy Optimization.
For the policy part, PPO combines clipped optimization target and pessimistic bound to update policy. For the value function part, PPO usually uses classical temporal difference methods (such as GAE).
This final target function is formulated as:

\min(\frac{\pi_{\theta}(a_{t}|s_{t})}{\pi_{\theta_k}(a_{t}|s_{t})}A^{\theta_k}(s_{t},a_{t}),\text{clip}(\frac{\pi_{\theta}(a_{t}|s_{t})}{\pi_{\theta_k}(a_{t}|s_{t})}, 1-\epsilon,1+\epsilon)A^{\theta_k}(s_{t},a_{t}))

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

Overview
Implementation of Proximal Policy Optimization (PPO) Related Link with entropy bonus, value_clip and dual_clip.


x
 
1
from typing import Optional, Tuple
2
from collections import namedtuple
3
import torch
4
import numpy as np
5
6
7
ppo_policy_data = namedtuple('ppo_policy_data', ['logit_new', 'logit_old', 'action', 'adv', 'weight'])
8
ppo_policy_loss = namedtuple('ppo_policy_loss', ['policy_loss', 'entropy_loss'])
9
ppo_info = namedtuple('ppo_info', ['approx_kl', 'clipfrac'])
10
11
12
def ppo_policy_error(data: namedtuple,
13
                     clip_ratio: float = 0.2,
14
                     dual_clip: Optional[float] = None) -> Tuple[namedtuple, namedtuple]:

Unpack data: $<\pi_{new}(a|s), \pi_{old}(a|s), a, A^{\pi_{old}}(s, a), w>$


xxxxxxxxxx
15
 
15
    logit_new, logit_old, action, adv, weight = data

Prepare weight for default cases.


xxxxxxxxxx
17
 
16
    if weight is None:
17
        weight = torch.ones_like(adv)

Prepare policy distribution from logit and get log propability.


xxxxxxxxxx
21
 
18
    dist_new = torch.distributions.categorical.Categorical(logits=logit_new)
19
    dist_old = torch.distributions.categorical.Categorical(logits=logit_old)
20
    logp_new = dist_new.log_prob(action)
21
    logp_old = dist_old.log_prob(action)

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


xxxxxxxxxx
23
 
22
    dist_new_entropy = dist_new.entropy()
23
    entropy_loss = (dist_new_entropy * weight).mean()

Importance sampling weight: $r(\theta) = \frac{\pi_{new}(a|s)}{\pi_{old}(a|s)}$


xxxxxxxxxx
24
 
24
    ratio = torch.exp(logp_new - logp_old)

Original surrogate objective: $r(\theta) A^{\pi_{old}}(s, a)$


xxxxxxxxxx
25
 
25
    surr1 = ratio * adv

Clipped surrogate objective: $clip(r(\theta), 1-\epsilon, 1+\epsilon) A^{\pi_{old}}(s, a)$


xxxxxxxxxx
26
 
26
    surr2 = ratio.clamp(1 - clip_ratio, 1 + clip_ratio) * adv

Dual clip proposed by Related Link .
Only use dual_clip when adv < 0.


xxxxxxxxxx
30
 
27
    if dual_clip is not None:
28
        clip1 = torch.min(surr1, surr2)
29
        clip2 = torch.max(clip1, dual_clip * adv)
30
        policy_loss = -(torch.where(adv < 0, clip2, clip1) * weight).mean()

PPO-Clipped Loss: $min(r(\theta) A^{\pi_{old}}(s, a), clip(r(\theta), 1-\epsilon, 1+\epsilon) A^{\pi_{old}}(s, a))$
Multiply sample-wise weight and reduce mean in batch dimension.


xxxxxxxxxx
32
 
31
    else:
32
        policy_loss = (-torch.min(surr1, surr2) * weight).mean()

Add some visualization metrics to monitor optimization status.


xxxxxxxxxx
36
 
33
    with torch.no_grad():
34
        approx_kl = (logp_old - logp_new).mean().item()
35
        clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio)
36
        clipfrac = torch.as_tensor(clipped).float().mean().item()

Return final loss items and information.


xxxxxxxxxx
39
 
37
    return ppo_policy_loss(policy_loss, entropy_loss), ppo_info(approx_kl, clipfrac)
38
39

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


xxxxxxxxxx
40
 
40
def test_ppo(clip_ratio, dual_clip):

batch size=4, action=32


xxxxxxxxxx
41
 
41
    B, N = 4, 32

Generate logit_new, logit_old, action, adv.


xxxxxxxxxx
46
 
42
    logit_new = torch.randn(B, N).requires_grad_(True)
43
    logit_old = logit_new + torch.rand_like(logit_new) * 0.1
44
    action = torch.randint(0, N, size=(B, ))
45
    adv = torch.rand(B)
46
    data = ppo_policy_data(logit_new, logit_old, action, adv, None)

Compute PPO error.


xxxxxxxxxx
47
 
47
    loss, info = ppo_policy_error(data, clip_ratio=clip_ratio, dual_clip=dual_clip)

Assert the loss is differentiable.


xxxxxxxxxx
54
 
48
    assert all([np.isscalar(i) for i in info])
49
    assert logit_new.grad is None
50
    total_loss = sum(loss)
51
    total_loss.backward()
52
    assert isinstance(logit_new.grad, torch.Tensor)
53
54

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