Annonated Algorithm Visualization

View code on GitHub

PyTorch implementation of Pop-Art algorithm for adaptive normalization techniques.
Related Link

Pop-Art is an adaptive normalization algorithm to normalized the targets used in the learning updates.
It can be used in value normalization in PPO algorithm to address multi-magnitude reward problem.

The two main components in Pop-Art are:
- ART to update scale and shift such that the return is appropriately normalized
- POP to preserve the outputs of the unnormalized function when we change the scale and shift.

Overview
The definition of Pop-Art layer, i.e., a linear layer with popart normalization, which should be
used as the last layer of a network.
For more information, you can refer to the paper Related Link


x
 
1
import pickle
2
import math
3
import torch
4
import torch.nn as nn
5
import treetensor.torch as ttorch
6
from torch.optim import AdamW
7
from torch.utils.data import DataLoader
8
9
10
class PopArt(nn.Module):

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


xxxxxxxxxx
12
 
11
        super(PopArt, self).__init__()
12

Define soft-update parameter beta.


xxxxxxxxxx
13
 
13
        self.beta = beta

Define the input and output feature dimensions of the linear layer.


xxxxxxxxxx
15
 
14
        self.input_features = input_features
15
        self.output_features = output_features

Initialize the linear layer parameters, weight and bias.


xxxxxxxxxx
17
 
16
        self.weight = nn.Parameter(torch.Tensor(output_features, input_features))
17
        self.bias = nn.Parameter(torch.Tensor(output_features))

Register a buffer for normalization parameters which can not be considered as model parameters.
Therefore, the tensor registered in buffer will not refer to gradient propagation but still can
be saved in state_dict.
The normalization parameters will be used later to save the target value's scale and shift.


xxxxxxxxxx
21
 
18
        self.register_buffer('mu', torch.zeros(output_features, requires_grad=False))
19
        self.register_buffer('sigma', torch.ones(output_features, requires_grad=False))
20
        self.register_buffer('v', torch.ones(output_features, requires_grad=False))
21

Reset the learned parameters, i.e., weight and bias.


xxxxxxxxxx
23
 
22
        self.reset_parameters()
23

Overview
The parameters initialization of the linear layer (i.e. weight and bias).


xxxxxxxxxx
24
 
24
    def reset_parameters(self) -> None:

In Kaiming Initialization, the mean of weights increment slowly and the std is close to 1,
which avoid the vanishing gradient problem and exploding gradient problem of deep models.
Specifically, the Kaiming Initialization funciton is as follows:
$std = \sqrt{\frac{2}{(1+a^2)\times fan\_in}}$
where a is the the negative slope of the rectifier used after this layer (0 for ReLU by default),
and fan_in is the number of input dimension.
For more kaiming intialization info, you can refer to the paper:
Related Link


xxxxxxxxxx
30
 
25
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
26
        if self.bias is not None:
27
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
28
            bound = 1 / math.sqrt(fan_in)
29
            nn.init.uniform_(self.bias, -bound, bound)
30

Overview
The computation graph of the linear layer with popart mechanism, which outputs both the output and the normalized output of the layer.


xxxxxxxxxx
31
 
31
    def forward(self, x: torch.Tensor) -> ttorch.Tensor:

Execute the linear layer computation $y=Wx+b$ , note here we use expand and broadcast to add bias.


xxxxxxxxxx
33
 
32
        normalized_output = x.mm(self.weight.t())
33
        normalized_output += self.bias.unsqueeze(0).expand_as(normalized_output)

Unnormalize the output for more convenient usage.


xxxxxxxxxx
38
 
34
        with torch.no_grad():
35
            output = normalized_output * self.sigma + self.mu
36
37
        return ttorch.as_tensor({'output': output, 'normalized_output': normalized_output})
38

Overview
The parameters update defined in Pop-Art, which outputs both the output and the normalized output of the layer.


xxxxxxxxxx
39
 
39
    def update_parameters(self, value: torch.Tensor) -> ttorch.Tensor:

Tensor device conversion of the normalization parameters.


xxxxxxxxxx
43
 
40
        self.mu = self.mu.to(value.device)
41
        self.sigma = self.sigma.to(value.device)
42
        self.v = self.v.to(value.device)
43

Store the old normalization parameters for later usage.


xxxxxxxxxx
45
 
44
        old_mu = self.mu
45
        old_std = self.sigma

Calculate the first and second moments (mean and variance) of the target value:
$\mu = \frac{G_t}{B}$
$v = \frac{G_t^2}{B}$ .


xxxxxxxxxx
47
 
46
        batch_mean = torch.mean(value, 0)
47
        batch_v = torch.mean(torch.pow(value, 2), 0)

Replace the nan value with old value for more stable training.


xxxxxxxxxx
49
 
48
        batch_mean[torch.isnan(batch_mean)] = self.mu[torch.isnan(batch_mean)]
49
        batch_v[torch.isnan(batch_v)] = self.v[torch.isnan(batch_v)]

Soft update the normalization parameters according to:
$\mu_t = (1-\beta)\mu_{t-1}+\beta G^v_t$
$v_t = (1-\beta)v_{t-1}+\beta(G^v_t)^2$ .


xxxxxxxxxx
51
 
50
        batch_mean = (1 - self.beta) * self.mu + self.beta * batch_mean
51
        batch_v = (1 - self.beta) * self.v + self.beta * batch_v

Calculate the standard deviation with the mean and variance:
$\sigma = \sqrt{v-\mu^2}$


xxxxxxxxxx
52
 
52
        batch_std = torch.sqrt(batch_v - (batch_mean ** 2))

Clip the standard deviation to reject the outlier data.


xxxxxxxxxx
53
 
53
        batch_std = torch.clamp(batch_std, min=1e-4, max=1e+6)

Replace the nan value with old value.


xxxxxxxxxx
55
 
54
        batch_std[torch.isnan(batch_std)] = self.sigma[torch.isnan(batch_std)]
55

Update the normalization parameters.


xxxxxxxxxx
58
 
56
        self.mu = batch_mean
57
        self.v = batch_v
58
        self.sigma = batch_std

Update weight and bias with mean and standard deviation to preserve unnormalised outputs:
$w'_i = \frac{\sigma_i}{\sigma'_i}w_i$
$b'_i = \frac{\sigma_i b_i + \mu_i-\mu'_i}{\sigma'_i}$


xxxxxxxxxx
61
 
59
        self.weight.data = (self.weight.t() * old_std / self.sigma).t()
60
        self.bias.data = (old_std * self.bias + old_mu - self.mu) / self.sigma
61

Return treetensor-type statistics.


xxxxxxxxxx
64
 
62
        return ttorch.as_tensor({'new_mean': batch_mean, 'new_std': batch_std})
63
64

Overview
A MLP network with popart as the final layer.
Input: observations and actions
Output: Estimated Q value
cat(obs,actions) -> encoder -> popart .


xxxxxxxxxx
67
 
65
class MLP(nn.Module):
66
67
    def __init__(self, obs_shape: int, action_shape: int) -> None:

Define the encoder and popart layer.
Here we use MLP with two layer and ReLU as activate function. The final layer is popart layer.


xxxxxxxxxx
75
 
68
        self.encoder = nn.Sequential(
69
            nn.Linear(obs_shape + action_shape, 16),
70
            nn.ReLU(),
71
            nn.Linear(16, 32),
72
            nn.ReLU(),
73
        )
74
        self.popart = PopArt(32, 1)
75

Overview
Forward computation of the MLP network with popart layer.


xxxxxxxxxx
76
 
76
    def forward(self, obs: torch.Tensor, actions: torch.Tensor) -> ttorch.Tensor:

The encoder first concatenate the observation vectors and actions,
then map the input to an embedding vector.


xxxxxxxxxx
78
 
77
        x = torch.cat((obs, actions), 1)
78
        x = self.encoder(x)

The popart layer maps the embedding vector to a normalized value.


xxxxxxxxxx
82
 
79
        x = self.popart(x)
80
        return x
81
82

Overview
Example training function for using MLP network with Pop-Art layer in fixed Q value approximation.


xxxxxxxxxx
83
 
83
def train(obs_shape: int, action_shape: int, NUM_EPOCH: int, train_data):

Define the MLP network and optimizer, and loss function.


xxxxxxxxxx
86
 
84
    model = MLP(obs_shape, action_shape)
85
    optimizer = AdamW(model.parameters(), lr=0.0001, weight_decay=0.0001)
86
    MSEloss = nn.MSELoss()

Read the preprocessed data of trained agent on lunarlander environment.
Each sample in the datasets should be a dict with following format:
$key\quad dim$
$observations\quad (*,8)$
$actions\quad (*,)$
$returns\quad (*,)$
where the returns is the discounted return from the current state.


xxxxxxxxxx
88
 
87
    train_data = DataLoader(train_data, batch_size=64, shuffle=True)
88

For loop 1: train MLP network for NUM_EPOCH epochs.


xxxxxxxxxx
90
 
89
    running_loss = 0.0
90
    for epoch in range(NUM_EPOCH):

For loop 2: Inside each epoch, split the entire dataset into mini-batches, then train on each mini-batch.


xxxxxxxxxx
91
 
91
        for idx, data in enumerate(train_data):

Compute the original output and the normalized output.


xxxxxxxxxx
94
 
92
            output = model(data['observations'], data['actions'])
93
            mu = model.popart.mu
94
            sigma = model.popart.sigma

Normalize the target return to align with the normalized Q value.


xxxxxxxxxx
96
 
95
            with torch.no_grad():
96
                normalized_return = (data['returns'] - mu) / sigma

The loss is calculated as the MSE loss between normalized Q value and normalized target return.


xxxxxxxxxx
97
 
97
            loss = MSEloss(output.normalized_output, normalized_return)

Loss backward and optimizer update step.


xxxxxxxxxx
100
 
98
            optimizer.zero_grad()
99
            loss.backward()
100
            optimizer.step()

After the model parameters are updated with the gradient,
the weights and bias should be updated to preserve unnormalised outputs.


xxxxxxxxxx
102
 
101
            model.popart.update_parameters(data['returns'])
102

Use item method to get the pure Python scalar of the loss, then add it into running_loss .


xxxxxxxxxx
104
 
103
            running_loss += loss.item()
104

Print the loss every 100 epochs.


xxxxxxxxxx
109
 
105
        if epoch % 100 == 99:
106
            print('Epoch [%d] loss: %.6f' % (epoch + 1, running_loss / 100))
107
            running_loss = 0.0
108
109