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

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. 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 and 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:

# 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:

# 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:

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	Abstract base and generic plumbing (init, learn, close).	workflow.py
evorl/workflows/rl_workflow.py	Provides RLWorkflow, OnPolicyWorkflow, OffPolicyWorkflow – the main scaffolding for most RL algorithms.	rl_workflow.py
evorl/workflows/ec_workflow.py	Base for evolutionary‑computation workflows; useful if you need population handling.	ec_workflow.py
evorl/algorithms/ppo.py	Real‑world example of a concrete OnPolicyWorkflow (PPOWorkflow).	ppo.py
evorl/algorithms/td3.py	Example of an OffPolicyWorkflow (TD3Workflow).	td3.py
evorl/metrics.py	Metric containers (WorkflowMetric, ECWorkflowMetric, etc.) that are returned from step/evaluate.	metrics.py
evorl/types.py	Definition of the State PyTree used throughout the workflow.	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 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, 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 and td3.py. Create a new file such as my_custom_workflow.py in this directory. To make the workflow discoverable by the training script, import the class in evorl/algorithms/__init__.py using from .my_custom_workflow import MyCustomWorkflow.

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