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

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. The OnPolicyWorkflow class implements the generic training loop shared by PPO, A2C, and IMPALA:

• Rollout collection via 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

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
• PPOAgent in evorl/algorithms/ppo.py
• IMPALAAgent in 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 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) 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 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) 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 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) 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 collects trajectories using agent.compute_actions
3. Statistics Update – Optional observation normalization updates
4. Advantage Computation –
   • compute_gae in evorl/utils/rl_toolkits.py for PPO/A2C
   • compute_vtrace and compute_pg_advantage in evorl/algorithms/impala.py for IMPALA
5. Loss Calculation – agent.loss() computes gradients
6. Optimization – agent_gradient_update in 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 uses Hydra for configuration. Create a YAML config (e.g., configs/ppo_cartpole.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:

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:

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:

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 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, A2C uses standard actor-critic loss in evorl/algorithms/a2c.py, and IMPALA uses V-Trace correction in evorl/algorithms/impala.py.
• Distributed training is supported via agent_gradient_update in evorl/distributed/__init__.py and automatic multi-device scaling through the workflow configuration.
• Entry point 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 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 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 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 (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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}