wehub/labmlai--annotated_deep_learning_paper_implementations
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:19:01 +08:00
commit 3b90d1192f
2172 changed files with 594509 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
labml_nn/rl/__init__.py
+21
View File
@@ -0,0 +1,21 @@
"""
---
title: Reinforcement Learning Algorithms
summary: >
This is a collection of PyTorch implementations/tutorials of reinforcement learning algorithms.
It currently includes Proximal Policy Optimization, Generalized Advantage Estimation, and
Deep Q Networks.
---
# Reinforcement Learning Algorithms
* [Proximal Policy Optimization](ppo)
* [This is an experiment](ppo/experiment.html) that runs a PPO agent on Atari Breakout.
* [Generalized advantage estimation](ppo/gae.html)
* [Deep Q Networks](dqn)
* [This is an experiment](dqn/experiment.html) that runs a DQN agent on Atari Breakout.
* [Model](dqn/model.html) with dueling network
* [Prioritized Experience Replay Buffer](dqn/replay_buffer.html)
[This is the implementation for OpenAI game wrapper](game.html) using `multiprocessing`.
"""
labml_nn/rl/dqn/__init__.py
+164
View File
@@ -0,0 +1,164 @@
"""
---
title: Deep Q Networks (DQN)
summary: >
This is a PyTorch implementation/tutorial of Deep Q Networks (DQN) from paper
Playing Atari with Deep Reinforcement Learning.
This includes dueling network architecture, a prioritized replay buffer and
double-Q-network training.
---
# Deep Q Networks (DQN)
This is a [PyTorch](https://pytorch.org) implementation of paper
[Playing Atari with Deep Reinforcement Learning](https://arxiv.org/abs/1312.5602)
along with [Dueling Network](model.html), [Prioritized Replay](replay_buffer.html)
and Double Q Network.
Here is the [experiment](experiment.html) and [model](model.html) implementation.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/dqn/experiment.ipynb)
"""
from typing import Tuple
import torch
from torch import nn
from labml import tracker
from labml_nn.rl.dqn.replay_buffer import ReplayBuffer
class QFuncLoss(nn.Module):
"""
## Train the model
We want to find optimal action-value function.
\begin{align}
Q^*(s,a) &= \max_\pi \mathbb{E} \Big[
r_t + \gamma r_{t + 1} + \gamma^2 r_{t + 2} + ... | s_t = s, a_t = a, \pi
\Big]
\\
Q^*(s,a) &= \mathop{\mathbb{E}}_{s' \sim \large{\varepsilon}} \Big[
r + \gamma \max_{a'} Q^* (s', a') | s, a
\Big]
\end{align}
### Target network 🎯
In order to improve stability we use experience replay that randomly sample
from previous experience $U(D)$. We also use a Q network
with a separate set of parameters $\textcolor{orange}{\theta_i^{-}}$ to calculate the target.
$\textcolor{orange}{\theta_i^{-}}$ is updated periodically.
This is according to paper
[Human Level Control Through Deep Reinforcement Learning](https://deepmind.com/research/dqn/).
So the loss function is,
$$
\mathcal{L}_i(\theta_i) = \mathop{\mathbb{E}}_{(s,a,r,s') \sim U(D)}
\bigg[
\Big(
r + \gamma \max_{a'} Q(s', a'; \textcolor{orange}{\theta_i^{-}}) - Q(s,a;\theta_i)
\Big) ^ 2
\bigg]
$$
### Double $Q$-Learning
The max operator in the above calculation uses same network for both
selecting the best action and for evaluating the value.
That is,
$$
\max_{a'} Q(s', a'; \theta) = \textcolor{cyan}{Q}
\Big(
s', \mathop{\operatorname{argmax}}_{a'}
\textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta}); \textcolor{cyan}{\theta}
\Big)
$$
We use [double Q-learning](https://arxiv.org/abs/1509.06461), where
the $\operatorname{argmax}$ is taken from $\textcolor{cyan}{\theta_i}$ and
the value is taken from $\textcolor{orange}{\theta_i^{-}}$.
And the loss function becomes,
\begin{align}
\mathcal{L}_i(\theta_i) = \mathop{\mathbb{E}}_{(s,a,r,s') \sim U(D)}
\Bigg[
\bigg(
&r + \gamma \textcolor{orange}{Q}
\Big(
s',
\mathop{\operatorname{argmax}}_{a'}
\textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i}); \textcolor{orange}{\theta_i^{-}}
\Big)
\\
- &Q(s,a;\theta_i)
\bigg) ^ 2
\Bigg]
\end{align}
"""
def __init__(self, gamma: float):
super().__init__()
self.gamma = gamma
self.huber_loss = nn.SmoothL1Loss(reduction='none')
def forward(self, q: torch.Tensor, action: torch.Tensor, double_q: torch.Tensor,
target_q: torch.Tensor, done: torch.Tensor, reward: torch.Tensor,
weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
* `q` - $Q(s;\theta_i)$
* `action` - $a$
* `double_q` - $\textcolor{cyan}Q(s';\textcolor{cyan}{\theta_i})$
* `target_q` - $\textcolor{orange}Q(s';\textcolor{orange}{\theta_i^{-}})$
* `done` - whether the game ended after taking the action
* `reward` - $r$
* `weights` - weights of the samples from prioritized experienced replay
"""
# $Q(s,a;\theta_i)$
q_sampled_action = q.gather(-1, action.to(torch.long).unsqueeze(-1)).squeeze(-1)
tracker.add('q_sampled_action', q_sampled_action)
# Gradients shouldn't propagate gradients
# $$r + \gamma \textcolor{orange}{Q}
# \Big(s',
# \mathop{\operatorname{argmax}}_{a'}
# \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i}); \textcolor{orange}{\theta_i^{-}}
# \Big)$$
with torch.no_grad():
# Get the best action at state $s'$
# $$\mathop{\operatorname{argmax}}_{a'}
# \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i})$$
best_next_action = torch.argmax(double_q, -1)
# Get the q value from the target network for the best action at state $s'$
# $$\textcolor{orange}{Q}
# \Big(s',\mathop{\operatorname{argmax}}_{a'}
# \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i}); \textcolor{orange}{\theta_i^{-}}
# \Big)$$
best_next_q_value = target_q.gather(-1, best_next_action.unsqueeze(-1)).squeeze(-1)
# Calculate the desired Q value.
# We multiply by `(1 - done)` to zero out
# the next state Q values if the game ended.
#
# $$r + \gamma \textcolor{orange}{Q}
# \Big(s',
# \mathop{\operatorname{argmax}}_{a'}
# \textcolor{cyan}{Q}(s', a'; \textcolor{cyan}{\theta_i}); \textcolor{orange}{\theta_i^{-}}
# \Big)$$
q_update = reward + self.gamma * best_next_q_value * (1 - done)
tracker.add('q_update', q_update)
# Temporal difference error $\delta$ is used to weigh samples in replay buffer
td_error = q_sampled_action - q_update
tracker.add('td_error', td_error)
# We take [Huber loss](https://en.wikipedia.org/wiki/Huber_loss) instead of
# mean squared error loss because it is less sensitive to outliers
losses = self.huber_loss(q_sampled_action, q_update)
# Get weighted means
loss = torch.mean(weights * losses)
tracker.add('loss', loss)
return td_error, loss
labml_nn/rl/dqn/experiment.ipynb
+243
View File
@@ -0,0 +1,243 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/dqn/experiment.ipynb) \n",
"\n",
"## Deep Q Networks (DQN)\n",
"\n",
"This is an experiment training an agent to play Atari Breakout game using Deep Q Networks (DQN)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ZCzmCrAIVg0L",
"outputId": "6c416266-1e99-4e60-a665-06ff9fba22a6"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "3-G5kplRFmsO"
},
"source": [
"Add Atari ROMs (Doesn't work without this in Google Colab)"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "SByhklD1FlSj",
"outputId": "74075a5e-ec1c-43dc-8859-8f7c3b3b8402"
},
"source": [
"! wget http://www.atarimania.com/roms/Roms.rar\n",
"! mkdir /content/ROM/\n",
"! unrar e /content/Roms.rar /content/ROM/\n",
"! python -m atari_py.import_roms /content/ROM/"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"from labml import experiment\n",
"from labml.configs import FloatDynamicHyperParam\n",
"from labml_nn.rl.dqn.experiment import Trainer"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"dqn\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Hw6uVl1_GaPv"
},
"source": [
"### Configurations\n",
"\n",
"`FloatDynamicHyperParam` is a dynamic hyper-parameter\n",
"that you can change while the experiment is running."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "L8bUtLD6GksC",
"outputId": "c7d4efe7-490e-4153-e691-ca31df1e1275"
},
"source": [
"configs = {\n",
" # Number of updates\n",
" 'updates': 1_000_000,\n",
" # Number of epochs to train the model with sampled data.\n",
" 'epochs': 8,\n",
" # Number of worker processes\n",
" 'n_workers': 8,\n",
" # Number of steps to run on each process for a single update\n",
" 'worker_steps': 4,\n",
" # Mini batch size\n",
" 'mini_batch_size': 32,\n",
" # Target model updating interval\n",
" 'update_target_model': 250,\n",
" # Learning rate.\n",
" 'learning_rate': FloatDynamicHyperParam(1e-4, (0, 1e-3)),\n",
"}"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Set experiment configurations"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"experiment.configs(configs)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "qYQCFt_JYsjd"
},
"source": [
"Create trainer"
]
},
{
"cell_type": "code",
"metadata": {
"id": "8LB7XVViYuPG"
},
"source": [
"trainer = Trainer(**configs)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 520
},
"id": "aIAWo7Fw5DR8",
"outputId": "f2bca844-662d-4bfb-a295-d8529f538eaa"
},
"source": [
"with experiment.start():\n",
" trainer.run_training_loop()"
],
"outputs": [],
"execution_count": null
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Deep Q Networks (DQN)",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.5"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
labml_nn/rl/dqn/experiment.py
+290
View File
@@ -0,0 +1,290 @@
"""
---
title: DQN Experiment with Atari Breakout
summary: Implementation of DQN experiment with Atari Breakout
---
# DQN Experiment with Atari Breakout
This experiment trains a Deep Q Network (DQN) to play Atari Breakout game on OpenAI Gym.
It runs the [game environments on multiple processes](../game.html) to sample efficiently.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/dqn/experiment.ipynb)
"""
import numpy as np
import torch
from labml import tracker, experiment, logger, monit
from labml.internal.configs.dynamic_hyperparam import FloatDynamicHyperParam
from labml_nn.helpers.schedule import Piecewise
from labml_nn.rl.dqn import QFuncLoss
from labml_nn.rl.dqn.model import Model
from labml_nn.rl.dqn.replay_buffer import ReplayBuffer
from labml_nn.rl.game import Worker
# Select device
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
def obs_to_torch(obs: np.ndarray) -> torch.Tensor:
"""Scale observations from `[0, 255]` to `[0, 1]`"""
return torch.tensor(obs, dtype=torch.float32, device=device) / 255.
class Trainer:
"""
## Trainer
"""
def __init__(self, *,
updates: int, epochs: int,
n_workers: int, worker_steps: int, mini_batch_size: int,
update_target_model: int,
learning_rate: FloatDynamicHyperParam,
):
# number of workers
self.n_workers = n_workers
# steps sampled on each update
self.worker_steps = worker_steps
# number of training iterations
self.train_epochs = epochs
# number of updates
self.updates = updates
# size of mini batch for training
self.mini_batch_size = mini_batch_size
# update target network every 250 update
self.update_target_model = update_target_model
# learning rate
self.learning_rate = learning_rate
# exploration as a function of updates
self.exploration_coefficient = Piecewise(
[
(0, 1.0),
(25_000, 0.1),
(self.updates / 2, 0.01)
], outside_value=0.01)
# $\beta$ for replay buffer as a function of updates
self.prioritized_replay_beta = Piecewise(
[
(0, 0.4),
(self.updates, 1)
], outside_value=1)
# Replay buffer with $\alpha = 0.6$. Capacity of the replay buffer must be a power of 2.
self.replay_buffer = ReplayBuffer(2 ** 14, 0.6)
# Model for sampling and training
self.model = Model().to(device)
# target model to get $\textcolor{orange}Q(s';\textcolor{orange}{\theta_i^{-}})$
self.target_model = Model().to(device)
# create workers
self.workers = [Worker(47 + i) for i in range(self.n_workers)]
# initialize tensors for observations
self.obs = np.zeros((self.n_workers, 4, 84, 84), dtype=np.uint8)
# reset the workers
for worker in self.workers:
worker.child.send(("reset", None))
# get the initial observations
for i, worker in enumerate(self.workers):
self.obs[i] = worker.child.recv()
# loss function
self.loss_func = QFuncLoss(0.99)
# optimizer
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=2.5e-4)
def _sample_action(self, q_value: torch.Tensor, exploration_coefficient: float):
"""
#### $\epsilon$-greedy Sampling
When sampling actions we use a $\epsilon$-greedy strategy, where we
take a greedy action with probabiliy $1 - \epsilon$ and
take a random action with probability $\epsilon$.
We refer to $\epsilon$ as `exploration_coefficient`.
"""
# Sampling doesn't need gradients
with torch.no_grad():
# Sample the action with highest Q-value. This is the greedy action.
greedy_action = torch.argmax(q_value, dim=-1)
# Uniformly sample and action
random_action = torch.randint(q_value.shape[-1], greedy_action.shape, device=q_value.device)
# Whether to chose greedy action or the random action
is_choose_rand = torch.rand(greedy_action.shape, device=q_value.device) < exploration_coefficient
# Pick the action based on `is_choose_rand`
return torch.where(is_choose_rand, random_action, greedy_action).cpu().numpy()
def sample(self, exploration_coefficient: float):
"""### Sample data"""
# This doesn't need gradients
with torch.no_grad():
# Sample `worker_steps`
for t in range(self.worker_steps):
# Get Q_values for the current observation
q_value = self.model(obs_to_torch(self.obs))
# Sample actions
actions = self._sample_action(q_value, exploration_coefficient)
# Run sampled actions on each worker
for w, worker in enumerate(self.workers):
worker.child.send(("step", actions[w]))
# Collect information from each worker
for w, worker in enumerate(self.workers):
# Get results after executing the actions
next_obs, reward, done, info = worker.child.recv()
# Add transition to replay buffer
self.replay_buffer.add(self.obs[w], actions[w], reward, next_obs, done)
# update episode information.
# collect episode info, which is available if an episode finished;
# this includes total reward and length of the episode -
# look at `Game` to see how it works.
if info:
tracker.add('reward', info['reward'])
tracker.add('length', info['length'])
# update current observation
self.obs[w] = next_obs
def train(self, beta: float):
"""
### Train the model
"""
for _ in range(self.train_epochs):
# Sample from priority replay buffer
samples = self.replay_buffer.sample(self.mini_batch_size, beta)
# Get the predicted Q-value
q_value = self.model(obs_to_torch(samples['obs']))
# Get the Q-values of the next state for [Double Q-learning](index.html).
# Gradients shouldn't propagate for these
with torch.no_grad():
# Get $\textcolor{cyan}Q(s';\textcolor{cyan}{\theta_i})$
double_q_value = self.model(obs_to_torch(samples['next_obs']))
# Get $\textcolor{orange}Q(s';\textcolor{orange}{\theta_i^{-}})$
target_q_value = self.target_model(obs_to_torch(samples['next_obs']))
# Compute Temporal Difference (TD) errors, $\delta$, and the loss, $\mathcal{L}(\theta)$.
td_errors, loss = self.loss_func(q_value,
q_value.new_tensor(samples['action']),
double_q_value, target_q_value,
q_value.new_tensor(samples['done']),
q_value.new_tensor(samples['reward']),
q_value.new_tensor(samples['weights']))
# Calculate priorities for replay buffer $p_i = |\delta_i| + \epsilon$
new_priorities = np.abs(td_errors.cpu().numpy()) + 1e-6
# Update replay buffer priorities
self.replay_buffer.update_priorities(samples['indexes'], new_priorities)
# Set learning rate
for pg in self.optimizer.param_groups:
pg['lr'] = self.learning_rate()
# Zero out the previously calculated gradients
self.optimizer.zero_grad()
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)
# Update parameters based on gradients
self.optimizer.step()
def run_training_loop(self):
"""
### Run training loop
"""
# Last 100 episode information
tracker.set_queue('reward', 100, True)
tracker.set_queue('length', 100, True)
# Copy to target network initially
self.target_model.load_state_dict(self.model.state_dict())
for update in monit.loop(self.updates):
# $\epsilon$, exploration fraction
exploration = self.exploration_coefficient(update)
tracker.add('exploration', exploration)
# $\beta$ for prioritized replay
beta = self.prioritized_replay_beta(update)
tracker.add('beta', beta)
# Sample with current policy
self.sample(exploration)
# Start training after the buffer is full
if self.replay_buffer.is_full():
# Train the model
self.train(beta)
# Periodically update target network
if update % self.update_target_model == 0:
self.target_model.load_state_dict(self.model.state_dict())
# Save tracked indicators.
tracker.save()
# Add a new line to the screen periodically
if (update + 1) % 1_000 == 0:
logger.log()
def destroy(self):
"""
### Destroy
Stop the workers
"""
for worker in self.workers:
worker.child.send(("close", None))
def main():
# Create the experiment
experiment.create(name='dqn')
# Configurations
configs = {
# Number of updates
'updates': 1_000_000,
# Number of epochs to train the model with sampled data.
'epochs': 8,
# Number of worker processes
'n_workers': 8,
# Number of steps to run on each process for a single update
'worker_steps': 4,
# Mini batch size
'mini_batch_size': 32,
# Target model updating interval
'update_target_model': 250,
# Learning rate.
'learning_rate': FloatDynamicHyperParam(1e-4, (0, 1e-3)),
}
# Configurations
experiment.configs(configs)
# Initialize the trainer
m = Trainer(**configs)
# Run and monitor the experiment
with experiment.start():
m.run_training_loop()
# Stop the workers
m.destroy()
# ## Run it
if __name__ == "__main__":
main()
labml_nn/rl/dqn/model.py
+105
View File
@@ -0,0 +1,105 @@
"""
---
title: Deep Q Network (DQN) Model
summary: Implementation of neural network model for Deep Q Network (DQN).
---
# Deep Q Network (DQN) Model
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/dqn/experiment.ipynb)
"""
import torch
from torch import nn
class Model(nn.Module):
"""
## Dueling Network ⚔️ Model for $Q$ Values
We are using a [dueling network](https://arxiv.org/abs/1511.06581)
to calculate Q-values.
Intuition behind dueling network architecture is that in most states
the action doesn't matter,
and in some states the action is significant. Dueling network allows
this to be represented very well.
\begin{align}
Q^\pi(s,a) &= V^\pi(s) + A^\pi(s, a)
\\
\mathop{\mathbb{E}}_{a \sim \pi(s)}
\Big[
A^\pi(s, a)
\Big]
&= 0
\end{align}
So we create two networks for $V$ and $A$ and get $Q$ from them.
$$
Q(s, a) = V(s) +
\Big(
A(s, a) - \frac{1}{|\mathcal{A}|} \sum_{a' \in \mathcal{A}} A(s, a')
\Big)
$$
We share the initial layers of the $V$ and $A$ networks.
"""
def __init__(self):
super().__init__()
self.conv = nn.Sequential(
# The first convolution layer takes a
# $84\times84$ frame and produces a $20\times20$ frame
nn.Conv2d(in_channels=4, out_channels=32, kernel_size=8, stride=4),
nn.ReLU(),
# The second convolution layer takes a
# $20\times20$ frame and produces a $9\times9$ frame
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2),
nn.ReLU(),
# The third convolution layer takes a
# $9\times9$ frame and produces a $7\times7$ frame
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1),
nn.ReLU(),
)
# A fully connected layer takes the flattened
# frame from third convolution layer, and outputs
# $512$ features
self.lin = nn.Linear(in_features=7 * 7 * 64, out_features=512)
self.activation = nn.ReLU()
# This head gives the state value $V$
self.state_value = nn.Sequential(
nn.Linear(in_features=512, out_features=256),
nn.ReLU(),
nn.Linear(in_features=256, out_features=1),
)
# This head gives the action value $A$
self.action_value = nn.Sequential(
nn.Linear(in_features=512, out_features=256),
nn.ReLU(),
nn.Linear(in_features=256, out_features=4),
)
def forward(self, obs: torch.Tensor):
# Convolution
h = self.conv(obs)
# Reshape for linear layers
h = h.reshape((-1, 7 * 7 * 64))
# Linear layer
h = self.activation(self.lin(h))
# $A$
action_value = self.action_value(h)
# $V$
state_value = self.state_value(h)
# $A(s, a) - \frac{1}{|\mathcal{A}|} \sum_{a' \in \mathcal{A}} A(s, a')$
action_score_centered = action_value - action_value.mean(dim=-1, keepdim=True)
# $Q(s, a) =V(s) + \Big(A(s, a) - \frac{1}{|\mathcal{A}|} \sum_{a' \in \mathcal{A}} A(s, a')\Big)$
q = state_value + action_score_centered
return q
labml_nn/rl/dqn/readme.md
+10
View File
@@ -0,0 +1,10 @@
# [Deep Q Networks (DQN)](https://nn.labml.ai/rl/dqn/index.html)
This is a [PyTorch](https://pytorch.org) implementation of paper
[Playing Atari with Deep Reinforcement Learning](https://arxiv.org/abs/1312.5602)
along with [Dueling Network](https://nn.labml.ai/rl/dqn/model.html), [Prioritized Replay](https://nn.labml.ai/rl/dqn/replay_buffer.html)
and Double Q Network.
Here is the [experiment](https://nn.labml.ai/rl/dqn/experiment.html) and [model](https://nn.labml.ai/rl/dqn/model.html) implementation.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/dqn/experiment.ipynb)
labml_nn/rl/dqn/replay_buffer.py
+277
View File
@@ -0,0 +1,277 @@
"""
---
title: Prioritized Experience Replay Buffer
summary: Annotated implementation of prioritized experience replay using a binary segment tree.
---
# Prioritized Experience Replay Buffer
This implements paper [Prioritized experience replay](https://arxiv.org/abs/1511.05952),
using a binary segment tree.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/dqn/experiment.ipynb)
"""
import random
import numpy as np
class ReplayBuffer:
"""
## Buffer for Prioritized Experience Replay
[Prioritized experience replay](https://arxiv.org/abs/1511.05952)
samples important transitions more frequently.
The transitions are prioritized by the Temporal Difference error (td error), $\delta$.
We sample transition $i$ with probability,
$$P(i) = \frac{p_i^\alpha}{\sum_k p_k^\alpha}$$
where $\alpha$ is a hyper-parameter that determines how much
prioritization is used, with $\alpha = 0$ corresponding to uniform case.
$p_i$ is the priority.
We use proportional prioritization $p_i = |\delta_i| + \epsilon$ where
$\delta_i$ is the temporal difference for transition $i$.
We correct the bias introduced by prioritized replay using
importance-sampling (IS) weights
$$w_i = \bigg(\frac{1}{N} \frac{1}{P(i)}\bigg)^\beta$$ in the loss function.
This fully compensates when $\beta = 1$.
We normalize weights by $\frac{1}{\max_i w_i}$ for stability.
Unbiased nature is most important towards the convergence at end of training.
Therefore we increase $\beta$ towards end of training.
### Binary Segment Tree
We use a binary segment tree to efficiently calculate
$\sum_k^i p_k^\alpha$, the cumulative probability,
which is needed to sample.
We also use a binary segment tree to find $\min p_i^\alpha$,
which is needed for $\frac{1}{\max_i w_i}$.
We can also use a min-heap for this.
Binary Segment Tree lets us calculate these in $\mathcal{O}(\log n)$
time, which is way more efficient that the naive $\mathcal{O}(n)$
approach.
This is how a binary segment tree works for sum;
it is similar for minimum.
Let $x_i$ be the list of $N$ values we want to represent.
Let $b_{i,j}$ be the $j^{\mathop{th}}$ node of the $i^{\mathop{th}}$ row
in the binary tree.
That is two children of node $b_{i,j}$ are $b_{i+1,2j}$ and $b_{i+1,2j + 1}$.
The leaf nodes on row $D = \left\lceil {1 + \log_2 N} \right\rceil$
will have values of $x$.
Every node keeps the sum of the two child nodes.
That is, the root node keeps the sum of the entire array of values.
The left and right children of the root node keep
the sum of the first half of the array and
the sum of the second half of the array, respectively.
And so on...
$$b_{i,j} = \sum_{k = (j -1) * 2^{D - i} + 1}^{j * 2^{D - i}} x_k$$
Number of nodes in row $i$,
$$N_i = \left\lceil{\frac{N}{D - i + 1}} \right\rceil$$
This is equal to the sum of nodes in all rows above $i$.
So we can use a single array $a$ to store the tree, where,
$$b_{i,j} \rightarrow a_{N_i + j}$$
Then child nodes of $a_i$ are $a_{2i}$ and $a_{2i + 1}$.
That is,
$$a_i = a_{2i} + a_{2i + 1}$$
This way of maintaining binary trees is very easy to program.
*Note that we are indexing starting from 1*.
We use the same structure to compute the minimum.
"""
def __init__(self, capacity, alpha):
"""
### Initialize
"""
# We use a power of $2$ for capacity because it simplifies the code and debugging
self.capacity = capacity
# $\alpha$
self.alpha = alpha
# Maintain segment binary trees to take sum and find minimum over a range
self.priority_sum = [0 for _ in range(2 * self.capacity)]
self.priority_min = [float('inf') for _ in range(2 * self.capacity)]
# Current max priority, $p$, to be assigned to new transitions
self.max_priority = 1.
# Arrays for buffer
self.data = {
'obs': np.zeros(shape=(capacity, 4, 84, 84), dtype=np.uint8),
'action': np.zeros(shape=capacity, dtype=np.int32),
'reward': np.zeros(shape=capacity, dtype=np.float32),
'next_obs': np.zeros(shape=(capacity, 4, 84, 84), dtype=np.uint8),
'done': np.zeros(shape=capacity, dtype=np.bool)
}
# We use cyclic buffers to store data, and `next_idx` keeps the index of the next empty
# slot
self.next_idx = 0
# Size of the buffer
self.size = 0
def add(self, obs, action, reward, next_obs, done):
"""
### Add sample to queue
"""
# Get next available slot
idx = self.next_idx
# store in the queue
self.data['obs'][idx] = obs
self.data['action'][idx] = action
self.data['reward'][idx] = reward
self.data['next_obs'][idx] = next_obs
self.data['done'][idx] = done
# Increment next available slot
self.next_idx = (idx + 1) % self.capacity
# Calculate the size
self.size = min(self.capacity, self.size + 1)
# $p_i^\alpha$, new samples get `max_priority`
priority_alpha = self.max_priority ** self.alpha
# Update the two segment trees for sum and minimum
self._set_priority_min(idx, priority_alpha)
self._set_priority_sum(idx, priority_alpha)
def _set_priority_min(self, idx, priority_alpha):
"""
#### Set priority in binary segment tree for minimum
"""
# Leaf of the binary tree
idx += self.capacity
self.priority_min[idx] = priority_alpha
# Update tree, by traversing along ancestors.
# Continue until the root of the tree.
while idx >= 2:
# Get the index of the parent node
idx //= 2
# Value of the parent node is the minimum of it's two children
self.priority_min[idx] = min(self.priority_min[2 * idx], self.priority_min[2 * idx + 1])
def _set_priority_sum(self, idx, priority):
"""
#### Set priority in binary segment tree for sum
"""
# Leaf of the binary tree
idx += self.capacity
# Set the priority at the leaf
self.priority_sum[idx] = priority
# Update tree, by traversing along ancestors.
# Continue until the root of the tree.
while idx >= 2:
# Get the index of the parent node
idx //= 2
# Value of the parent node is the sum of it's two children
self.priority_sum[idx] = self.priority_sum[2 * idx] + self.priority_sum[2 * idx + 1]
def _sum(self):
"""
#### $\sum_k p_k^\alpha$
"""
# The root node keeps the sum of all values
return self.priority_sum[1]
def _min(self):
"""
#### $\min_k p_k^\alpha$
"""
# The root node keeps the minimum of all values
return self.priority_min[1]
def find_prefix_sum_idx(self, prefix_sum):
"""
#### Find largest $i$ such that $\sum_{k=1}^{i} p_k^\alpha \le P$
"""
# Start from the root
idx = 1
while idx < self.capacity:
# If the sum of the left branch is higher than required sum
if self.priority_sum[idx * 2] > prefix_sum:
# Go to left branch of the tree
idx = 2 * idx
else:
# Otherwise go to right branch and reduce the sum of left
# branch from required sum
prefix_sum -= self.priority_sum[idx * 2]
idx = 2 * idx + 1
# We are at the leaf node. Subtract the capacity by the index in the tree
# to get the index of actual value
return idx - self.capacity
def sample(self, batch_size, beta):
"""
### Sample from buffer
"""
# Initialize samples
samples = {
'weights': np.zeros(shape=batch_size, dtype=np.float32),
'indexes': np.zeros(shape=batch_size, dtype=np.int32)
}
# Get sample indexes
for i in range(batch_size):
p = random.random() * self._sum()
idx = self.find_prefix_sum_idx(p)
samples['indexes'][i] = idx
# $\min_i P(i) = \frac{\min_i p_i^\alpha}{\sum_k p_k^\alpha}$
prob_min = self._min() / self._sum()
# $\max_i w_i = \bigg(\frac{1}{N} \frac{1}{\min_i P(i)}\bigg)^\beta$
max_weight = (prob_min * self.size) ** (-beta)
for i in range(batch_size):
idx = samples['indexes'][i]
# $P(i) = \frac{p_i^\alpha}{\sum_k p_k^\alpha}$
prob = self.priority_sum[idx + self.capacity] / self._sum()
# $w_i = \bigg(\frac{1}{N} \frac{1}{P(i)}\bigg)^\beta$
weight = (prob * self.size) ** (-beta)
# Normalize by $\frac{1}{\max_i w_i}$,
# which also cancels off the $\frac{1}{N}$ term
samples['weights'][i] = weight / max_weight
# Get samples data
for k, v in self.data.items():
samples[k] = v[samples['indexes']]
return samples
def update_priorities(self, indexes, priorities):
"""
### Update priorities
"""
for idx, priority in zip(indexes, priorities):
# Set current max priority
self.max_priority = max(self.max_priority, priority)
# Calculate $p_i^\alpha$
priority_alpha = priority ** self.alpha
# Update the trees
self._set_priority_min(idx, priority_alpha)
self._set_priority_sum(idx, priority_alpha)
def is_full(self):
"""
### Whether the buffer is full
"""
return self.capacity == self.size
labml_nn/rl/game.py
+169
View File
@@ -0,0 +1,169 @@
"""
---
title: Atari wrapper with multi-processing
summary: This implements the Atari games with multi-processing.
---
# Atari wrapper with multi-processing
"""
import multiprocessing
import multiprocessing.connection
import cv2
import gym
import numpy as np
class Game:
"""
<a id="GameEnvironment"></a>
## Game environment
This is a wrapper for OpenAI gym game environment.
We do a few things here:
1. Apply the same action on four frames and get the last frame
2. Convert observation frames to gray and scale it to (84, 84)
3. Stack four frames of the last four actions
4. Add episode information (total reward for the entire episode) for monitoring
5. Restrict an episode to a single life (game has 5 lives, we reset after every single life)
#### Observation format
Observation is tensor of size (4, 84, 84). It is four frames
(images of the game screen) stacked on first axis.
i.e, each channel is a frame.
"""
def __init__(self, seed: int):
# create environment
self.env = gym.make('BreakoutNoFrameskip-v4')
self.env.seed(seed)
# tensor for a stack of 4 frames
self.obs_4 = np.zeros((4, 84, 84))
# buffer to keep the maximum of last 2 frames
self.obs_2_max = np.zeros((2, 84, 84))
# keep track of the episode rewards
self.rewards = []
# and number of lives left
self.lives = 0
def step(self, action):
"""
### Step
Executes `action` for 4 time steps and
returns a tuple of (observation, reward, done, episode_info).
* `observation`: stacked 4 frames (this frame and frames for last 3 actions)
* `reward`: total reward while the action was executed
* `done`: whether the episode finished (a life lost)
* `episode_info`: episode information if completed
"""
reward = 0.
done = None
# run for 4 steps
for i in range(4):
# execute the action in the OpenAI Gym environment
obs, r, done, info = self.env.step(action)
if i >= 2:
self.obs_2_max[i % 2] = self._process_obs(obs)
reward += r
# get number of lives left
lives = self.env.unwrapped.ale.lives()
# reset if a life is lost
if lives < self.lives:
done = True
break
# maintain rewards for each step
self.rewards.append(reward)
if done:
# if finished, set episode information if episode is over, and reset
episode_info = {"reward": sum(self.rewards), "length": len(self.rewards)}
self.reset()
else:
episode_info = None
# get the max of last two frames
obs = self.obs_2_max.max(axis=0)
# push it to the stack of 4 frames
self.obs_4 = np.roll(self.obs_4, shift=-1, axis=0)
self.obs_4[-1] = obs
return self.obs_4, reward, done, episode_info
def reset(self):
"""
### Reset environment
Clean up episode info and 4 frame stack
"""
# reset OpenAI Gym environment
obs = self.env.reset()
# reset caches
obs = self._process_obs(obs)
for i in range(4):
self.obs_4[i] = obs
self.rewards = []
self.lives = self.env.unwrapped.ale.lives()
return self.obs_4
@staticmethod
def _process_obs(obs):
"""
#### Process game frames
Convert game frames to gray and rescale to 84x84
"""
obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
return obs
def worker_process(remote: multiprocessing.connection.Connection, seed: int):
"""
##Worker Process
Each worker process runs this method
"""
# create game
game = Game(seed)
# wait for instructions from the connection and execute them
while True:
cmd, data = remote.recv()
if cmd == "step":
remote.send(game.step(data))
elif cmd == "reset":
remote.send(game.reset())
elif cmd == "close":
remote.close()
break
else:
raise NotImplementedError
class Worker:
"""
Creates a new worker and runs it in a separate process.
"""
def __init__(self, seed):
self.child, parent = multiprocessing.Pipe()
self.process = multiprocessing.Process(target=worker_process, args=(parent, seed))
self.process.start()
labml_nn/rl/ppo/__init__.py
+206
View File
@@ -0,0 +1,206 @@
"""
---
title: Proximal Policy Optimization - PPO
summary: >
An annotated implementation of Proximal Policy Optimization - PPO algorithm in PyTorch.
---
# Proximal Policy Optimization - PPO
This is a [PyTorch](https://pytorch.org) implementation of
[Proximal Policy Optimization - PPO](https://arxiv.org/abs/1707.06347).
PPO is a policy gradient method for reinforcement learning.
Simple policy gradient methods do a single gradient update per sample (or a set of samples).
Doing multiple gradient steps for a single sample causes problems
because the policy deviates too much, producing a bad policy.
PPO lets us do multiple gradient updates per sample by trying to keep the
policy close to the policy that was used to sample data.
It does so by clipping gradient flow if the updated policy
is not close to the policy used to sample the data.
You can find an experiment that uses it [here](experiment.html).
The experiment uses [Generalized Advantage Estimation](gae.html).
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb)
"""
import torch
from labml_nn.rl.ppo.gae import GAE
from torch import nn
class ClippedPPOLoss(nn.Module):
"""
## PPO Loss
Here's how the PPO update rule is derived.
We want to maximize policy reward
$$\max_\theta J(\pi_\theta) =
\mathop{\mathbb{E}}_{\tau \sim \pi_\theta}\Biggl[\sum_{t=0}^\infty \gamma^t r_t \Biggr]$$
where $r$ is the reward, $\pi$ is the policy, $\tau$ is a trajectory sampled from policy,
and $\gamma$ is the discount factor between $[0, 1]$.
\begin{align}
\mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
\sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
\Biggr] &=
\\
\mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
\sum_{t=0}^\infty \gamma^t \Bigl(
Q^{\pi_{OLD}}(s_t, a_t) - V^{\pi_{OLD}}(s_t)
\Bigr)
\Biggr] &=
\\
\mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
\sum_{t=0}^\infty \gamma^t \Bigl(
r_t + V^{\pi_{OLD}}(s_{t+1}) - V^{\pi_{OLD}}(s_t)
\Bigr)
\Biggr] &=
\\
\mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
\sum_{t=0}^\infty \gamma^t \Bigl(
r_t
\Bigr)
\Biggr]
- \mathbb{E}_{\tau \sim \pi_\theta}
\Biggl[V^{\pi_{OLD}}(s_0)\Biggr] &=
J(\pi_\theta) - J(\pi_{\theta_{OLD}})
\end{align}
So,
$$\max_\theta J(\pi_\theta) =
\max_\theta \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
\sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
\Biggr]$$
Define discounted-future state distribution,
$$d^\pi(s) = (1 - \gamma) \sum_{t=0}^\infty \gamma^t P(s_t = s | \pi)$$
Then,
\begin{align}
J(\pi_\theta) - J(\pi_{\theta_{OLD}})
&= \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
\sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
\Biggr]
\\
&= \frac{1}{1 - \gamma}
\mathbb{E}_{s \sim d^{\pi_\theta}, a \sim \pi_\theta} \Bigl[
A^{\pi_{OLD}}(s, a)
\Bigr]
\end{align}
Importance sampling $a$ from $\pi_{\theta_{OLD}}$,
\begin{align}
J(\pi_\theta) - J(\pi_{\theta_{OLD}})
&= \frac{1}{1 - \gamma}
\mathbb{E}_{s \sim d^{\pi_\theta}, a \sim \pi_\theta} \Bigl[
A^{\pi_{OLD}}(s, a)
\Bigr]
\\
&= \frac{1}{1 - \gamma}
\mathbb{E}_{s \sim d^{\pi_\theta}, a \sim \pi_{\theta_{OLD}}} \Biggl[
\frac{\pi_\theta(a|s)}{\pi_{\theta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
\Biggr]
\end{align}
Then we assume $d^\pi_\theta(s)$ and $d^\pi_{\theta_{OLD}}(s)$ are similar.
The error we introduce to $J(\pi_\theta) - J(\pi_{\theta_{OLD}})$
by this assumption is bound by the KL divergence between
$\pi_\theta$ and $\pi_{\theta_{OLD}}$.
[Constrained Policy Optimization](https://arxiv.org/abs/1705.10528)
shows the proof of this. I haven't read it.
\begin{align}
J(\pi_\theta) - J(\pi_{\theta_{OLD}})
&= \frac{1}{1 - \gamma}
\mathop{\mathbb{E}}_{s \sim d^{\pi_\theta} \atop a \sim \pi_{\theta_{OLD}}} \Biggl[
\frac{\pi_\theta(a|s)}{\pi_{\theta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
\Biggr]
\\
&\approx \frac{1}{1 - \gamma}
\mathop{\mathbb{E}}_{\textcolor{orange}{s \sim d^{\pi_{\theta_{OLD}}}}
\atop a \sim \pi_{\theta_{OLD}}} \Biggl[
\frac{\pi_\theta(a|s)}{\pi_{\theta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
\Biggr]
\\
&= \frac{1}{1 - \gamma} \mathcal{L}^{CPI}
\end{align}
"""
def __init__(self):
super().__init__()
def forward(self, log_pi: torch.Tensor, sampled_log_pi: torch.Tensor,
advantage: torch.Tensor, clip: float) -> torch.Tensor:
# ratio $r_t(\theta) = \frac{\pi_\theta (a_t|s_t)}{\pi_{\theta_{OLD}} (a_t|s_t)}$;
# *this is different from rewards* $r_t$.
ratio = torch.exp(log_pi - sampled_log_pi)
# ### Cliping the policy ratio
#
# \begin{align}
# \mathcal{L}^{CLIP}(\theta) =
# \mathbb{E}_{a_t, s_t \sim \pi_{\theta{OLD}}} \biggl[
# min \Bigl(r_t(\theta) \bar{A_t},
# clip \bigl(
# r_t(\theta), 1 - \epsilon, 1 + \epsilon
# \bigr) \bar{A_t}
# \Bigr)
# \biggr]
# \end{align}
#
# The ratio is clipped to be close to 1.
# We take the minimum so that the gradient will only pull
# $\pi_\theta$ towards $\pi_{\theta_{OLD}}$ if the ratio is
# not between $1 - \epsilon$ and $1 + \epsilon$.
# This keeps the KL divergence between $\pi_\theta$
# and $\pi_{\theta_{OLD}}$ constrained.
# Large deviation can cause performance collapse;
# where the policy performance drops and doesn't recover because
# we are sampling from a bad policy.
#
# Using the normalized advantage
# $\bar{A_t} = \frac{\hat{A_t} - \mu(\hat{A_t})}{\sigma(\hat{A_t})}$
# introduces a bias to the policy gradient estimator,
# but it reduces variance a lot.
clipped_ratio = ratio.clamp(min=1.0 - clip,
max=1.0 + clip)
policy_reward = torch.min(ratio * advantage,
clipped_ratio * advantage)
self.clip_fraction = (abs((ratio - 1.0)) > clip).to(torch.float).mean()
return -policy_reward.mean()
class ClippedValueFunctionLoss(nn.Module):
"""
## Clipped Value Function Loss
Similarly we clip the value function update also.
\begin{align}
V^{\pi_\theta}_{CLIP}(s_t)
&= clip\Bigl(V^{\pi_\theta}(s_t) - \hat{V_t}, -\epsilon, +\epsilon\Bigr)
\\
\mathcal{L}^{VF}(\theta)
&= \frac{1}{2} \mathbb{E} \biggl[
max\Bigl(\bigl(V^{\pi_\theta}(s_t) - R_t\bigr)^2,
\bigl(V^{\pi_\theta}_{CLIP}(s_t) - R_t\bigr)^2\Bigr)
\biggr]
\end{align}
Clipping makes sure the value function $V_\theta$ doesn't deviate
significantly from $V_{\theta_{OLD}}$.
"""
def forward(self, value: torch.Tensor, sampled_value: torch.Tensor, sampled_return: torch.Tensor, clip: float):
clipped_value = sampled_value + (value - sampled_value).clamp(min=-clip, max=clip)
vf_loss = torch.max((value - sampled_return) ** 2, (clipped_value - sampled_return) ** 2)
return 0.5 * vf_loss.mean()
labml_nn/rl/ppo/experiment.ipynb
+238
View File
@@ -0,0 +1,238 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb) \n",
"\n",
"## Proximal Policy Optimization - PPO\n",
"\n",
"This is an experiment training an agent to play Atari Breakout game using Proximal Policy Optimization - PPO"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ZCzmCrAIVg0L",
"outputId": "028e759e-0c9f-472e-b4b8-fdcf3e4604ee"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Add Atari ROMs (Doesn't work without this in Google Colab)"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"! wget http://www.atarimania.com/roms/Roms.rar\n",
"! mkdir /content/ROM/\n",
"! unrar e /content/Roms.rar /content/ROM/\n",
"! python -m atari_py.import_roms /content/ROM/"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"from labml import experiment\n",
"from labml.configs import FloatDynamicHyperParam, IntDynamicHyperParam\n",
"from labml_nn.rl.ppo.experiment import Trainer"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"ppo\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"### Configurations\n",
"\n",
"`IntDynamicHyperParam` and `FloatDynamicHyperParam` are dynamic hyper parameters\n",
"that you can change while the experiment is running."
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"configs = {\n",
" # number of updates\n",
" 'updates': 10000,\n",
" # number of epochs to train the model with sampled data\n",
" 'epochs': IntDynamicHyperParam(8),\n",
" # number of worker processes\n",
" 'n_workers': 8,\n",
" # number of steps to run on each process for a single update\n",
" 'worker_steps': 128,\n",
" # number of mini batches\n",
" 'batches': 4,\n",
" # Value loss coefficient\n",
" 'value_loss_coef': FloatDynamicHyperParam(0.5),\n",
" # Entropy bonus coefficient\n",
" 'entropy_bonus_coef': FloatDynamicHyperParam(0.01),\n",
" # Clip range\n",
" 'clip_range': FloatDynamicHyperParam(0.1),\n",
" # Learning rate\n",
" 'learning_rate': FloatDynamicHyperParam(2.5e-4, (0, 1e-3)),\n",
"}"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "0e978879-5dcd-4140-ec53-24a3fbd547de"
},
"source": [
"experiment.configs(configs)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "qYQCFt_JYsjd"
},
"source": [
"Create trainer"
]
},
{
"cell_type": "code",
"metadata": {
"id": "8LB7XVViYuPG"
},
"source": [
"trainer = Trainer(**configs)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"id": "aIAWo7Fw5DR8"
},
"source": [
"with experiment.start():\n",
" trainer.run_training_loop()"
],
"outputs": [],
"execution_count": null
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Proximal Policy Optimization - PPO",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.5"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
labml_nn/rl/ppo/experiment.py
+396
View File
@@ -0,0 +1,396 @@
"""
---
title: PPO Experiment with Atari Breakout
summary: Annotated implementation to train a PPO agent on Atari Breakout game.
---
# PPO Experiment with Atari Breakout
This experiment trains Proximal Policy Optimization (PPO) agent Atari Breakout game on OpenAI Gym.
It runs the [game environments on multiple processes](../game.html) to sample efficiently.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb)
"""
from typing import Dict
import numpy as np
import torch
from torch import nn
from torch import optim
from torch.distributions import Categorical
from labml import monit, tracker, logger, experiment
from labml.configs import FloatDynamicHyperParam, IntDynamicHyperParam
from labml_nn.rl.game import Worker
from labml_nn.rl.ppo import ClippedPPOLoss, ClippedValueFunctionLoss
from labml_nn.rl.ppo.gae import GAE
# Select device
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
class Model(nn.Module):
"""
## Model
"""
def __init__(self):
super().__init__()
# The first convolution layer takes a
# 84x84 frame and produces a 20x20 frame
self.conv1 = nn.Conv2d(in_channels=4, out_channels=32, kernel_size=8, stride=4)
# The second convolution layer takes a
# 20x20 frame and produces a 9x9 frame
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2)
# The third convolution layer takes a
# 9x9 frame and produces a 7x7 frame
self.conv3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1)
# A fully connected layer takes the flattened
# frame from third convolution layer, and outputs
# 512 features
self.lin = nn.Linear(in_features=7 * 7 * 64, out_features=512)
# A fully connected layer to get logits for $\pi$
self.pi_logits = nn.Linear(in_features=512, out_features=4)
# A fully connected layer to get value function
self.value = nn.Linear(in_features=512, out_features=1)
#
self.activation = nn.ReLU()
def forward(self, obs: torch.Tensor):
h = self.activation(self.conv1(obs))
h = self.activation(self.conv2(h))
h = self.activation(self.conv3(h))
h = h.reshape((-1, 7 * 7 * 64))
h = self.activation(self.lin(h))
pi = Categorical(logits=self.pi_logits(h))
value = self.value(h).reshape(-1)
return pi, value
def obs_to_torch(obs: np.ndarray) -> torch.Tensor:
"""Scale observations from `[0, 255]` to `[0, 1]`"""
return torch.tensor(obs, dtype=torch.float32, device=device) / 255.
class Trainer:
"""
## Trainer
"""
def __init__(self, *,
updates: int, epochs: IntDynamicHyperParam,
n_workers: int, worker_steps: int, batches: int,
value_loss_coef: FloatDynamicHyperParam,
entropy_bonus_coef: FloatDynamicHyperParam,
clip_range: FloatDynamicHyperParam,
learning_rate: FloatDynamicHyperParam,
):
# #### Configurations
# number of updates
self.updates = updates
# number of epochs to train the model with sampled data
self.epochs = epochs
# number of worker processes
self.n_workers = n_workers
# number of steps to run on each process for a single update
self.worker_steps = worker_steps
# number of mini batches
self.batches = batches
# total number of samples for a single update
self.batch_size = self.n_workers * self.worker_steps
# size of a mini batch
self.mini_batch_size = self.batch_size // self.batches
assert (self.batch_size % self.batches == 0)
# Value loss coefficient
self.value_loss_coef = value_loss_coef
# Entropy bonus coefficient
self.entropy_bonus_coef = entropy_bonus_coef
# Clipping range
self.clip_range = clip_range
# Learning rate
self.learning_rate = learning_rate
# #### Initialize
# create workers
self.workers = [Worker(47 + i) for i in range(self.n_workers)]
# initialize tensors for observations
self.obs = np.zeros((self.n_workers, 4, 84, 84), dtype=np.uint8)
for worker in self.workers:
worker.child.send(("reset", None))
for i, worker in enumerate(self.workers):
self.obs[i] = worker.child.recv()
# model
self.model = Model().to(device)
# optimizer
self.optimizer = optim.Adam(self.model.parameters(), lr=2.5e-4)
# GAE with $\gamma = 0.99$ and $\lambda = 0.95$
self.gae = GAE(self.n_workers, self.worker_steps, 0.99, 0.95)
# PPO Loss
self.ppo_loss = ClippedPPOLoss()
# Value Loss
self.value_loss = ClippedValueFunctionLoss()
def sample(self) -> Dict[str, torch.Tensor]:
"""
### Sample data with current policy
"""
rewards = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
actions = np.zeros((self.n_workers, self.worker_steps), dtype=np.int32)
done = np.zeros((self.n_workers, self.worker_steps), dtype=np.bool)
obs = np.zeros((self.n_workers, self.worker_steps, 4, 84, 84), dtype=np.uint8)
log_pis = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
values = np.zeros((self.n_workers, self.worker_steps + 1), dtype=np.float32)
with torch.no_grad():
# sample `worker_steps` from each worker
for t in range(self.worker_steps):
# `self.obs` keeps track of the last observation from each worker,
# which is the input for the model to sample the next action
obs[:, t] = self.obs
# sample actions from $\pi_{\theta_{OLD}}$ for each worker;
# this returns arrays of size `n_workers`
pi, v = self.model(obs_to_torch(self.obs))
values[:, t] = v.cpu().numpy()
a = pi.sample()
actions[:, t] = a.cpu().numpy()
log_pis[:, t] = pi.log_prob(a).cpu().numpy()
# run sampled actions on each worker
for w, worker in enumerate(self.workers):
worker.child.send(("step", actions[w, t]))
for w, worker in enumerate(self.workers):
# get results after executing the actions
self.obs[w], rewards[w, t], done[w, t], info = worker.child.recv()
# collect episode info, which is available if an episode finished;
# this includes total reward and length of the episode -
# look at `Game` to see how it works.
if info:
tracker.add('reward', info['reward'])
tracker.add('length', info['length'])
# Get value of after the final step
_, v = self.model(obs_to_torch(self.obs))
values[:, self.worker_steps] = v.cpu().numpy()
# calculate advantages
advantages = self.gae(done, rewards, values)
#
samples = {
'obs': obs,
'actions': actions,
'values': values[:, :-1],
'log_pis': log_pis,
'advantages': advantages
}
# samples are currently in `[workers, time_step]` table,
# we should flatten it for training
samples_flat = {}
for k, v in samples.items():
v = v.reshape(v.shape[0] * v.shape[1], *v.shape[2:])
if k == 'obs':
samples_flat[k] = obs_to_torch(v)
else:
samples_flat[k] = torch.tensor(v, device=device)
return samples_flat
def train(self, samples: Dict[str, torch.Tensor]):
"""
### Train the model based on samples
"""
# It learns faster with a higher number of epochs,
# but becomes a little unstable; that is,
# the average episode reward does not monotonically increase
# over time.
# May be reducing the clipping range might solve it.
for _ in range(self.epochs()):
# shuffle for each epoch
indexes = torch.randperm(self.batch_size)
# for each mini batch
for start in range(0, self.batch_size, self.mini_batch_size):
# get mini batch
end = start + self.mini_batch_size
mini_batch_indexes = indexes[start: end]
mini_batch = {}
for k, v in samples.items():
mini_batch[k] = v[mini_batch_indexes]
# train
loss = self._calc_loss(mini_batch)
# Set learning rate
for pg in self.optimizer.param_groups:
pg['lr'] = self.learning_rate()
# Zero out the previously calculated gradients
self.optimizer.zero_grad()
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)
# Update parameters based on gradients
self.optimizer.step()
@staticmethod
def _normalize(adv: torch.Tensor):
"""#### Normalize advantage function"""
return (adv - adv.mean()) / (adv.std() + 1e-8)
def _calc_loss(self, samples: Dict[str, torch.Tensor]) -> torch.Tensor:
"""
### Calculate total loss
"""
# $R_t$ returns sampled from $\pi_{\theta_{OLD}}$
sampled_return = samples['values'] + samples['advantages']
# $\bar{A_t} = \frac{\hat{A_t} - \mu(\hat{A_t})}{\sigma(\hat{A_t})}$,
# where $\hat{A_t}$ is advantages sampled from $\pi_{\theta_{OLD}}$.
# Refer to sampling function in [Main class](#main) below
# for the calculation of $\hat{A}_t$.
sampled_normalized_advantage = self._normalize(samples['advantages'])
# Sampled observations are fed into the model to get $\pi_\theta(a_t|s_t)$ and $V^{\pi_\theta}(s_t)$;
# we are treating observations as state
pi, value = self.model(samples['obs'])
# $-\log \pi_\theta (a_t|s_t)$, $a_t$ are actions sampled from $\pi_{\theta_{OLD}}$
log_pi = pi.log_prob(samples['actions'])
# Calculate policy loss
policy_loss = self.ppo_loss(log_pi, samples['log_pis'], sampled_normalized_advantage, self.clip_range())
# Calculate Entropy Bonus
#
# $\mathcal{L}^{EB}(\theta) =
# \mathbb{E}\Bigl[ S\bigl[\pi_\theta\bigr] (s_t) \Bigr]$
entropy_bonus = pi.entropy()
entropy_bonus = entropy_bonus.mean()
# Calculate value function loss
value_loss = self.value_loss(value, samples['values'], sampled_return, self.clip_range())
# $\mathcal{L}^{CLIP+VF+EB} (\theta) =
# \mathcal{L}^{CLIP} (\theta) +
# c_1 \mathcal{L}^{VF} (\theta) - c_2 \mathcal{L}^{EB}(\theta)$
loss = (policy_loss
+ self.value_loss_coef() * value_loss
- self.entropy_bonus_coef() * entropy_bonus)
# for monitoring
approx_kl_divergence = .5 * ((samples['log_pis'] - log_pi) ** 2).mean()
# Add to tracker
tracker.add({'policy_reward': -policy_loss,
'value_loss': value_loss,
'entropy_bonus': entropy_bonus,
'kl_div': approx_kl_divergence,
'clip_fraction': self.ppo_loss.clip_fraction})
return loss
def run_training_loop(self):
"""
### Run training loop
"""
# last 100 episode information
tracker.set_queue('reward', 100, True)
tracker.set_queue('length', 100, True)
for update in monit.loop(self.updates):
# sample with current policy
samples = self.sample()
# train the model
self.train(samples)
# Save tracked indicators.
tracker.save()
# Add a new line to the screen periodically
if (update + 1) % 1_000 == 0:
logger.log()
def destroy(self):
"""
### Destroy
Stop the workers
"""
for worker in self.workers:
worker.child.send(("close", None))
def main():
# Create the experiment
experiment.create(name='ppo')
# Configurations
configs = {
# Number of updates
'updates': 10000,
# ⚙️ Number of epochs to train the model with sampled data.
# You can change this while the experiment is running.
'epochs': IntDynamicHyperParam(8),
# Number of worker processes
'n_workers': 8,
# Number of steps to run on each process for a single update
'worker_steps': 128,
# Number of mini batches
'batches': 4,
# ⚙️ Value loss coefficient.
# You can change this while the experiment is running.
'value_loss_coef': FloatDynamicHyperParam(0.5),
# ⚙️ Entropy bonus coefficient.
# You can change this while the experiment is running.
'entropy_bonus_coef': FloatDynamicHyperParam(0.01),
# ⚙️ Clip range.
'clip_range': FloatDynamicHyperParam(0.1),
# You can change this while the experiment is running.
# ⚙️ Learning rate.
'learning_rate': FloatDynamicHyperParam(1e-3, (0, 1e-3)),
}
experiment.configs(configs)
# Initialize the trainer
m = Trainer(**configs)
# Run and monitor the experiment
with experiment.start():
m.run_training_loop()
# Stop the workers
m.destroy()
# ## Run it
if __name__ == "__main__":
main()
labml_nn/rl/ppo/gae.py
+82
View File
@@ -0,0 +1,82 @@
"""
---
title: Generalized Advantage Estimation (GAE)
summary: A PyTorch implementation/tutorial of Generalized Advantage Estimation (GAE).
---
# Generalized Advantage Estimation (GAE)
This is a [PyTorch](https://pytorch.org) implementation of paper
[Generalized Advantage Estimation](https://arxiv.org/abs/1506.02438).
You can find an experiment that uses it [here](experiment.html).
"""
import numpy as np
class GAE:
def __init__(self, n_workers: int, worker_steps: int, gamma: float, lambda_: float):
self.lambda_ = lambda_
self.gamma = gamma
self.worker_steps = worker_steps
self.n_workers = n_workers
def __call__(self, done: np.ndarray, rewards: np.ndarray, values: np.ndarray) -> np.ndarray:
"""
### Calculate advantages
\begin{align}
\hat{A_t^{(1)}} &= r_t + \gamma V(s_{t+1}) - V(s)
\\
\hat{A_t^{(2)}} &= r_t + \gamma r_{t+1} +\gamma^2 V(s_{t+2}) - V(s)
\\
...
\\
\hat{A_t^{(\infty)}} &= r_t + \gamma r_{t+1} +\gamma^2 r_{t+2} + ... - V(s)
\end{align}
$\hat{A_t^{(1)}}$ is high bias, low variance, whilst
$\hat{A_t^{(\infty)}}$ is unbiased, high variance.
We take a weighted average of $\hat{A_t^{(k)}}$ to balance bias and variance.
This is called Generalized Advantage Estimation.
$$\hat{A_t} = \hat{A_t^{GAE}} = \frac{\sum_k w_k \hat{A_t^{(k)}}}{\sum_k w_k}$$
We set $w_k = \lambda^{k-1}$, this gives clean calculation for
$\hat{A_t}$
\begin{align}
\delta_t &= r_t + \gamma V(s_{t+1}) - V(s_t)
\\
\hat{A_t} &= \delta_t + \gamma \lambda \delta_{t+1} + ... +
(\gamma \lambda)^{T - t + 1} \delta_{T - 1}
\\
&= \delta_t + \gamma \lambda \hat{A_{t+1}}
\end{align}
"""
# advantages table
advantages = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
last_advantage = 0
# $V(s_{t+1})$
last_value = values[:, -1]
for t in reversed(range(self.worker_steps)):
# mask if episode completed after step $t$
mask = 1.0 - done[:, t]
last_value = last_value * mask
last_advantage = last_advantage * mask
# $\delta_t$
delta = rewards[:, t] + self.gamma * last_value - values[:, t]
# $\hat{A_t} = \delta_t + \gamma \lambda \hat{A_{t+1}}$
last_advantage = delta + self.gamma * self.lambda_ * last_advantage
#
advantages[:, t] = last_advantage
last_value = values[:, t]
# $\hat{A_t}$
return advantages
labml_nn/rl/ppo/readme.md
+18
View File
@@ -0,0 +1,18 @@
# [Proximal Policy Optimization - PPO](https://nn.labml.ai/rl/ppo/index.html)
This is a [PyTorch](https://pytorch.org) implementation of
[Proximal Policy Optimization - PPO](https://arxiv.org/abs/1707.06347).
PPO is a policy gradient method for reinforcement learning.
Simple policy gradient methods one do a single gradient update per sample (or a set of samples).
Doing multiple gradient steps for a singe sample causes problems
because the policy deviates too much producing a bad policy.
PPO lets us do multiple gradient updates per sample by trying to keep the
policy close to the policy that was used to sample data.
It does so by clipping gradient flow if the updated policy
is not close to the policy used to sample the data.
You can find an experiment that uses it [here](https://nn.labml.ai/rl/ppo/experiment.html).
The experiment uses [Generalized Advantage Estimation](https://nn.labml.ai/rl/ppo/gae.html).
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb)