# How to Implement On-Policy Algorithms (PPO, A2C, IMPALA) in EvoRL

> Implement PPO A2C and IMPALA on policy algorithms in EvoRL easily. Discover how EvoRL's modular workflow simplifies training pipelines with custom agent factories and loss functions.

- Repository: [EMI-Group/evorl](https://github.com/emi-group/evorl)
- Tags: how-to-guide
- Published: 2026-03-01

---

**EvoRL provides a modular architecture where PPO, A2C, and IMPALA share a common `OnPolicyWorkflow` base class, requiring only algorithm-specific agent factories and loss functions to implement complete training pipelines.**

EvoRL is a JAX-based framework designed for scalable evolutionary and reinforcement learning research. When you implement on-policy algorithms like PPO, A2C, and IMPALA in EvoRL, you leverage a unified workflow system that handles distributed rollouts, advantage estimation, and optimization while allowing each algorithm to customize its loss computation and value estimation strategy.

## Core Architecture of On-Policy Algorithms in EvoRL

### The OnPolicyWorkflow Base Class

The foundation for all on-policy implementations resides in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py). The `OnPolicyWorkflow` class implements the generic training loop shared by PPO, A2C, and IMPALA:

- **Rollout collection** via [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py)
- **Observation normalization** updates (optional)
- **Advantage estimation** (GAE for PPO/A2C, V-Trace for IMPALA)
- **Loss computation** and **optimizer steps** via `agent_gradient_update` in [`evorl/distributed/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/__init__.py)

Algorithm-specific workflows (`PPOWorkflow`, `A2CWorkflow`, `IMPALAWorkflow`) inherit from this base and only provide a `name()` method and a builder that creates the appropriate agent.

### Agent Structure and Network Factories

Each algorithm defines an `Agent` subclass stored in its respective file:

- `A2CAgent` in [`evorl/algorithms/a2c.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/a2c.py)
- `PPOAgent` in [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py)
- `IMPALAAgent` in [`evorl/algorithms/impala.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/impala.py)

These agents encapsulate:
- **Policy and value networks** (built via `make_policy_network` and `make_v_network` in `evorl/networks/`)
- **Algorithm-specific hyperparameters** (clip epsilon for PPO, V-Trace parameters for IMPALA)
- **Loss functions** (`loss` method) and **action computation** (`compute_actions`)

Factory functions `make_mlp_ppo_agent`, `make_mlp_a2c_agent`, and `make_mlp_impala_agent` provide convenient construction with configurable hidden layer sizes and optional observation normalization.

## Implementing PPO, A2C, and IMPALA

### PPO (Proximal Policy Optimization)

PPO implementation in [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) uses the clipped surrogate objective. The `PPOAgent.loss` method computes:

1. **Policy loss**: Clipped importance sampling ratio with epsilon (default 0.2)
2. **Value loss**: MSE between predicted and target values
3. **Entropy bonus**: Encourages exploration
4. **Optional KL penalty**: For early stopping or constraint

The `PPOWorkflow` class (line 61 in [`ppo.py`](https://github.com/emi-group/evorl/blob/main/ppo.py)) inherits from `OnPolicyWorkflow` and uses `make_mlp_ppo_agent` to construct the agent with GAE advantage estimation configured in the workflow base.

### A2C (Advantage Actor-Critic)

A2C in [`evorl/algorithms/a2c.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/a2c.py) provides the simplest on-policy implementation. The `A2CAgent` computes standard actor-critic loss without importance sampling corrections.

Key characteristics:
- **Single-step or n-step returns** (configurable)
- **GAE support** via the base workflow
- **Synchronous updates** (no replay buffer)

The `A2CWorkflow` (line 29 in [`a2c.py`](https://github.com/emi-group/evorl/blob/main/a2c.py)) is a minimal subclass that instantiates `make_mlp_a2c_agent` and relies entirely on the base `OnPolicyWorkflow` for training logic.

### IMPALA (Importance Weighted Actor-Learner Architecture)

IMPALA implementation in [`evorl/algorithms/impala.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/impala.py) introduces **V-Trace** correction for off-policy actor-critic learning with distributed actors. Unlike PPO and A2C, IMPALA uses:

- **V-Trace targets** (`compute_vtrace` function) to correct for policy lag between actors and learners
- **Policy gradient advantages** (`compute_pg_advantage`) for stable learning
- **Decoupled actor-learner** architecture support via the distributed utilities

The `IMPALAAgent` loss combines:
1. **V-Trace corrected policy loss**
2. **Baseline loss** (value function)
3. **Entropy regularization**
4. **KL divergence penalty**

The `IMPALAWorkflow` (line 91 in [`impala.py`](https://github.com/emi-group/evorl/blob/main/impala.py)) configures the agent using `make_mlp_impala_agent` and ensures the workflow uses the correct V-Trace advantage computation path in the base class.

## Data Flow and Training Loop

Understanding the data flow helps when customizing these algorithms. The `OnPolicyWorkflow` orchestrates:

1. **Environment Creation** – `evorl.envs.create_env` instantiates vectorized environments
2. **Rollout Phase** – [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py) collects trajectories using `agent.compute_actions`
3. **Statistics Update** – Optional observation normalization updates
4. **Advantage Computation** – 
   - `compute_gae` in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) for PPO/A2C
   - `compute_vtrace` and `compute_pg_advantage` in [`evorl/algorithms/impala.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/impala.py) for IMPALA
5. **Loss Calculation** – `agent.loss()` computes gradients
6. **Optimization** – `agent_gradient_update` in [`evorl/distributed/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/__init__.py) applies updates
7. **Metrics & Checkpointing** – `RLWorkflow` handles logging via recorders and periodic checkpointing

This unified pipeline means switching between PPO, A2C, and IMPALA requires only changing the agent factory and workflow class in your configuration.

## Practical Implementation Examples

### Training PPO with a Configuration File

The standard entry point [`scripts/train.py`](https://github.com/emi-group/evorl/blob/main/scripts/train.py) uses Hydra for configuration. Create a YAML config (e.g., [`configs/ppo_cartpole.yaml`](https://github.com/emi-group/evorl/blob/main/configs/ppo_cartpole.yaml)):

```yaml
defaults:
  - _self_
  - override hydra/launcher: joblib

workflow_cls: evorl.workflows.PPOWorkflow
seed: 42

env:
  env_name: CartPole-v1
  env_type: gymnax

agent:
  actor_hidden_layer_sizes: [256, 256]
  critic_hidden_layer_sizes: [256, 256]
  clip_epsilon: 0.2
  normalize_obs: true

training:
  num_iterations: 1000
  rollout_length: 128
  num_envs: 8
  learning_rate: 3e-4

```

Run with:

```bash
python scripts/train.py --config-name=ppo_cartpole

```

### Direct Agent Instantiation for Custom Experiments

For research requiring custom network architectures or loss modifications, instantiate agents directly:

```python
import jax
from evorl.algorithms.ppo import make_mlp_ppo_agent
from evorl.envs import create_env

# Setup environment

env = create_env(
    dict(env_name="CartPole-v1", env_type="gymnax"),
    episode_length=200,
    parallel=8
)

# Build PPO agent with custom architecture

agent = make_mlp_ppo_agent(
    action_space=env.action_space,
    clip_epsilon=0.2,
    actor_hidden_layer_sizes=(256, 256),
    critic_hidden_layer_sizes=(256, 256),
    normalize_obs=True,
    normalize_gae=True,
)

# Access network parameters

key = jax.random.PRNGKey(0)
agent_state = agent.init(key)
print(f"Policy network initialized with parameters: {agent_state.params.keys()}")

```

### Running Distributed IMPALA

IMPALA benefits from distributed training. The workflow automatically handles device parallelism:

```python
from evorl.algorithms.impala import IMPALAWorkflow
import jax

# Configuration should specify IMPALA-specific parameters like:

# - vtrace_rho_bar: 1.0

# - vtrace_c_bar: 1.0

# - baseline_cost: 0.5

cfg = {
    "workflow_cls": "evorl.algorithms.impala.IMPALAWorkflow",
    "env": {"env_name": "Breakout", "env_type": "gymnax"},
    "agent": {
        "actor_hidden_layer_sizes": [256, 256],
        "vtrace_rho_bar": 1.0,
        "vtrace_c_bar": 1.0
    },
    "training": {
        "num_iterations": 10000,
        "rollout_length": 20,
        "num_envs": 32
    }
}

# Build workflow with multi-device support

workflow = IMPALAWorkflow.build_from_config(cfg, enable_multi_devices=True)

# Initialize and train

state = workflow.init(jax.random.PRNGKey(cfg["training"].get("seed", 42)))
final_state = workflow.learn(state)

```

## Summary

- **EvoRL** provides a unified `OnPolicyWorkflow` base class in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) that handles the complete training loop for PPO, A2C, and IMPALA.
- **Agent factories** (`make_mlp_ppo_agent`, `make_mlp_a2c_agent`, `make_mlp_impala_agent`) in their respective algorithm files construct policy/value networks with configurable architectures.
- **Algorithm-specific losses** differentiate implementations: PPO uses clipped surrogate loss in [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py), A2C uses standard actor-critic loss in [`evorl/algorithms/a2c.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/a2c.py), and IMPALA uses V-Trace correction in [`evorl/algorithms/impala.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/impala.py).
- **Distributed training** is supported via `agent_gradient_update` in [`evorl/distributed/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/__init__.py) and automatic multi-device scaling through the workflow configuration.
- **Entry point** [`scripts/train.py`](https://github.com/emi-group/evorl/blob/main/scripts/train.py) uses Hydra configurations to instantiate workflows, making it easy to switch between algorithms by changing the `workflow_cls` parameter.

## Frequently Asked Questions

### What is the difference between PPO and IMPALA implementations in EvoRL?

The primary difference lies in the **advantage estimation and policy correction mechanisms**. PPO uses **Generalized Advantage Estimation (GAE)** via `compute_gae` in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) and a **clipped surrogate objective** to limit policy updates. IMPALA uses **V-Trace** correction via `compute_vtrace` and `compute_pg_advantage` in [`evorl/algorithms/impala.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/impala.py) to handle off-policy data from distributed actors, making it suitable for decoupled actor-learner architectures where actors lag behind the learner.

### How do I switch from PPO to A2C in my training script?

Switching requires only two configuration changes in your Hydra config or Python script. First, change the `workflow_cls` parameter from `evorl.workflows.PPOWorkflow` (or `evorl.algorithms.ppo.PPOWorkflow`) to `evorl.workflows.A2CWorkflow` (or `evorl.algorithms.a2c.A2CWorkflow`). Second, ensure your agent configuration uses the A2C-specific factory by setting the appropriate agent parameters (removing PPO-specific fields like `clip_epsilon`). The base `OnPolicyWorkflow` handles the rest of the training loop automatically.

### Can I use custom neural network architectures with these algorithms?

Yes, you can replace the default MLP architectures by modifying the network factories or manually constructing agents. Instead of using `make_mlp_ppo_agent`, you can build custom policy networks using `make_policy_network` and value networks using `make_v_network` from `evorl/networks/`, then pass these directly to the `Agent` constructor (e.g., `PPOAgent`). This allows you to use CNNs for visual observations, custom activation functions, or specialized layer configurations while retaining the algorithmic logic of PPO, A2C, or IMPALA.

### Where is the distributed training logic handled for IMPALA?

Distributed training logic is centralized in [`evorl/distributed/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/__init__.py) through the `agent_gradient_update` function, which handles gradient aggregation across devices using `psum` for distributed sums and `all_reduce` for metrics. For IMPALA specifically, the V-Trace correction in [`evorl/algorithms/impala.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/impala.py) (`compute_vtrace` and `compute_pg_advantage`) handles the off-policy corrections necessary when actors and learners operate with different policy parameters. The workflow supports multi-device training via the `enable_multi_devices=True` flag in `build_from_config`, which automatically scales the rollout and update steps across available JAX devices.