The 32 Implementation Details of Proximal Policy Optimization (PPO) Algorithm

DEPRECATION NOTICE:

Deprecated Content:

Important links before the main content

• Source code
  • For Atari 2600 games:
    • https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari.py
    • https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_visual.py (with action probability distribution visualization)
  • For Mujoco, Pybullet and others:
    • https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_continuous_action.py
• Our experiments metrics and videos of agents playing the game
  • https://app.wandb.ai/cleanrl/cleanrl.benchmark/reports/PPO-Reproduction--VmlldzoxMzAzNTQ

General implementation details

1. Clipped surrogate objective
   • Obviously, this is the core feature of PPO, so please implement it 😉.
2. Generalized Advantage Estimation (GAE)
   • Although the original PPO paper just uses the abstraction of advantage estimate in the PPO's objective, the implementation does use GAE.
3. Normalization of Advantages
   • After calculating the advantages based on GAE, the advantages vector is normalized by subtracting its mean and divided by its standard deviation. In particular, notice this normalization happens at the minibatch level instead of the whole batch level!
4. Value Function Loss Clipping
   • The value function loss is clipped in a manner that is similar to the PPO’s clipped surrogate objective.
5. Overall Loss Includes Entropy Loss
   • The overall loss is calculated as loss = pg_loss - entropy * ent_coef + vf_loss * vf_coef, which includes entropy maximization, which intuitively encourages the exploration by encouraging the action probability distribution to be more chaotic.
6. Adam Learning Rate Annealing
   • The Adam optimizer’s learning rate could be either constant or set to decay. In particular, the hyper-parameters for training agents playing Atari games set the learning rate to decay from 2.5e-4 to 0 as the number of timesteps agent trained increases. In mujoco, it decays from 3e-4 to 0.
7. Mini-batch Updates
   • The PPO implementation of the openai/baselines also uses mini-batches to compute the gradient and update the policy instead of the whole batch such as in openai/spinningup.
8. Global Gradient Clipping
   • For each update iteration in an epoch, the gradients of the policy and value network are clipped so that the “global l2 norm” (i.e. the norm of the concatenated gradients of all parameters) does not exceed 0.5
9. Orthogonal Initialization of Weights and Constant Initialization of biases
   • The related code is actually scattered all over the places in the openai/baselines library. The code to do such initialization is somehow in baselines/baselines/a2c/utils.py, when in fact it is used for other algorithms such as PPO. In general, the weights of hidden layers use orthogonal initialization of weights with scaling np.sqrt(2), and the biases are set to 0, as shown in the CNN initialization for Atari, and the MLP initialization for Mujoco, which are located at baselines/baselines/common/models.py. However, the weights of the policy output layer are initialized with the scale of 0.01, and the weights of the value output layer are initialized with the scale of 1, and their corresponding code is located at baselines/baselines/common/policies.py. By the way, it seems the implementation of the orthogonal initialization of openai/baselines seems different from that of pytorch/pytorch, but I consider this to be a very low-level detail that should not have much impact on the performance.
10. The Use of Parallelized Environment
    • PPO uses the parallelized environment to speed up execution. As mentioned in the paper, it uses ``fixed-length trajectory segments''. I think the use of such environments is actually under-appreciated by most people, as it kind of changes the underlying MDP somehow. I will dedicate a section below to discuss it in detail.
11. The Epsilon Parameter of Adam Optimizer
    • By default, this parameter is not in the list of configurable hyper-parameters of PPO. However, it is set to 1e-5, Which is different than the default epsilon of 1e-8 in PyTorch and TensorFlow.

Mujoco specific implementation details

# https://github.com/openai/baselines/blob/master/baselines/ppo2/defaults.py
def mujoco():
    return dict(
        nsteps=2048,
        nminibatches=32,
        lam=0.95,
        gamma=0.99,
        noptepochs=10,
        log_interval=1,
        ent_coef=0.0,
        lr=lambda f: 3e-4 * f,
        cliprange=0.2,
        value_network='copy'
    )

1. Normalization of Observation
   • At each timestep, the VecNormalize wrapper pre-processes the observation before feeding it to the PPO agent. The raw observation was normalized by subtracting its running mean and divided by its variance.
2. Observation Clipping
   • Followed by the normalization of observation, the normalized observation is further clipped by VecNormalize to a range, usually [−10, 10].
3. Reward Scaling
   • The VecNormalize also applies a certain discount-based scaling scheme, where the rewards are divided by the standard deviation of a rolling discounted sum of the rewards (without subtracting and re-adding the mean).
4. Reward Clipping
   • Followed by the scaling of reward, the scaled reward is further clipped by VecNormalize to a range, usually [−10, 10].
5. The Way Standard Deviation is Paramterized
   • Policy gradient methods (including PPO) assume the continuous actions are sampled from a normal distribution. So to create such distribution, the neural network needs to output the mean and standard deviation of the continuous action. The implementation outputs the logits for the mean, but instead of outputting the logits for the standard deviation, it outputs the logarithm of the standard deviation. In addition, this log std is set to be state-independent and initialized to be 0.
6. Hyperbolic Tangent Activations
   • The activation functions of the hidden layers are always set to use the hyperbolic tangent function.
7. Not Sharing Hidden Layers for Policy and Value Functions
   • The network of the policy and value functions are initialized separately.
8. Handling of action clipping to valid range and storage
   • After a continuous action is sampled, Such action could be invalid because it could exceed the valid range of continuous actions in the environment. To avoid this, add applies the rapper to clip the action into the valid range. However, the original unclipped action is stored as part of the episodic data.

Atari specific implementation details

# https://github.com/openai/baselines/blob/master/baselines/ppo2/defaults.py
def atari():
    return dict(
        nsteps=128, nminibatches=4,
        lam=0.95, gamma=0.99, noptepochs=4, log_interval=1,
        ent_coef=.01,
        lr=lambda f : f * 2.5e-4,
        cliprange=0.1,
    )

1. The Use of NoopResetEnv
   • This wrapper samples initial states by taking a random number of no-ops on reset. No-op is assumed to be action 0.
2. The Use of FireResetEnv
   • This wrapper takes action of FIRE on reset for environments that are fixed until firing.
3. The Use of EpisodicLifeEnv
   • This wrapper makes end-of-life == end-of-episode but only resets on the true game over. Done by DeepMind for the DQN and co. since it helps value estimation. This is an arguably very interesting wrapper in my opinion because it changes the initial state distribution of the game.
4. The Use of MaxAndSkipEnv
   • This wrapper repeats actions, sums rewards, and returns the max-pooled observation.
5. The Use of WarpFrame
   • This wrapper warps frames to 84x84 as done in the Nature paper and later work. If the environment uses dictionary observations
6. The Use of FrameStack
   • This wrapper stacks k last frames such that the agent can infer the velocity and directions of moving objects.
7. The Use of ClipRewardEnv
   • This wrapper bins reward to {+1, 0, -1} by its sign. According to our initial experiments, it seems to have a huge impact on the PPO's performance on Breakout.
8. Scaling the Images to Range [0, 1]
   • The input data has the range of [0,255], but it is divided by 255 to be in the range of [0,1]. My initial experiments found this scaling to be extremely important. Without it, the first policy update results in the Kullback–Leibler divergence explodes, likely due to how the layers are initialized.
9. Rectified Linear Unit (ReLU) Activations
   • The activation functions for the hidden and output layers are always ReLU.
10. Sharing Hidden Layers for Policy and Value Functions
    • The hidden layers for the policy and value functions share the same weights and biases. Although I don't think it makes much difference for the performance, the computational cost is definitely reduced since only one set of hidden layers needs to be optimized.

Auxiliary implementation details

1. Clip Range Annealing
   • The clip coefficient of PPO can be annealed similar to how the learning rate is annealed. However, as shown by the hyperparameters above, the clip range annealing is actually not used.
2. Parallellized Gradient Update
   • The policy gradient is calculated in parallel using multiple processes. However, I consider this as an auxiliary detail because it is difficult to implement and according to my experience does not improve the performance, measured in episode rewards achieved.
3. Early Stopping
   • This is not actually an implementation detail of openai/baselines, but rather an implementation detail in openai/spinningup. This early stopping optimization measures that mean KL divergence between the target and the current policy of PPO, and stops the policy updates of the current epoch if the mean KL divergence exceeds some preset threshold. This feature is implemented and can be toggled by using --kle-stop

A dedicated note to the use of parallelized environment

1. True termination. The episode of the game really terminates. As an example, the true termination of the Breakout game in Atari 2600 comes when you lose all of your lives.
2. Time limit termination. To avoid the episode of the game going on indefinitely, the episode of the game is set to finish with a maximum of N time steps.

import gym
import numpy as np
from gym import error, spaces, utils
from gym.envs.registration import register

"""
First kind of termination: true termination
As an example, the true termination of the Breakout game in Atari 2600
comes when you lose all of your lives.
"""
class TestEnv(gym.Env):
    """
    A simple env that always ends after 10 timesteps, which 
    can be considered as the ``true termination'' from the environment.
    At each timestep, its observation is its internal timestep
    """
    def __init__(self):
        self.action_space = spaces.Discrete(2)
        self.observation_space = spaces.Box(low=np.array([-1.]),
            high=np.array([10.]))
    def step(self, action):
        self.t += 1
        return np.array([0.])+self.t, 1, self.t==10, {}
    def reset(self):
        self.t = 0
        return np.array([0.])+self.t

if "TestEnv-v0" not in gym.envs.registry.env_specs:
    register(
        "TestEnv-v0",
        entry_point='__main__:TestEnv'
    )

env = gym.make("TestEnv-v0")
print(f"env is {env}")
for i in range(2):
    all_obs = [env.reset()]
    while True:
        obs, reward, done, info = env.step(env.action_space.sample())
        all_obs += [obs]
        if done:
            print(f"all observation in episode {i}:")
            print(all_obs)
            print("true termination")
            print()
            break
print("=========")

env is <TestEnv<TestEnv-v0>>
all observation in episode 0:
[array([0.]), array([1.]), array([2.]), array([3.]), array([4.]), 
array([5.]), array([6.]), array([7.]), array([8.]), array([9.]), array([10.])]
true termination

all observation in episode 1:
[array([0.]), array([1.]), array([2.]), array([3.]), array([4.]), 
array([5.]), array([6.]), array([7.]), array([8.]), array([9.]), array([10.])]
true termination

=========

"""
Second kind of termination: TimeLimit termination
As an example, TimeLimit termination comes when the episode of "CartPole-v0"
exceeds length 200.
"""
if "TestEnvTimeLimit3-v0" not in gym.envs.registry.env_specs:
    register(
        "TestEnvTimeLimit3-v0",
        entry_point='__main__:TestEnv',
        max_episode_steps=8
    )

env = gym.make("TestEnvTimeLimit3-v0")
# equivalent to below
# env = TestEnv()
# env = gym.wrappers.TimeLimit(env, max_episode_steps=8)
print(f"env is {env}")
print(f"env's timelimit is {env._max_episode_steps}")
for i in range(2):
    all_obs = [env.reset()]
    while True:
        obs, reward, done, info = env.step(env.action_space.sample())
        all_obs += [obs]
        if done:
            print(f"all observation in episode {i}:")
            print(all_obs)
            print("TimeLimit termination")
            print()
            break
print("=========")

env is <TimeLimit<TestEnv<TestEnvTimeLimit3-v0>>>
env's timelimit is 8
all observation in episode 0:
[array([0.]), array([1.]), array([2.]), array([3.]), array([4.]), 
array([5.]), array([6.]), array([7.]), array([8.])]
time limit termination

all observation in episode 1:
[array([0.]), array([1.]), array([2.]), array([3.]), array([4.]), 
array([5.]), array([6.]), array([7.]), array([8.])]
time limit termination

=========

from stable_baselines3.common.vec_env import DummyVecEnv
"""
Third kind of termination: early termination induced by fixed 
trajectory length of `n_steps`
This is usually combined with TimeLimit wrapped env,
but you can use it without the TimeLimit
"""
n_steps = 5
envs = DummyVecEnv([
    lambda: gym.make("TestEnvTimeLimit3-v0")])
print(f"envs is {envs}")
print(f"envs' timelimit is {envs.envs[0]._max_episode_steps}")
obss = envs.reset()
for i in range(3):
    all_obss = []
    for j in range(n_steps):
        all_obss += [obss.astype("float")]
        obss, rewards, dones, infos = envs.step(np.array([1.,1.]))
        # print(infos)
    
    print(f"all observation in trajectory {i}:")
    print(all_obss)
    print("early termination by `n_steps`")
    print()
print("=========")

envs is <stable_baselines3.common.vec_env.dummy_vec_env.DummyVecEnv 
object at 0x7f7381256f10>
envs' timelimit is 8
all observation in trajectory 0:
[array([[0.]]), array([[1.]]), array([[2.]]), array([[3.]]), array([[4.]])]
early termination by `n_steps`

all observation in trajectory 1:
[array([[5.]]), array([[6.]]), array([[7.]]), array([[0.]]), array([[1.]])]
early termination by `n_steps`

all observation in trajectory 2:
[array([[2.]]), array([[3.]]), array([[4.]]), array([[5.]]), array([[6.]])]
early termination by `n_steps`

=========

1. The initial observation of the trajectory is no longer always array([0.]), which means the initial state distribution of the trajectory is no longer the same as that of the original MDP.
2. The time-limited episode of the game is broken into pieces presented in different trajectories.
3. The terminal observation array([8.]) is not present in any trajectory.

Results

Results on Atari

Results on Mujoco

Results on Pybullet and others

A case study on HalfCheetah-v2

Fun visualization of action probability distribution

Reproduction of our results

wandb login

for seed in {1..2}
do
    (sleep 0.3 && nohup xvfb-run -a python ppo_atari_visual.py \
    --gym-id BeamRiderNoFrameskip-v4 \
    --total-timesteps 10000000 \
    --wandb-project-name cleanrl.benchmark \
    --wandb-entity cleanrl \
    --prod-mode \
    --capture-video \
    --seed $seed
    ) >& /dev/null &
done

for seed in {1..2}
do
    (sleep 0.3 && nohup xvfb-run -a python ppo_continuous_action.py \
    --gym-id Hopper-v2 \
    --total-timesteps 2000000 \
    --wandb-project-name cleanrl.benchmark \
    --wandb-entity cleanrl \
    --prod-mode \
    --capture-video \
    --seed $seed
    ) >& /dev/null &
done

docker run -d --cpuset-cpus="0" -e WANDB={REPLACE_WITH_YOUR_WANDB_KEY} \
 vwxyzjn/cleanrl:latest python ppo_atari_visual.py \
--gym-id BeamRiderNoFrameskip-v4 --total-timesteps 10000000 \
--wandb-project-name cleanrl.benchmark --wandb-entity cleanrl \
--prod-mode --capture-video --seed 1

About our CleanRL

Research-friendly features

1. Seeding of everything
   • Our implementation always has a section the seeds the following modules. Although there are still sources of indeterminism for some environments, results are mostly reproducible within a single machine. This makes debugging a lot easier.
2. Cloud experiment management
   • We use a cloud experiment management service named weights and biases. It logs the hyperparameters, training metrics (e.g. episode reward, training losses), system metrics (e.g. CPU utilization, memory utilization), stdout, stderr, source code (this is especially helpful since we have single file implementation, which makes version management very easy), and videos of the agents playing games.
3. Cloud video logging of agents playing the game
   • I want to additionally stress the importance of logging the videos of the agents playing games. Visually inspected the agents' actual behavior gives you far more insight than the metric of episode reward could ever give you.
4. Integration with Docker and AWS Batch
   • CleanRL actually scales quite well if your experiments individually don't run for billions of time steps. We package the files into docker containers, and by leveraging AWS batch, we have finished tasks about 8000 CPU-hours in four hours, costing about $100 via spot instances. One of those days I'll find time to write instructions on how to do that 😅.
   • For now, just know you can run experiments using docker run -d --cpuset-cpus="1" -e WANDB=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx vwxyzjn/cleanrl:latest python ppo_atari_visual.py --gym-id BreakoutNoFrameskip-v4 --total-timesteps 10000000 --wandb-project-name cleanrl.benchmark --wandb-entity cleanrl --prod-mode --capture-video --seed 2

   

Long-term goal: Open RL Benchmark

Getting involved

Support us