On-Policy Training

In online reinforcement learning, an agent is able to gather data by directly interacting with its environment. It can then use this experience to learn from and update its policy. To enable our agent to interact in this way, the agent needs to act either in the real world, or in a simulation.

AgileRL’s online training framework enables agents to learn in environments, using the standard Gym interface, 10x faster than SOTA by using our Evolutionary Hyperparameter Optimization algorithm.

On-policy reinforcement learning involves learning from experiences gathered by following a single policy. In these algorithms, the data collection policy and the learning policy are the same, meaning that the agent learns from its own actions and their outcomes in the environment. This approach often leads to more stable learning as the agent directly interacts with the environment based on its current policy, continuously updating and improving it. However, on-policy algorithms can be less sample-efficient compared to off-policy methods since they are constrained to learn from the data generated by the current policy, potentially limiting exploration and the use of past experiences.

Algorithms	Tutorials
PPO	Acrobot Masked Velocity LunarLander-v3

Training with LocalTrainer

The recommended way to train a fully-customised on-policy agent is through a YAML manifest and the LocalTrainer. This handles population creation, rollout collection, evolutionary HPO, and the training loop automatically.

Here is an example manifest to train PPO on LunarLander-v3:

ppo.yaml

algorithm:
  name: PPO
  batch_size: 128
  lr: 0.001
  learn_step: 2048
  gamma: 0.99
  gae_lambda: 0.95
  action_std_init: 0.6
  clip_coef: 0.2
  ent_coef: 0.01
  vf_coef: 0.5
  max_grad_norm: 0.5
  update_epochs: 4

environment:
  name: LunarLander-v3
  num_envs: 16

training:
  max_steps: 6_000_000
  target_score: 250.0
  pop_size: 4
  evo_steps: 10_240

network:
  latent_dim: 64
  encoder_config:
    hidden_size: [64]
    activation: ReLU
    layer_norm: true
  head_config:
    hidden_size: [64]
    activation: ReLU
    output_vanish: true
    layer_norm: true

mutation:
  probabilities:
    no_mut: 0.4
    arch_mut: 0.2
    new_layer: 0.2
    params_mut: 0.2
    rl_hp_mut: 0.2
  rl_hp_selection:
    lr:
      min: 0.0001
      max: 0.01
    batch_size:
      min: 8
      max: 1024
    learn_step:
      min: 256
      max: 8192
    ent_coef:
      min: 0.001
      max: 0.1
  mutation_sd: 0.1
  rand_seed: 42

tournament_selection:
  tournament_size: 2
  elitism: true

Python

from agilerl import LocalTrainer

trainer = LocalTrainer.from_manifest("ppo.yaml")
population, fitnesses = trainer.train()

CLI

python -m agilerl.train ppo.yaml

See also

Trainers for full manifest reference and additional options.

Population Creation

To perform evolutionary HPO, we require a population of agents. Individuals in this population learn individually, allowing us to determine the efficacy of certain hyperparameters. Individual agents which learn best are more likely to survive until the next generation, and so their hyperparameters are more likely to remain present in the population. The sequence of evolution (tournament selection followed by mutation) is detailed further below.

import torch
from agilerl.algorithms import PPO
from agilerl.utils.utils import make_vect_envs

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

num_envs = 16
env = make_vect_envs("LunarLander-v3", num_envs=num_envs)  # Create environment

observation_space = env.single_observation_space
action_space = env.single_action_space

# Configure network architecture
net_config = {
    "encoder_config": {"hidden_size": [64, 64]},
    "head_config": {"hidden_size": [64]},
}

# Algorithm hyperparameters
init_hp = {
    "batch_size": 128,
    "lr": 1e-3,
    "learn_step": 128,
    "gamma": 0.99,
    "gae_lambda": 0.95,
    "action_std_init": 0.6,
    "clip_coef": 0.2,
    "ent_coef": 0.01,
    "vf_coef": 0.5,
    "max_grad_norm": 0.5,
    "target_kl": None,
    "update_epochs": 4,
    "num_envs": num_envs,
}

# Initialize population
population_size = 6
pop = PPO.population(
    size=population_size,
    observation_space=observation_space,
    action_space=action_space,
    net_config=net_config,
    device=device,
    **init_hp,
)

Evolutionary HPO

Tournament selection is used to select the agents from a population which will make up the next generation of agents. If elitism is used, the best agent from a population is automatically preserved and becomes a member of the next generation. Mutation is periodically used to explore the hyperparameter space.

from agilerl.hpo.mutation import Mutations
from agilerl.hpo.tournament import TournamentSelection

tournament = TournamentSelection(
    tournament_size=2,  # Tournament selection size
    elitism=True,  # Elitism in tournament selection
    population_size=6,  # Population size
)

mutations = Mutations(
    no_mutation=0.4,  # No mutation
    architecture=0.2,  # Architecture mutation
    new_layer_prob=0.2,  # New layer mutation
    parameters=0.2,  # Network parameters mutation
    activation=0,  # Activation layer mutation
    rl_hp=0.2,  # Learning HP mutation
    mutation_sd=0.1,  # Mutation strength
    rand_seed=1,  # Random seed
    device=device,
)

See also

Evolutionary Hyperparameter Optimization for details on how evolutionary HPO works.

Training Loop

While off-policy RL algorithms can be considered more sample-efficient than on-policy algorithms, due to their ability to learn from experiences collected using a different or previous policy, on-policy algorithms often do better in practice due to the improved stability during training. Currently, AgileRL includes an evolvable implementation of Proximal Policy Optimisation (PPO). This algorithm can be used in a variety of settings and is widely popular across domains including robotics, games, finance, and RLHF.

The setup for PPO is very similar to the off-policy example above, except it does not require the use of an experience replay buffer. It also requires some different hyperparameters, shown below in the custom loop.

You can use our off-the-shelf on-policy training function to train a population of agents using PPO:

from agilerl.training.train_on_policy import train_on_policy

trained_pop, pop_fitnesses = train_on_policy(
    env=env,                              # Gym-style environment
    env_name="LunarLander-v3",  # Environment name
    pop=agent_pop,  # Population of agents
    max_steps=200000,  # Max number of training steps
    evo_steps=10000,  # Evolution frequency
    eval_steps=None,  # Number of steps in evaluation episode
    eval_loop=1,  # Number of evaluation episodes
    target=200.,  # Target score for early stopping
    tournament=tournament,  # Tournament selection object
    mutation=mutations,  # Mutations object
    wb=True,  # Weights and Biases tracking
)

Note

Known Gymnasium issue - running vectorize environments as top-level code (without if __name__ == "__main__":) may cause multiprocessing errors. To fix, run the above as a method under main, e.g.

def train_agent():
    # ... training code

if __name__ == "__main__":
    train_agent()

Alternatively, use a custom on-policy training loop:

Example Custom Training Loop

import numpy as np
import torch
from tqdm import trange
from agilerl.algorithms import PPO
from agilerl.hpo.mutation import Mutations
from agilerl.hpo.tournament import TournamentSelection
from agilerl.rollouts.on_policy import collect_rollouts
from agilerl.utils.utils import make_vect_envs, default_progress_bar

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

num_envs = 16
env = make_vect_envs("LunarLander-v3", num_envs=num_envs)  # Create environment
observation_space = env.single_observation_space
action_space = env.single_action_space

# Configure network architecture
net_config = {
    "encoder_config": {
        "hidden_size": [32, 32], # Encoder hidden size
        "activation": "ReLU"
        },
    "head_config": {
        "hidden_size": [32]  # Head hidden size
    }
}

# Algorithm hyperparameters
init_hp = {
    "batch_size": 128,
    "lr": 1e-3,
    "learn_step": 128,
    "gamma": 0.99,
    "gae_lambda": 0.95,
    "action_std_init": 0.6,
    "clip_coef": 0.2,
    "ent_coef": 0.01,
    "vf_coef": 0.5,
    "max_grad_norm": 0.5,
    "target_kl": None,
    "update_epochs": 4,
    "num_envs": num_envs,
}

# Initialize population
pop = PPO.population(
    size=6,
    observation_space=observation_space,
    action_space=action_space,
    net_config=net_config,
    device=device,
    **init_hp,
)

tournament = TournamentSelection(
    tournament_size=2,  # Tournament selection size
    elitism=True,  # Elitism in tournament selection
    population_size=6,  # Population size
)

mutations = Mutations(
    no_mutation=0.4,  # No mutation
    architecture=0.2,  # Architecture mutation
    new_layer_prob=0.2,  # New layer mutation
    parameters=0.2,  # Network parameters mutation
    activation=0,  # Activation layer mutation
    rl_hp=0.2,  # Learning HP mutation
    mutation_sd=0.1,  # Mutation strength  # Network architecture
    rand_seed=1,  # Random seed
    device=device,
)

max_steps = 200000  # Max steps
evo_steps = 10000  # Evolution frequency
eval_steps = None  # Evaluation steps per episode - go until done
eval_loop = 1  # Number of evaluation episodes
total_steps = 0

# TRAINING LOOP
pbar = default_progress_bar(max_steps)
while np.less([agent.steps for agent in pop], max_steps).all():
    pop_episode_scores = []
    for agent in pop:  # Loop through population
        agent.set_training_mode(True)

        completed_episode_scores = []
        steps = 0

        for _ in range(-(evo_steps // -agent.learn_step)):
            # Collect rollouts and save in the agent's rollout buffer
            episode_scores = collect_rollouts(agent, env)

            agent.learn()  # Learn from rollout buffer

            # Update step counter and scores
            total_steps += agent.learn_step
            steps += agent.learn_step
            agent.steps += agent.learn_step
            completed_episode_scores += episode_scores

        pbar.update(steps // len(pop))
        pop_episode_scores.append(completed_episode_scores)

    # Evaluate population
    fitnesses = [
        agent.test(
            env,
            max_steps=eval_steps,
            loop=eval_loop,
        )
        for agent in pop
    ]
    mean_scores = [
        (
            np.mean(episode_scores)
            if len(episode_scores) > 0
            else "0 completed episodes"
        )
        for episode_scores in pop_episode_scores
    ]

    pbar.write(
        f"--- Global steps {total_steps} ---\n"
        f"Steps: {[agent.steps for agent in pop]}\n"
        f"Scores: {mean_scores}\n"
        f"Fitnesses: {['%.2f' % fitness for fitness in fitnesses]}\n"
        f"5 fitness avgs: {['%.2f' % np.mean(agent.fitness[-5:]) for agent in pop]}\n"
    )

    # Tournament selection and population mutation
    elite, pop = tournament.select(pop)
    pop = mutations.mutation(pop)

pbar.close()
env.close()

Training Loop for Recurrent On-Policy Algorithms

Recurrent on-policy algorithms require a different training loop to the standard on-policy algorithms. This is because the agent needs to maintain a hidden state between steps, which is not possible with the standard training loop. AgileRL currently supports recurrent policies to be used with PPO. To use a recurrent policy, users must set recurrent=True when creating the algorithm.

End-to-end example: Recurrent PPO on LunarLander-v3

import torch
from agilerl.algorithms import PPO
from agilerl.rollouts.on_policy import collect_rollouts_recurrent
from agilerl.utils.utils import make_vect_envs, default_progress_bar

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create environment
num_envs = 16
env = make_vect_envs("LunarLander-v3", num_envs=num_envs)

observation_space = env.single_observation_space
action_space = env.single_action_space

# Configure network architecture
net_config = {
    "encoder_config": {
        "hidden_state_size": 64,
        "num_layers": 1,
    },
    "head_config": {
        "hidden_size": [64],
    }
}

# Algorithm hyperparameters
init_hp = {
    "batch_size": 128,
    "lr": 1e-3,
    "learn_step": 128,
    "gamma": 0.99,
    "gae_lambda": 0.95,
    "action_std_init": 0.6,
    "clip_coef": 0.2,
    "ent_coef": 0.01,
    "vf_coef": 0.5,
    "max_grad_norm": 0.5,
    "recurrent": True,
    "max_seq_len": 512,
    "target_kl": None,
    "update_epochs": 4,
    "num_envs": num_envs,
}

# Initialize population
pop = PPO.population(
    size=6,
    observation_space=observation_space,
    action_space=action_space,
    net_config=net_config,
    device=device,
    **init_hp,
)

tournament = TournamentSelection(
    tournament_size=2,  # Tournament selection size
    elitism=True,  # Elitism in tournament selection
    population_size=6,  # Population size
)

mutations = Mutations(
    no_mutation=0.4,  # No mutation
    architecture=0.2,  # Architecture mutation
    new_layer_prob=0.2,  # New layer mutation
    parameters=0.2,  # Network parameters mutation
    activation=0,  # Activation layer mutation
    rl_hp=0.2,  # Learning HP mutation
    mutation_sd=0.1,  # Mutation strength  # Network architecture
    rand_seed=1,  # Random seed
    device=device,
)

max_steps = 200000  # Max steps
evo_steps = 10000  # Evolution frequency
eval_steps = None  # Evaluation steps per episode - go until done
eval_loop = 1  # Number of evaluation episodes
total_steps = 0

# TRAINING LOOP
pbar = default_progress_bar(max_steps)
while np.less([agent.steps for agent in pop], max_steps).all():
    pop_episode_scores = []
    for agent in pop:  # Loop through population
        steps = 0
        completed_episodes = []
        last_obs, last_done, last_scores, last_info = None, None, None, None
        for _ in range(-(evo_steps // -agent.learn_step)):
            # Collect rollouts and save in buffer
            episode_scores, last_obs, last_done, last_scores, last_info = (
                collect_rollouts_recurrent(
                    agent,
                    env,
                    last_obs=last_obs,
                    last_done=last_done,
                    last_scores=last_scores,
                    last_info=last_info,
                )
            )

            agent.learn()  # Learn from rollout buffer

            # Update step counter and scores
            total_steps += agent.learn_step
            steps += agent.learn_step
            agent.steps += agent.learn_step
            completed_episodes += episode_scores

        pop_episode_scores.append(
            np.mean(completed_episodes)
            if len(completed_episodes) > 0
            else "0 completed episodes"
        )
        pbar.update(steps // len(pop))

    # Evaluate population
    fitnesses = [
        agent.test(
            env,
            max_steps=eval_steps,
            loop=eval_loop,
        )
        for agent in pop
    ]

    pbar.write(
        f"--- Global steps {total_steps} ---\n"
        f"Steps: {[agent.steps for agent in pop]}\n"
        f"Scores: {pop_episode_scores}\n"
        f"Fitnesses: {['%.2f' % fitness for fitness in fitnesses]}\n"
        f"5 fitness avgs: {['%.2f' % np.mean(agent.fitness[-5:]) for agent in pop]}\n"
    )

    if any(score >= required_score for score in pop_episode_scores):
        print(
            f"\nAgent achieved required score {required_score}. Stopping training."
        )
        elite, _ = tournament.select(pop)
        break

    # Tournament selection and population mutation
    elite, pop = tournament.select(pop)
    pop = mutations.mutation(pop)

pbar.close()
env.close()

Tutorial

Acrobot with PPO

Evolve a PPO population on Acrobot-v1.

Partially Observable Pendulum-v1 with Recurrent PPO

Recurrent PPO on masked-velocity LunarLander-v3.