Proximal Policy Optimization (PPO)¶

PPO is an on-policy policy gradient algorithm that uses a clipped objective to constrain policy updates. It aims to combine the stability of Trust Region Policy Optimization (TRPO) with the simplicity and scalability of vanilla policy gradients, effectively maintaining a balance between exploration and exploitation.

AgileRL offers support for recurrent policies in PPO to solve Partially Observable Markov Decision Processes (POMDPs). For more information, please refer to the Partially Observable Markov Decision Processes (POMDPs) documentation, or our tutorial on solving Pendulum-v1 with masked angular velocity observations here.

Compatible Action Spaces¶

`Discrete`	`Box`	`MultiDiscrete`	`MultiBinary`
✔️	✔️	✔️	✔️

LunarLanderContinuous-v3 Example¶

import numpy as np
from tqdm import tqdm

from agilerl.algorithms.ppo import PPO
from agilerl.rollouts.on_policy import collect_rollouts
from agilerl.utils.utils import make_vect_envs

# Create environment
num_envs = 16
max_steps = 100000
env = make_vect_envs('LunarLanderContinuous-v3', num_envs=num_envs)
observation_space = env.single_observation_space
action_space = env.single_action_space

# Create PPO agent
agent = PPO(
    observation_space,
    action_space,
    lr=1e-3,
    batch_size=128,
    learn_step=2048
)

pbar = tqdm(total=max_steps)
total_steps = 0
while total_steps < max_steps:
    agent.set_training_mode(True)

    # Collect rollouts and save in the agent's rollout buffer
    episode_scores = collect_rollouts(agent, env)

    agent.learn()    # Learn from rollout buffer

    total_steps += agent.learn_step
    agent.steps += agent.learn_step
    pbar.update(agent.learn_step)
    if episode_scores:
        pbar.set_description(f"Score: {np.mean(episode_scores)}")

Neural Network Configuration¶

To configure the architecture of the network’s encoder / head, pass a kwargs dict to the PPO net_config field. Full arguments can be found in the documentation of EvolvableMLP, EvolvableCNN, EvolvableMultiInput, and EvolvableLSTM.

For discrete / vector observations:

NET_CONFIG = {
      "encoder_config": {'hidden_size': [32, 32]},  # Network head hidden size
      "head_config": {'hidden_size': [32]}      # Network head hidden size
  }

For image observations:

NET_CONFIG = {
    "encoder_config": {
      'channel_size': [32, 32], # CNN channel size
      'kernel_size': [8, 4],   # CNN kernel size
      'stride_size': [4, 2],   # CNN stride size
    },
    "head_config": {'hidden_size': [32]}  # Network head hidden size
  }

For dictionary / tuple observations containing any combination of image, discrete, and vector observations:

CNN_CONFIG = {
    "channel_size": [32, 32], # CNN channel size
    "kernel_size": [8, 4],   # CNN kernel size
    "stride_size": [4, 2],   # CNN stride size
}

NET_CONFIG = {
    "encoder_config": {
      "latent_dim": 32,
      # Config for nested EvolvableCNN objects
      "cnn_config": CNN_CONFIG,
      # Config for nested EvolvableMLP objects
      "mlp_config": {
          "hidden_size": [32, 32]
      },
      "vector_space_mlp": True # Process vector observations with an MLP
    },
    "head_config": {'hidden_size': [32]}  # Network head hidden size
  }

For recurrent observations:

NET_CONFIG = {
  "encoder_config": {
    "hidden_state_size": 64,
    "num_layers": 1,
    "max_seq_len": 512,
  },
  "head_config": {
    "hidden_size": [64],
  }
}

# Create PPO agent
agent = PPO(
  observation_space=observation_space,
  action_space=action_space,
  net_config=NET_CONFIG
  )

Evolutionary Hyperparameter Optimization¶

AgileRL allows for efficient hyperparameter optimization during training to provide state-of-the-art results in a fraction of the time. For more information on how this is done, please refer to the Evolutionary Hyperparameter Optimization documentation.

Saving and Loading Agents¶

To save an agent, use the save_checkpoint method:

from agilerl.algorithms.ppo import PPO

agent = PPO(observation_space, action_space)   # Create PPO agent

checkpoint_path = "path/to/checkpoint"
agent.save_checkpoint(checkpoint_path)

To load a saved agent, use the load method:

from agilerl.algorithms.ppo import PPO

checkpoint_path = "path/to/checkpoint"
agent = PPO.load(checkpoint_path)

Parameters¶

class agilerl.algorithms.ppo.PPO(*args: Any, **kwargs: Any)¶

Proximal Policy Optimization (PPO).

Paper: https://arxiv.org/abs/1707.06347v2

Parameters:

observation_space (gym.spaces.Space) – Observation space of the environment
action_space (gym.spaces.Space) – Action space of the environment
index (int, optional) – Index to keep track of object instance during tournament selection and mutation, defaults to 0
hp_config (HyperparameterConfig, optional) – RL hyperparameter mutation configuration, defaults to None, whereby algorithm mutations are disabled.
net_config (dict, optional) – Network configuration, defaults to None
batch_size (int, optional) – Size of batched sample from replay buffer for learning, defaults to 64
lr (float, optional) – Learning rate for optimizer, defaults to 1e-4
learn_step (int, optional) – Learning frequency, defaults to 4096
gamma (float, optional) – Discount factor, defaults to 0.99
gae_lambda (float, optional) – Lambda for general advantage estimation, defaults to 0.95
mut (str, optional) – Most recent mutation to agent, defaults to None
action_std_init (float, optional) – Initial action standard deviation, defaults to 0.6
clip_coef (float, optional) – Surrogate clipping coefficient, defaults to 0.2
ent_coef (float, optional) – Entropy coefficient, defaults to 0.01
vf_coef (float, optional) – Value function coefficient, defaults to 0.5
max_grad_norm (float, optional) – Maximum norm for gradient clipping, defaults to 0.5
target_kl (float, optional) – Target KL divergence threshold, defaults to None
normalize_images (bool, optional) – Flag to normalize images, defaults to True
update_epochs (int, optional) – Number of policy update epochs, defaults to 4
actor_network (nn.Module, optional) – Custom actor network, defaults to None
critic_network (nn.Module, optional) – Custom critic network, defaults to None
share_encoders (bool, optional) – Flag to share encoder parameters between actor and critic, defaults to False
num_envs (int, optional) – Number of parallel environments. Omitted means 1.
rollout_buffer_config (dict[str, Any] | None, optional) – Extra keyword arguments forwarded to the rollout buffer constructor, defaults to None (treated as an empty dict).
recurrent (bool, optional) – Flag to use hidden states for recurrent policies, defaults to False
device (str, optional) – Device for accelerated computing, ‘cpu’ or ‘cuda’, defaults to ‘cpu’
accelerator (accelerate.Accelerator(), optional) – Accelerator for distributed computing, defaults to None
wrap (bool, optional) – Wrap models for distributed training upon creation, defaults to True
bptt_sequence_type (BPTTSequenceType, optional) – How recurrent rollouts are cut into training sequences. Omitted means chunked.
max_seq_len (int, optional) – Maximum sequence length for truncated BPTT, defaults to None, where complete episodes are used as sequences.

add_scores(scores: Sequence[float | list[float]]) → None¶

Add scores to the metrics.

Parameters:: scores (Sequence[float | list[float]]) – List of scores (or per-agent score rows) to add.

clean_up() → None¶

Clean up the algorithm by deleting the networks and optimizers.

Returns:: None
Return type:: None

clone(index: int | None = None, wrap: bool = True) → Self¶

Create a clone of the algorithm.

Parameters:

index (int | None, optional) – The index of the clone, defaults to None
wrap (bool, optional) – If True, wrap the models in the clone with the accelerator, defaults to False

Returns:

A clone of the algorithm

Return type:

EvolvableAlgorithm

static copy_attributes(agent: IndividualT, clone: IndividualT, exclude: Iterable[str] = ()) → IndividualT¶

Copy the non-evolvable attributes of the algorithm to a clone.

Parameters:

agent (EvolvableAlgorithm) – The algorithm to copy attributes from.
clone (EvolvableAlgorithm) – The clone of the algorithm.
exclude (Iterable[str]) – Attribute names to leave on clone / agent.

Returns:

The clone of the algorithm.

Return type:

EvolvableAlgorithm

create_rollout_buffer() → None¶: Create a rollout buffer with the current configuration.

eval_policy_network_ids() → set[int]¶

Return the id of every evaluation network in the agent’s policy group.

Returns:: Identities of the policy’s evaluation networks.
Return type:: set[int]

Evaluate the actions.

Parameters:

obs (ObservationType) – Environment observation, or multiple observations in a batch
actions (torch.Tensor) – Actions to evaluate
hidden_state (dict[str, torch.Tensor] | None) – Hidden state for recurrent policies, defaults to None. Expected shape: dict with tensors of shape (batch_size, 1, hidden_size).
action_mask (ActionMaskInput) – Mask of legal actions 1=legal 0=illegal, defaults to None

Returns:

Log probability, entropy, state values

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

evolvable_attributes(networks_only: bool = False) → dict[str, Any]¶

Return the attributes related to the evolvable networks in the algorithm. Includes attributes that are either EvolvableModule or ModuleDict objects, as well as the optimizers associated with the networks.

Parameters:: networks_only (bool, optional) – If True, only include evolvable networks, defaults to False
Returns:: A dictionary of network attributes.
Return type:: dict[str, Any]

finalize_training_step(num_steps: int) → None¶

Close the agent’s training block, storing any captured GraMa scores.

Parameters:: num_steps (int) – Number of steps taken during the training step.
Returns:: None.
Return type:: None

property fitness: list[float | ndarray[tuple[int, ...], dtype[_ScalarType_co]]]¶: Fitness history (scalars, or per-sub-agent rows for multi-agent).

get_action(obs: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | dict[str, ndarray[tuple[int, ...], dtype[_ScalarType_co]]] | tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ...] | Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor] | Number | list[RolloutPrompt] | RolloutPrompt, action_mask: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | Sequence[ndarray[tuple[int, ...], dtype[_ScalarType_co]]] | Tensor | None = None, *, hidden_state: dict[str, Tensor], **kwargs: Any) → tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], dict[str, Tensor] | None]¶

get_action(obs: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | dict[str, ndarray[tuple[int, ...], dtype[_ScalarType_co]]] | tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ...] | Tensor | TensorDict | tuple[Tensor, ...] | dict[str, Tensor] | Number | list[RolloutPrompt] | RolloutPrompt, action_mask: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | Sequence[ndarray[tuple[int, ...], dtype[_ScalarType_co]]] | Tensor | None = None, hidden_state: None = None, *args: Any, **kwargs: Any) → tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]]]

Return the next action to take in the environment.

Parameters:

obs (ObservationType) – Environment observation, or multiple observations in a batch
action_mask (ActionMaskInput) – Mask of legal actions 1=legal 0=illegal, defaults to None
hidden_state (dict[str, torch.Tensor] | None) – Hidden state for recurrent policies, defaults to None

Returns:

Action, log probability, entropy, state values, and (if recurrent) next hidden state

Return type:

tuple[npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray, dict[str, torch.Tensor] | None] | tuple[npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray]

Return the dimension of the action space as it pertains to the underlying networks (i.e. the output size of the networks).

Parameters:: action_space (spaces.Space or list[spaces.Space].) – The action space of the environment.
Returns:: The dimension of the action space.
Return type:: int | dict[str, int] | tuple[int | dict[str, int], …]

get_eval_modules(cloning: bool = True) → tuple[dict[str, EvolvableModule], dict[str, EvolvableModule]]¶

Get the offsprings of all of the evaluation modules in the individual.

Parameters:: cloning (bool, optional) – Whether to clone each evaluation module before returning it, defaults to True.
Returns:: Tuple of offspring policy and the rest of the evaluation modules
Return type:: tuple[dict[str, EvolvableModule], dict[str, EvolvableModule]]

get_hidden_state_architecture() → dict[str, tuple[int, ...]]¶

Get the hidden state architecture for the environment.

Returns:: Dictionary describing the hidden state architecture (name to (num_layers, num_envs, hidden_size) shape)
Return type:: dict[str, tuple[int, …]]

get_initial_hidden_state(num_envs: int = 1) → dict[str, Tensor]¶

Get the initial hidden state for the environment.

The hidden states are generally cached on a per Module basis. The reason the Cache is per Module is because the user might want to have a custom initialization for the hidden states.

Parameters:: num_envs (int, optional) – Number of environments, defaults to 1
Returns:: Initial hidden state dictionary
Return type:: dict[str, torch.Tensor]

get_lr_names() → list[str | tuple[str, str]]¶: Return the learning-rate attribute name(s) of each optimizer.

get_policy() → EvolvableModuleProtocol¶: Return the policy network of the algorithm.

Return the dimension of the state space as it pertains to the underlying networks (i.e. the input size of the networks).

Parameters:: observation_space (spaces.Space or list[spaces.Space].) – The observation space of the environment.
Returns:: The dimension of the state space.
Return type:: tuple[int, …] | dict[str, tuple[int, …]]

property hp_config: HyperparameterConfig¶: Return the hyperparameter configuration for Evo-HPO mutations.

property index: int¶: Return the index of the algorithm.

init_training_step(capture_grama: bool = False) → None¶

Open the agent’s training block: metrics tracking, and GraMa capture.

Hooks are registered afresh each cycle, so they follow the agent through architecture mutations, checkpoint reloads and accelerator re-wrapping. Opening a block implicitly closes one that an earlier call left open.

Parameters:: capture_grama (bool) – Whether to register GraMa capture hooks for this training step. Defaults to False since the LLM finetuners never run ReGraMa.
Returns:: None.
Return type:: None

static inspect_attributes(agent: EvolvableAlgorithmProtocol | AgentWrapperProtocol[Any], input_args_only: bool = False, exclude: Iterable[str] = ()) → dict[str, Any]¶

Inspect and retrieve the attributes of the current object, excluding attributes related to the underlying evolvable networks (i.e. EvolvableModule, torch.optim.Optimizer) and with an option to include only the attributes that are input arguments to the constructor.

Parameters:

input_args_only (bool) – If True, only include attributes that are input arguments to the constructor. Defaults to False.
exclude (Iterable[str], optional) – Extra attribute names to drop from the result, on top of the standard exclusions below. For a caller-specific reason to leave an attribute out of its own view.

Returns:

A dictionary of attribute names and their values.

Return type:

dict[str, Any]

learn(experiences: TensorDict | None = None) → float¶

Update agent network parameters to learn from experiences.

Parameters:: experiences (TensorDict | None) – Optional pre-collected rollout batch. When None (the default), samples are drawn from the agent’s internal rollout buffer.
Returns:: Mean loss value from training.
Return type:: float

classmethod load(path: str, device: str | device = 'cpu', accelerator: Accelerator | None = None) → Self¶

Load an algorithm from a checkpoint.

Parameters:

path (string) – Location to load checkpoint from.
device (str, optional) – Device to load the algorithm on, defaults to ‘cpu’
accelerator (Accelerator | None, optional) – Accelerator object for distributed computing, defaults to None

Returns:

An instance of the algorithm

Return type:

SingleAgentAlgorithm

load_checkpoint(path: str) → None¶

Load saved agent properties and network weights from checkpoint.

Restores full training state (weights, optimizer, LR schedule, hyperparameters) to resume a run; load_weights() takes weights only.

Parameters:: path (string) – Location to load checkpoint from

load_weights(path: str) → None¶

Load only the network weights from a checkpoint.

Warm-starts a new run from prior weights; optimizer, LR schedule, training progress and hyperparameters are not loaded.

Parameters:: path (string) – Location to load checkpoint from

property mut: str | None¶: Return the mutation object of the algorithm.

mutation_hook() → None¶: Execute the hooks registered with the algorithm.

Create a population of algorithms.

Parameters:

size (int) – The size of the population.
observation_space (GymSpaceType) – The observation space.
action_space (GymSpaceType) – The action space.
device (DeviceType) – Torch device. Defaults to "cpu".
wrapper_cls (type | None) – Optional wrapper class to apply to each agent.
wrapper_kwargs (dict[str, Any] | None) – Keyword arguments for the wrapper class.
resume_from_checkpoint (str | None) – Path to checkpoint to resume from.
kwargs (Any) – Additional keyword arguments to pass to the algorithm constructor.

Returns:

A list of algorithms.

Return type:

list[SingleAgentAlgorithm]

Preprocesses observations for forward pass through neural network.

Parameters:: observation (ObservationType) – Observations of environment
Returns:: Preprocessed observations
Return type:: torch.Tensor[float] or dict[str, torch.Tensor[float]] or tuple[torch.Tensor[float], …]

recompile() → None¶: Recompiles the evolvable modules in the algorithm with the specified torch compiler.

register_mutation_hook(hook: LambdaType | MethodType) → None¶

Parameters:: hook (MutationHook) – The hook to be executed after mutation.

register_network_group(group: NetworkGroup) → None¶

Set the evaluation network for the algorithm.

Parameters:: name (str) – The name of the evaluation network.

reinit_optimizers(optimizer: OptimizerConfig | None = None) → None¶

Reinitialize the optimizers of an algorithm. If no optimizer is passed, all optimizers are reinitialized.

Parameters:: optimizer (OptimizerConfig | None, optional) – The optimizer to reinitialize, defaults to None, in which case all optimizers are reinitialized.

save_checkpoint(path: str) → None¶

Save a checkpoint of agent properties and network weights to path.

Parameters:: path (string) – Location to save checkpoint at

property scores: list[float | list[float]]¶: Per-episode scores (per-group score rows for multi-agent metrics).

set_training_mode(training: bool) → None¶

Set the training mode of the algorithm.

Parameters:: training (bool) – If True, set the algorithm to training mode.

share_encoder_parameters() → None¶: Shares the encoder parameters between the actor and critic.

property steps: int¶: Cumulative global step count.

test(env: Env | VectorEnv, max_steps: int | None = None, loop: int = 3, vectorized: bool = True, callback: Callable[[float, dict[str, Any]], None] | None = None) → float¶

Return mean test score of agent in environment with epsilon-greedy policy.

Parameters:

env (gym.Env | gym.vector.VectorEnv) – The environment to be tested in
max_steps (int, optional) – Maximum number of testing steps, defaults to None
loop (int, optional) – Number of testing loops/episodes to complete. The returned score is the mean. Defaults to 3
vectorized (bool, optional) – Whether the environment is vectorized, defaults to True
callback (Callable[[float, dict[str, Any]], None] | None) – Optional callback function that takes the sum of rewards and the last info dictionary as input, defaults to None

Returns:

Mean test score of agent in environment

Return type:

float

Move experiences to the device.

Parameters:: experiences (tuple[torch.Tensor[float], ...]) – Experiences to move to device
Returns:: Experiences on the device
Return type:: tuple[torch.Tensor[float], …]

unrolled_eval_networks() → list[tuple[str | None, Module]]¶

Return the agent’s evaluation networks as (network_id, network) pairs.

Returns:: One (network_id, network) pair per measured network.
Return type:: list[tuple[str | None, torch.nn.Module]]

unwrap_models() → None¶: Unwraps the models in the algorithm from the accelerator.

wrap_models() → None¶: Wrap the models in the algorithm with the accelerator.