# How to Implement a Custom Workflow Class for Training Loop Logic in EvoRL

> Learn to implement a custom Workflow class for training loop logic in EvoRL. Subclass base workflows, implement factory methods, and register your class. Optimize your RL training.

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

---

**To implement a custom Workflow class in EvoRL, subclass a base workflow such as `RLWorkflow` or `ECWorkflow`, implement the `build_from_config` factory method along with `setup` and `step`, and register the class by importing it in the algorithms package.**

EvoRL is an open-source evolutionary reinforcement learning framework that cleanly separates algorithmic logic from training orchestration through a workflow abstraction. Creating a custom Workflow class allows you to define specialized training loop logic while inheriting JIT compilation, multi-device distribution, and automatic checkpointing from the framework's base implementations.

## Understanding the Workflow Hierarchy in EvoRL

The workflow system is defined in [`evorl/workflows/workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/workflow.py). At the top sits `AbstractWorkflow`, which declares the core contract that every training loop must fulfill:

- `init(self, key: chex.PRNGKey) -> State`: Initializes the workflow state.
- `step(self, state: State) -> tuple[Any, State]`: Executes one training iteration.
- `name(cls) -> str`: Returns the workflow's registered identifier.

The `Workflow` class inherits from `AbstractWorkflow` and implements generic infrastructure including the `Recorder` for logging, checkpoint management via `setup_checkpoint_manager`, and the high-level `learn` method that orchestrates the training loop. Subclasses must override `setup` to create the initial `State` and `step` to define the training logic.

## Choosing the Right Base Class for Your Custom Workflow

EvoRL provides specialized base classes in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) and [`evorl/workflows/ec_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/ec_workflow.py) that handle common patterns:

### RLWorkflow for Reinforcement Learning

`RLWorkflow` is the foundation for model-free RL algorithms. It provides:

- **Multi-device support**: Automatic `pmap` handling via `enable_multi_devices`.
- **JIT compilation**: Toggle via `enable_jit`.
- **Common setup**: Standardized initialization for environment-agent-optimizer bundles.

### OnPolicyWorkflow and OffPolicyWorkflow

These subclasses of `RLWorkflow` provide ready-made templates:

- **OnPolicyWorkflow**: For algorithms like PPO that collect fresh trajectories each iteration. Implements `setup` and `evaluate`; you only override `step`.
- **OffPolicyWorkflow**: For algorithms like TD3 or SAC that use replay buffers. Manages buffer state initialization and interaction loops.

### ECWorkflow for Evolutionary Computation

`ECWorkflow` is tailored for evolutionary computation methods that maintain and evolve populations. It includes distributed population evaluation helpers and integrates with `EvoOptimizer` classes for ask/tell patterns.

## Required Methods for a Custom Workflow Implementation

Regardless of which base class you choose, you must implement these specific methods:

### Factory Methods

- **`build_from_config(cls, config, ...)`**: The public entry point used by training scripts. It deep-copies the configuration, optionally enables `pmap` or `jit`, and delegates to `_build_from_config`.

- **`_build_from_config(cls, config)`**: Constructs the actual workflow instance. Responsible for creating the environment, agent, optimizer, evaluator, and replay buffer using the configuration. Returns `cls(env, agent, ..., config)`.

### State Initialization

- **`setup(self, key)`**: Initializes the workflow state. Must initialize agent and optimizer states, reset the environment, and construct the initial `State` PyTree containing `key`, `metrics`, `agent_state`, `env_state`, and optional `replay_buffer_state`.

### Training Logic

- **`step(self, state)`**: The core training iteration. Responsible for sampling trajectories, computing losses, performing gradient updates or evolutionary ask/tell operations, updating metrics, and returning `(metrics, new_state)`.

### Optional Evaluation

- **`evaluate(self, state)`**: Runs the current policy in a separate evaluation environment to compute `EvaluateMetric` without affecting training state. Most RL workflows implement this for periodic evaluation.

## Step-by-Step: Implementing a Custom On-Policy Workflow

The following example demonstrates a minimal custom Workflow class extending `OnPolicyWorkflow`. It adds a custom logging metric while reusing standard PPO-style update logic.

Create the workflow file:

```python

# my_custom_workflow.py

from omegaconf import DictConfig
import chex
import jax

from evorl.workflows.rl_workflow import OnPolicyWorkflow
from evorl.metrics import WorkflowMetric
from evorl.types import State
from evorl.utils.rl_toolkits import flatten_rollout_trajectory, tree_stop_gradient


class MyCustomWorkflow(OnPolicyWorkflow):
    """A toy on-policy workflow that records the mean of the last-step reward."""

    @classmethod
    def name(cls):
        return "MyCustomOnPolicy"

    @classmethod
    def _build_from_config(cls, config: DictConfig):
        # Re-use helpers that build the environment, agent, optimizer and evaluator

        env = cls._make_env(config)
        agent = cls._make_agent(config, env)
        optimizer = cls._make_optimizer(config)
        evaluator = cls._make_evaluator(config, env)
        return cls(env, agent, optimizer, evaluator, config)

    def _setup_workflow_metrics(self) -> WorkflowMetric:
        # Extend base metrics with a new field

        base = super()._setup_workflow_metrics()
        return base.replace(mean_last_reward=0.0)

    def step(self, state: State):
        # 1️⃣ Roll out a single trajectory

        key, rollout_key = jax.random.split(state.key)
        trajectory, env_state = self.env.step(
            state.env_state,
            self.agent.compute_actions,
            rollout_key,
            self.config.rollout_length,
        )
        
        # 2️⃣ Compute loss and update policy (reuse parent logic)

        metrics, new_state = super().step(state.replace(key=key, env_state=env_state))
        
        # 3️⃣ Add custom metric (mean of last reward in trajectory)

        last_reward = trajectory.rewards[-1].mean()
        metrics = metrics.replace(mean_last_reward=last_reward)
        
        return metrics, new_state

```

**Explanation of key components:**

- **`name`**: Returns `"MyCustomOnPolicy"`, used by the CLI to locate this workflow.
- **`_build_from_config`**: Uses inherited helpers (`_make_env`, `_make_agent`, etc.) from `OnPolicyWorkflow` to construct components, ensuring consistency with the framework's configuration schema.
- **`_setup_workflow_metrics`**: Extends the generic `WorkflowMetric` with a custom field to track the mean of the last-step reward.
- **`step`**: Orchestrates environment interaction, delegates policy updates to the parent class, and injects custom metric computation.

## Registering and Running Your Custom Workflow

After implementing your custom Workflow class, register it so that the training script can instantiate it via Hydra configuration.

Import the class in the algorithms package initialization:

```python

# evorl/algorithms/__init__.py

from .my_custom_workflow import MyCustomWorkflow  # noqa: F401

```

This registration makes the workflow discoverable by the training script. You can now reference it in your configuration or command line:

```bash
python scripts/train.py \
  --config-name my_custom_config \
  workflow.name=MyCustomOnPolicy

```

The `workflow.name` parameter must match the string returned by your `name()` classmethod.

## Key Source Files Reference

| File | Role | Link |
|------|------|------|
| [`evorl/workflows/workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/workflow.py) | Abstract base and generic plumbing (`init`, `learn`, `close`). | [workflow.py](https://github.com/emi-group/evorl/blob/main/evorl/workflows/workflow.py) |
| [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) | Provides `RLWorkflow`, `OnPolicyWorkflow`, `OffPolicyWorkflow` – the main scaffolding for most RL algorithms. | [rl_workflow.py](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) |
| [`evorl/workflows/ec_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/ec_workflow.py) | Base for evolutionary‑computation workflows; useful if you need population handling. | [ec_workflow.py](https://github.com/emi-group/evorl/blob/main/evorl/workflows/ec_workflow.py) |
| [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) | Real‑world example of a concrete `OnPolicyWorkflow` (`PPOWorkflow`). | [ppo.py](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) |
| [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) | Example of an `OffPolicyWorkflow` (`TD3Workflow`). | [td3.py](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) |
| [`evorl/metrics.py`](https://github.com/emi-group/evorl/blob/main/evorl/metrics.py) | Metric containers (`WorkflowMetric`, `ECWorkflowMetric`, etc.) that are returned from `step`/`evaluate`. | [metrics.py](https://github.com/emi-group/evorl/blob/main/evorl/metrics.py) |
| [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py) | Definition of the `State` PyTree used throughout the workflow. | [types.py](https://github.com/emi-group/evorl/blob/main/evorl/types.py) |

## Summary

Implementing a custom Workflow class in EvoRL allows you to define specialized training loop logic while leveraging the framework's infrastructure for distributed training and checkpointing. The key steps include:

- **Select the appropriate base class**: Use `RLWorkflow` for gradient-based RL, `ECWorkflow` for evolutionary methods, or specialized templates like `OnPolicyWorkflow` and `OffPolicyWorkflow`.
- **Implement factory methods**: Provide `build_from_config` and `_build_from_config` to handle configuration-based instantiation.
- **Define state initialization**: Override `setup` to construct the initial `State` PyTree with all necessary components.
- **Implement the training step**: Override `step` to execute rollouts, compute updates, and return metrics.
- **Register the workflow**: Import the class in [`evorl/algorithms/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/__init__.py) to make it discoverable by the training script.

By adhering to this structure, your custom training loop automatically gains JIT compilation via JAX, multi-device support through `pmap`, and integrated checkpointing through the framework's `CheckpointManager`.

## Frequently Asked Questions

### What is the difference between RLWorkflow and ECWorkflow in EvoRL?

`RLWorkflow` is designed for model-free reinforcement learning algorithms that optimize policies via gradient descent, providing utilities for multi-device training and JIT compilation. `ECWorkflow`, located in [`evorl/workflows/ec_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/ec_workflow.py), is tailored for evolutionary computation methods that maintain populations, supplying distributed population evaluation helpers and integration with `EvoOptimizer` classes for ask/tell patterns.

### Do I need to implement the evaluate method for my custom Workflow?

Implementing `evaluate` is optional but recommended for reinforcement learning workflows. The method runs the current policy in a separate evaluation environment to compute `EvaluateMetric` without affecting training state. Most RL workflows in EvoRL implement this for periodic evaluation during training, but if your custom workflow does not require separate evaluation phases, you can omit this method.

### How does EvoRL handle multi-device training in custom Workflows?

Multi-device training is handled automatically by the base classes when you set `enable_multi_devices=True` in the `build_from_config` factory method. The `RLWorkflow` base class manages device replication via `pmap`, slices populations or batches across devices, and performs necessary `all_gather` or `psum` reductions. Your custom `step` method operates on a single device batch, with the base class handling parallelization.

### Where should I place my custom Workflow file in the EvoRL repository?

Place your custom Workflow class in the `evorl/algorithms/` directory alongside existing algorithm implementations like [`ppo.py`](https://github.com/emi-group/evorl/blob/main/ppo.py) and [`td3.py`](https://github.com/emi-group/evorl/blob/main/td3.py). Create a new file such as [`my_custom_workflow.py`](https://github.com/emi-group/evorl/blob/main/my_custom_workflow.py) in this directory. To make the workflow discoverable by the training script, import the class in [`evorl/algorithms/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/__init__.py) using `from .my_custom_workflow import MyCustomWorkflow`.