Annonated Algorithm Visualization

View code on GitHub

PyTorch tutorial of Proximal Policy Optimization (PPO) algorithm for continuous action.
Related Link

PPO is one of the most popular policy gradient methods for deep reinforcement learning. It combines the classic Actor-Critic paradigm and the trust region policy optimization method into a simple yet effect algorithm design. Compared to some traditional RL algorithms like REINFORCE and A2C, PPO can deploy more stable and efficient policy optimization by using clipped surrogate objective mentioned below:

J(\theta) = \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}))

The final objective is a lower bound (i.e., a pessimistic bound) on the unclipped objective, which only ignore the change in probability ratio when it would make the objective improve, and we include it when it makes the objective worse.
Detailed notation definition can be found in Related Link.

Continuous action space, one of the most commonly used action spaces, is often used in practical decision applications such as robot manipulation and drone control. It contains serveral controllable continuous actions and RL agent needs to output proper and accuracy values every execution. Continuous action space is often directly predicted or modelled by gaussian distribution (regression problem).

This tutorial is mainly composed of the following three parts, you can learn from these demo codes step by step or using them as code segment in your own program:
- Policy Network Architecture
- Sample Action Function
- Main (Test) Function
More visulization results about PPO in continuous action space can be found in Related Link.

Overview
The definition of continuous action policy network used in PPO, which is mainly composed of three parts: encoder, mu and log_sigma.


x
 
1
from typing import Dict
2
import torch
3
import torch.nn as nn
4
from torch.distributions import Normal, Independent
5
6
7
class ContinuousPolicyNetwork(nn.Module):
8
    def __init__(self, obs_shape: int, action_shape: int) -> None:

PyTorch necessary requirements for extending nn.Module . Our network should also subclass this class.


xxxxxxxxxx
9
 
9
        super(ContinuousPolicyNetwork, self).__init__()

Define encoder module, which maps raw state into embedding vector.
It could be different for various state, such as Convolution Neural Network (CNN) for image state and Multilayer perceptron (MLP) for vector state, respectively.
Here we use two-layer MLP for vector state.
$y = max(W_2 max(W_1x+b_1, 0) + b_2, 0)$


xxxxxxxxxx
15
 
10
        self.encoder = nn.Sequential(
11
            nn.Linear(obs_shape, 16),
12
            nn.ReLU(),
13
            nn.Linear(16, 32),
14
            nn.ReLU(),
15
        )

Define mu module, which is a FC and outputs the argument mu for gaussian distribution.
$\mu = Wx + b$


xxxxxxxxxx
16
 
16
        self.mu = nn.Linear(32, action_shape)

Define log_sigma module, which is a learnable parameter but independent to state.
Here we set it as log_sigma for the convenience of optimization and usage. You can also adjust its initial value for your demands.
$\sigma = e^w$


xxxxxxxxxx
18
 
17
        self.log_sigma = nn.Parameter(torch.zeros(1, action_shape))
18

Overview
The computation graph of continuous action policy network used in PPO.
x -> encoder -> mu -> \mu .
log_sigma -> exp -> sigma .


xxxxxxxxxx
19
 
19
    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:

Transform original state into embedding vector, i.e. $(B, *) -> (B, N)$


xxxxxxxxxx
20
 
20
        x = self.encoder(x)

Output the argument mu depending on the embedding vector, i.e. $(B, N) -> (B, A)$


xxxxxxxxxx
21
 
21
        mu = self.mu(x)

Utilize broadcast mechanism to make the same shape between log_sigma and mu.
zeros_like operation doesn't pass gradient.
Related Link


xxxxxxxxxx
22
 
22
        log_sigma = self.log_sigma + torch.zeros_like(mu)

Utilize exponential operation to produce the actual sigma.
$\sigma = e^w$


xxxxxxxxxx
26
 
23
        sigma = torch.exp(log_sigma)
24
        return {'mu': mu, 'sigma': sigma}
25
26

Overview
The function of sampling continuous action, input is a dict with two keys mu and sigma ,
both of them has shape = (B, action_shape), output shape = (B, action_shape).
In this example, the distributions shapes are:
batch_shape = (B, ), event_shape = (action_shape, ), sample_shape = ().


xxxxxxxxxx
27
 
27
def sample_continuous_action(logit: Dict[str, torch.Tensor]) -> torch.Tensor:

Construct gaussian distribution, i.e.
$X \sim \mathcal{N}(\mu,\,\sigma^{2})$
Its probability density function is: $f(x) = \frac{1}{\sigma\sqrt{2\pi}} \exp\left( -\frac{1}{2}\left(\frac{x-\mu}{\sigma}\right)^{\!2}\,\right)$
Related Link


xxxxxxxxxx
28
 
28
    dist = Normal(logit['mu'], logit['sigma'])

Reinterpret action_shape gaussian distribution into a multivariate gaussian distribution with diagonal convariance matrix.
Ensure each event is independent with each other.
Related Link


xxxxxxxxxx
29
 
29
    dist = Independent(dist, 1)

Sample one action of the shape action_shape per sample (state input) and return it.


xxxxxxxxxx
32
 
30
    return dist.sample()
31
32

Overview
The function of testing sampling continuous action. Construct a standard continuous action
policy and sample a group of action.


xxxxxxxxxx
33
 
33
def test_sample_continuous_action():

Set batch_size = 4, obs_shape = 10, action_shape = 6.
action_shape is different from discrete and continuous action. The former is the possible
choice of a discrete action while the latter is the dimension of continuous action.


xxxxxxxxxx
34
 
34
    B, obs_shape, action_shape = 4, 10, 6

Generate state data from uniform distribution in [0, 1].


xxxxxxxxxx
35
 
35
    state = torch.rand(B, obs_shape)

Define continuous action network (which is similar to reparameterization) with encoder, mu and log_sigma.


xxxxxxxxxx
36
 
36
    policy_network = ContinuousPolicyNetwork(obs_shape, action_shape)

Policy network forward procedure, input state and output dict-type logit.
$\mu, \sigma = \pi(a|s)$


xxxxxxxxxx
40
 
37
    logit = policy_network(state)
38
    assert isinstance(logit, dict)
39
    assert logit['mu'].shape == (B, action_shape)
40
    assert logit['sigma'].shape == (B, action_shape)

Sample action accoding to corresponding logit (i.e., mu and sigma).


xxxxxxxxxx
44
 
41
    action = sample_continuous_action(logit)
42
    assert action.shape == (B, action_shape)
43
44

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