PPO¶

Proximal Policy Optimization collects a fixed-length rollout from n_envs parallel environments, then takes several passes of minibatch gradient steps on it. A clipped objective keeps each update close to the policy that collected the data, so the same rollout can be reused for n_epochs passes without the policy running away. PPO is the default on-policy agent in Ajax: it handles discrete and continuous actions, memory, image observations and the minibatch layouts of Brax's PPO.

Paper: Schulman, Wolski, Dhariwal, Radford and Klimov, Proximal Policy Optimization Algorithms, 2017 (arXiv:1707.06347).

Can I use?¶

Feature	PPO
Discrete actions	yes
Continuous actions	yes (Gaussian; `squash=True` for a tanh-bounded one)
Memory (`memory=`)	yes: GRU, LSTM, transformer, Mamba (Memory), without extensions
Environments	gymnax, Brax, MuJoCo Playground
Image observations	yes, a CNN encoder (`cnn_image_shape`)
PID policy	yes, `actor_architecture=PIDPolicy(...)` (Networks)
Several seeds in one call	yes, `train(seed=[...])`
Extension phases	`pretrain`, `post_update`, `eval_metrics`, `on_target`, `critic_loss`, `actor_loss`
Resume and checkpoints	yes; a split run equals the uninterrupted one (probe and resume contract)

Example¶

from ajax import PPO
from ajax.logging.wandb_logging import LoggingConfig

agent = PPO("CartPole-v1", n_envs=4, n_steps=128, batch_size=32)
logging = LoggingConfig(config={}, log_frequency=4_096, use_wandb=False)
state, evaluations = agent.train(seed=[1, 2], n_timesteps=20_480, logging_config=logging)
print(evaluations["Eval/episodic mean reward"])  # one row per seed, one column per evaluation

With memory, the update replays the rollout in contiguous chunks of bptt_length steps, each started from the memory the policy had there:

from ajax import PPO
from ajax.networks.memory import MemoryConfig

agent = PPO("CartPole-v1", n_envs=4, n_steps=128, batch_size=32,
            memory=MemoryConfig(kind="gru", hidden_size=32), bptt_length=32)
state, _ = agent.train(seed=0, n_timesteps=4_096)

Key hyperparameters¶

Defaults are those of the constructor.

Argument	Default	What it does
`n_envs`	4	environments stepped in parallel
`n_steps`	2048	steps per environment in each rollout
`batch_size`	64	sets the number of minibatches (see the notes)
`num_minibatches`	0	when positive, the number of minibatches, as in Brax
`n_epochs`	10	passes over the rollout per update, each reshuffled
`actor_learning_rate`, `critic_learning_rate`	3e-4	Adam step sizes (each a float or a schedule)
`adam_eps`	1e-5	Adam's epsilon (ppo2's)
`gamma`	0.99	discount
`gae_lambda`	0.95	GAE's lambda
`clip_range`	0.2	the ratio clip; a float or a schedule of the timestep
`ent_coef`	0	entropy bonus weight
`vf_coef`	0.5	critic loss weight, on half the squared error (ppo2's)
`max_grad_norm`	0.5	one gradient clip by global norm over both networks (`None` turns it off)
`normalize_advantage`	True	advantages standardised in each minibatch
`actor_architecture`, `critic_architecture`	`("128", "tanh", "128", "tanh")`	hidden layers, as width and activation strings
`normalize_observations`, `normalize_rewards`	False	running normalisation, as SB3's `VecNormalize`
`reward_normalization_gamma`	None	with `normalize_rewards`, divide by the spread of the discounted return
`memory`, `bptt_length`	None	memory block, and the length of the replayed chunks (by default the longest that split into the minibatches; see the notes)
`unroll_length`	None	use Brax's minibatching: fragments of this length, advantages per minibatch (feedforward only)
`log_std_state_independent`, `log_std_init`	True, 0.0	one learned log standard deviation per action dimension, and its initial value (see the notes)
`log_std_bounds`	None	the `(low, high)` clip of that log standard deviation; `None` leaves it unclipped
`actor_bias_init`, `critic_bias_init`	None	initialisers by name (Networks)
`encoder_kernel_init`	None	the hidden layers' initialiser; None is `"orthogonal(np.sqrt(2))"` (Baselines' ppo2; PureJaxRL's recurrent PPO, with `memory=`)
`cnn_image_shape`, `cnn_extra_obs_dim`, `cnn_spec`	None, 0, None	convolutional encoder for flat image observations
`episode_length`, `action_repeat`	None, 1	simulator steps per episode (1,000 when `None`) and per agent step, on Brax and Playground
`expose_recent_rollout`	False	keep the last rollout on the state, for measurement extensions
`extensions`	`()`	research features (see Extensions)

Notes¶

batch_size sets the number of minibatches, not their size. Unless num_minibatches is positive, the rollout of n_steps * n_envs samples splits into max(n_steps, batch_size) // min(n_steps, batch_size) minibatches. With the defaults that is 32 minibatches of 256 samples (batch_size * n_envs). This differs from Stable-Baselines3, where batch_size is the minibatch size.
Shuffling. Every epoch draws a new split of the rollout into minibatches, as OpenAI Baselines' and CleanRL's PPO do.
Where the advantages are computed. The advantages are computed once on the whole rollout and its n_steps * n_envs samples are shuffled flat into the minibatches, as in the PPO paper's Algorithm 1 and OpenAI Baselines' ppo2. unroll_length switches to Brax's PPO: the rollout is cut into fragments of that length that bootstrap at their ends, each minibatch holds whole fragments, and its advantages are recomputed with the current critic. unroll_length=n_steps keeps each environment's whole rollout as one fragment. PPO refuses at construction a length that does not split the rollout evenly into the minibatches.
Actions. A continuous PPO samples an unbounded Gaussian by default; Ajax clips the sample to [-1, 1] and maps that range onto the environment's bounds, so on gymnax Pendulum (torque in [-2, 2]) one unit of the Gaussian is two units of torque. squash=True bounds it with tanh instead. Evaluation uses the mean action.
Exploration noise. The Gaussian learns one log standard deviation per action dimension, whatever the observation, starting at 0 (a standard deviation of 1) and never clipped (log_std_bounds=None, which the run config records), as PPO's reference code does: OpenAI Baselines at ea25b9e (ppo2's common/distributions.py:96-106; ppo1's ppo1/mlp_policy.py:36-39, gaussian_fixed_var=True), CleanRL at fe8d8a0 (ppo_continuous_action.py:129) and PureJaxRL at 31756b1 (ppo_continuous_action.py:41). As in theirs, it is updated with the rest of the model, under one gradient clip over both networks (see separate networks). log_std_bounds clips every log standard deviation PPO can learn, so by default none is clipped. log_std_state_independent=False with log_std_init=-1.0 builds the head PPO had before: a log standard deviation computed from the observation, starting at -1; log_std_bounds=(-20.0, 2.0) adds the clip it then had, SAC's. A checkpoint saved with that head restores only into that network, and restoring it into the default one is refused with this advice. A PID policy keeps its own log standard deviation, which starts at log_std_init and is clipped to log_std_bounds.
Brax-style options. mean_kernel_init, squash, num_resets_per_eval and num_evals reproduce choices of Brax's PPO. They are off by default, as is encoder_layer_norm (a LayerNorm on the encoder's output).
Separate networks. The actor and the critic are separate networks. As ppo2's one optimizer, their gradients are clipped together to one global norm, then each takes an Adam step (Adam is per coordinate, so the two are one Adam over the whole model); the value loss is not clipped. As in Baselines' ppo2, their hidden layers start orthogonal with gain sqrt(2), the policy's output layer with gain 0.01 and the value layer with gain 1. An entropy bonus on a squashed policy is a one-sample estimate of the executed action's entropy, as in Brax.
Normalisation. With normalize_observations=True the loss sees the observations normalised with the statistics of the step that collected them, as SB3's VecNormalize and CleanRL do.
Extensions. PPO folds the loss and target phases, so instrumentation such as ConditioningMetrics works on it (pass expose_recent_rollout=True so it can read the latest rollout). It has no action phase: extensions that act during collection need SAC. A recurrent PPO refuses extensions. See Extensions.
Memory. A recurrent PPO minibatches whole chunks of bptt_length steps of the rollout. Libraries disagree on that length, so by default the number of minibatches chooses it: each environment's whole rollout when n_envs is a multiple of the number of minibatches, as OpenAI Baselines' ppo2, CleanRL and PureJaxRL do, otherwise the fewest equal chunks that split (256 steps with the defaults). PPO refuses at construction chunks that do not split. Memory compares the libraries and says how to set it: to reproduce theirs, choose num_minibatches to divide n_envs. unroll_length is for feedforward policies; recurrent ones use bptt_length.

GAIL: imitation by a discriminator¶

GAIL (Ho and Ermon, Generative Adversarial Imitation Learning, NeurIPS 2016, arXiv:1606.03476) trains a policy to imitate demonstrations without their rewards. A discriminator learns to tell the policy's rollouts from the demonstrations, and the policy is rewarded where it is mistaken for the expert. In Ajax it is PPO with the AdversarialReward extension, which folds the reward phase:

import jax
import jax.numpy as jnp

from ajax import PPO
from ajax.buffers.dataset import record
from ajax.extensions.adversarial import AdversarialReward


def pd_expert(obs):
    theta = jnp.arctan2(obs[..., 1], obs[..., 0])
    return jnp.clip(-(2.0 * theta + 0.5 * obs[..., 2]), -2.0, 2.0)[..., None]


env_args = PPO("Pendulum-v1").env_args
demonstrations = record(pd_expert, env_args, jax.random.PRNGKey(0), n_steps=1_000)
agent = PPO("Pendulum-v1", extensions=(AdversarialReward(demonstrations),))

After each rollout, the rewards become -log(1 - D(s, a)) from the current discriminator, and the discriminator takes n_updates (2) cross-entropy steps on the rollout's rows and as many demonstration rows. It follows HumanCompatibleAI/imitation at e5ef188 (algorithms/adversarial/gail.py, common.py): its 32-32 ReLU reward network and torch's Adam at 1e-3. PPO replaces the paper's TRPO, as in that code. The evaluations still report the environment's return. The discriminator reads the action the environment runs, clipped to [-1, 1], on the rollout and on the demonstrations, as the reference records the clipped action SB3 steps with. Deviations: it reads (s, a) without the running normalisation of the reference's examples, and trains on the whole rollout, where the reference draws demo_batch_size rows.

Results¶

The reproduction results set PPO beside Stable-Baselines3's on CartPole and Brax's on inverted_pendulum, with learning curves. Unit tests (tests/agents/PPO and the contracts every agent shares) check its pieces against exact contracts, without learning, such as the fresh split of each epoch. The probes (tests/probing) train it briefly on tiny environments with known answers: the exact values of small problems, episode ends, memory, extension phases, and a run split by a checkpoint matching the run left whole. On the CPU speed benchmark (benchmarks/agent_bench.py: CartPole, 4 environments, 8,000 steps) PPO ran at 17,817 steps per second of execution (4,298 with compilation) at main a8a13c2 (PERFORMANCE_REPORT.md).

API¶

ajax.agents.PPO.PPO.PPO ¶

PPO(
    env_id,
    n_envs=4,
    actor_learning_rate=0.0003,
    critic_learning_rate=0.0003,
    actor_architecture=("128", "tanh", "128", "tanh"),
    critic_architecture=("128", "tanh", "128", "tanh"),
    gamma=0.99,
    env_params=None,
    max_grad_norm=0.5,
    weight_decay=0.0,
    ent_coef=0,
    clip_range=0.2,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    num_minibatches=0,
    adam_eps=1e-05,
    vf_coef=0.5,
    unroll_length=None,
    num_resets_per_eval=0,
    num_evals=1,
    gae_lambda=0.95,
    normalize_advantage=True,
    memory=None,
    bptt_length=None,
    normalize_observations=False,
    normalize_rewards=False,
    reward_normalization_gamma=None,
    actor_bias_init=None,
    critic_bias_init=None,
    encoder_kernel_init=None,
    log_std_state_independent=True,
    log_std_init=0.0,
    log_std_bounds=None,
    mean_kernel_init=None,
    encoder_layer_norm=False,
    squash=False,
    episode_length=None,
    cnn_image_shape=None,
    cnn_extra_obs_dim=0,
    cnn_spec=None,
    extensions=(),
    expose_recent_rollout=False,
    action_repeat=1,
)

Bases: ActorCritic

Proximal Policy Optimization (Schulman et al., 2017) for discrete and continuous action spaces.