# How to Leverage EvoRL's Object-Oriented Functional Programming Model with JAX jit

> Learn how EvoRL uses JAX jit with its object-oriented functional programming model. Discover how class instances act as static containers for pure functional kernels and immutable pytrees.

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

---

**EvoRL combines object-oriented design with functional JAX primitives by treating class instances as static containers while JIT-compiling pure functional kernels that operate on immutable pytrees.**

EvoRL is an open-source evolutionary reinforcement learning framework that bridges Python's object-oriented patterns with JAX's functional programming requirements. By structuring **agents** and **workflows** as stateful objects while keeping computational kernels pure, EvoRL enables seamless JIT compilation without sacrificing code organization. This hybrid **object-oriented functional programming model with JAX jit** allows researchers to extend the framework using familiar class-based inheritance while automatically gaining XLA-accelerated performance.

## Core Architecture of EvoRL's Hybrid Model

EvoRL's architecture cleanly separates stateful object management from pure functional computation. This design appears throughout the codebase, from the base `Agent` class to the `RLWorkflow` orchestration layer.

### Stateful Objects as Static Containers

In EvoRL, high-level components like **agents** and **workflows** are implemented as Python classes that inherit from `Agent` and `RLWorkflow`. These classes reside in [`evorl/agent.py`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py) and [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) respectively. Rather than mutating internal state, these objects act as containers for immutable pytrees—holding **parameters**, **RNG keys**, and configuration while exposing pure functional methods.

The `Agent` base class defines a strict interface where methods like `compute_actions`, `evaluate_actions`, and `loss` receive explicit inputs (state, observations, keys) and return outputs without side effects. This functional purity ensures that instances can be treated as **static arguments** during JIT compilation.

### Pure Functional Kernels

The computational heavy lifting occurs in standalone functions found in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) and [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py). These kernels operate on immutable data structures like `SampleBatch`, `AgentState`, and `Params`. Because they avoid global state and mutate no inputs, these functions are trivially eligible for `jax.jit` compilation.

For example, utilities like `compute_gae` and `average_episode_discount_return` accept pytrees and return transformed pytrees, making them ideal candidates for XLA optimization.

### Automatic JIT and Pmap Wrappers

The `RLWorkflow` class provides class methods `enable_jit()` and `enable_pmap()` that automatically wrap the `step` and `evaluate` methods. As implemented in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py), `enable_jit` applies `jax.jit` with `static_argnums=(0,)` to treat `self` as a compile-time constant:

```python
@classmethod
def enable_jit(cls):
    cls.evaluate = jax.jit(cls.evaluate, static_argnums=(0,))
    cls.step = jax.jit(cls.step, static_argnums=(0,))

```

For custom methods, EvoRL provides the `jit_method` and `pmap_method` factories in [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py). These decorators return partial functions pre-configured with `static_argnums=0`, allowing seamless JIT compilation of instance methods. The implementation uses `functools.partial`:

```python
def jit_method(*, static_argnums=None, **kwargs):
    return partial(jax.jit,
                   static_argnums=static_argnums,
                   **kwargs)

```

## Implementing JIT Compilation in EvoRL

Depending on your use case, EvoRL offers multiple patterns for leveraging JIT compilation—from built-in workflow configuration to custom method decoration.

### Standard Algorithm Configuration

For built-in algorithms like PPO or SAC, enable JIT compilation at construction time by passing `enable_jit=True` to the workflow builder. This triggers `RLWorkflow.enable_jit()` which compiles the core training loop:

```python
from evorl.workflows import PPOWorkflow
from omegaconf import OmegaConf

cfg = OmegaConf.load("configs/ppo.yaml")
workflow = PPOWorkflow.build_from_config(cfg, enable_jit=True)
final_state = workflow.learn(workflow.setup(jax.random.PRNGKey(0)))

```

This single flag JIT-compiles both `workflow.step` and `workflow.evaluate`, eliminating Python overhead during training.

### Custom Workflow Methods

When extending `RLWorkflow` with custom training logic, use the `@jit_method` decorator to compile your methods. Import the decorator from `evorl/utils/jax_utils` and specify `static_argnums=0` to treat the instance as static:

```python
from evorl.workflows import RLWorkflow
from evorl.utils.jax_utils import jit_method
import jax.numpy as jnp

class MyWorkflow(RLWorkflow):
    @jit_method(static_argnums=0)
    def compute_advantages(self, rewards, values, dones):
        """Pure JAX implementation of generalized advantage estimation."""
        deltas = rewards + 0.99 * (1.0 - dones) * values[1:] - values[:-1]
        
        def scan_fun(acc, delta):
            return acc * 0.95 + delta, acc * 0.95 + delta
        
        _, advantages = jax.lax.scan(scan_fun, jnp.zeros_like(deltas[0]), deltas, reverse=True)
        return advantages

```

The decorator ensures `compute_advantages` runs as compiled XLA code while maintaining the object-oriented structure of your workflow.

### JIT-Compiling Agent Methods

Custom agents can also leverage JIT compilation for expensive operations like forward passes. Subclass `Agent` and apply `@jit_method` to methods that operate on immutable `AgentState`:

```python
from evorl.agent import Agent, AgentState
from evorl.utils.jax_utils import jit_method
import jax.numpy as jnp

class LinearAgent(Agent):
    @jit_method(static_argnums=0)
    def compute_actions(self, state: AgentState, batch, key):
        """JIT-compiled linear policy forward pass."""
        w = state.params["w"]  # (obs_dim, act_dim)

        obs = batch.obs        # (B, obs_dim)

        logits = jnp.dot(obs, w)
        actions = jax.random.categorical(key, logits)
        return actions, {}

```

Because `state` is an immutable pytree passed as an argument, JAX can cache the compiled kernel across training steps.

### Multi-Device Training with Pmap

For distributed training across multiple GPUs or TPUs, use `enable_multi_devices=True` when building the workflow. This invokes `RLWorkflow.enable_pmap()`, which wraps methods with `jax.pmap` instead of `jax.jit`:

```python
workflow = PPOWorkflow.build_from_config(
    cfg, enable_multi_devices=True, enable_jit=False
)

```

To write device-agnostic custom methods that work with both JIT and pmap, use the `pmap_method` decorator from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) with the same `static_argnums=0` pattern.

## Technical Foundations: Why This Design Works

EvoRL's hybrid model succeeds because it respects JAX's fundamental requirements while preserving Python's object-oriented ergonomics.

**Functional purity** ensures that all JIT-eligible code paths receive explicit inputs and return outputs without mutating global state. By passing **immutable pytrees** (structures registered with JAX's tree_util) rather than Python objects, the framework allows XLA to optimize memory layout and execution.

The critical technique is treating `self` as a **static argument** via `static_argnums=0`. This tells JAX to treat the Python instance as compile-time metadata rather than a traced value, allowing the instance to hold configuration and Python state while the compiled kernel operates on the dynamic pytree arguments.

## Summary

EvoRL's **object-oriented functional programming model** provides a template for high-performance ML research:

- **Stateful containers**: `Agent` and `RLWorkflow` classes in [`evorl/agent.py`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py) and [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) hold configuration and immutable state.
- **Pure kernels**: Computational logic resides in functional utilities like those in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py), operating only on pytrees.
- **Automatic compilation**: `enable_jit()` and `enable_pmap()` class methods compile the training loop without manual intervention.
- **Custom decorators**: `jit_method` and `pmap_method` in [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) extend JIT capabilities to user-defined methods.
- **Static self**: Using `static_argnums=0` allows instance methods to be JIT-compiled while preserving object-oriented design patterns.

## Frequently Asked Questions

### What is the object-oriented functional programming model in EvoRL?

The **object-oriented functional programming model** in EvoRL refers to the framework's hybrid architecture where Python classes (like `Agent` and `RLWorkflow`) provide structure and hold immutable state, while all computational operations are implemented as pure functions compatible with JAX's functional paradigm. This allows the framework to use object inheritance and encapsulation for code organization while leveraging `jax.jit` for performance.

### How does EvoRL handle JIT compilation of class methods?

EvoRL handles JIT compilation through the `static_argnums=0` pattern. The `enable_jit()` method in `RLWorkflow` wraps `step` and `evaluate` with `jax.jit`, specifying that the first argument (`self`) should be treated as a static compile-time constant rather than a traced value. For custom methods, the `jit_method` decorator in [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) automates this configuration.

### Can I use JIT with custom agents in EvoRL?

Yes. When implementing a custom agent by subclassing `Agent` from [`evorl/agent.py`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py), you can apply the `@jit_method(static_argnums=0)` decorator to any pure method that accepts immutable pytrees (like `AgentState`). This compiles the method to XLA while maintaining the object-oriented API, as demonstrated in the `LinearAgent` example above.

### What is the difference between enable_jit and enable_pmap in EvoRL?

`enable_jit()` compiles methods for single-device execution using `jax.jit`, while `enable_pmap()` distributes computation across multiple devices using `jax.pmap`. In [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py), `enable_jit` is used for standard single-GPU/CPU training, whereas `enable_pmap` replicates the workflow across devices for data-parallel training. You typically enable one or the other via `build_from_config()` parameters, not both simultaneously.