# How to Optimize Training Performance Using `jax.jit` and `jax.vmap` in EvoRL

> Optimize EvoRL training with jax.jit and jax.vmap. Compile loops and vectorize operations for faster execution and better GPU/TPU use. Eliminate Python overhead.

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

---

**Use `jax.jit` to compile training loops into XLA-optimized kernels and `jax.vmap` to vectorize operations across environments or populations, eliminating Python overhead and maximizing GPU/TPU utilization in EvoRL.**

EvoRL (emi-group/evorl) is an open-source Evolutionary Reinforcement Learning library built entirely on JAX. To achieve high-throughput training, the codebase leverages JAX’s functional transformation primitives—specifically `jax.jit` for compilation and `jax.vmap` for automatic batching. This guide explains how these mechanisms are implemented in EvoRL’s workflow classes and how you can apply them to your own training pipelines.

## Understanding `jax.jit` in EvoRL

`jax.jit` traces a Python function once and compiles it to XLA-optimized machine code. Subsequent calls reuse the compiled version, eliminating Python dispatch overhead and enabling aggressive fusion of array operations.

### How JIT Compilation Works in Workflows

In EvoRL, the core workflow classes wrap their `step`, `evaluate`, and internal helper methods with `jax.jit`. For example, in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py), the `RLWorkflow` class applies JIT to the training step:

```python

# In evorl/workflows/rl_workflow.py

class RLWorkflow:
    @classmethod
    def step(cls, state, batch, key):
        # Training logic here

        ...
    

# JIT compilation with static_argnums=(0,) treats 'cls' as compile-time constant

jit_step = jax.jit(RLWorkflow.step, static_argnums=(0,))

```

The `static_argnums=(0,)` parameter tells JAX that the class instance (or class itself) does not change across calls, allowing it to treat object fields as compile-time constants. This pattern is also used in [`evorl/workflows/ec_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/ec_workflow.py) for evolutionary computation workflows.

### JIT-Friendly Utilities

EvoRL provides helper utilities in [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to ensure functions remain JIT-compatible:

- **`jit_method`**: A decorator that automatically applies `jax.jit` with appropriate `static_argnums`.
- **`rng_split`**: Pure functional PRNG splitting that works inside JIT blocks.
- **`tree_stop_gradient`**: Applies `jax.lax.stop_gradient` across pytrees within compiled code.

```python
from evorl.utils.jax_utils import jit_method

@jit_method(static_argnums=(0,))
def my_training_step(self, state, batch, key):
    key, subkey = rng_split(key)
    # Pure computation here

    return new_state, metrics, key

```

## Leveraging `jax.vmap` for Parallelization

While `jax.jit` removes Python overhead, `jax.vmap` vectorizes functions across a leading batch axis, turning explicit Python loops into single XLA operations. EvoRL uses this for parallel environment simulation and population-level evolutionary operators.

### Batched Environment Simulation

The [`evorl/envs/wrappers/training_wrapper.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/wrappers/training_wrapper.py) file contains `TrainingWrapper` and `MA_TrainingWrapper`, which expose a `vmap_step` flag. When enabled, the wrapper calls the environment’s `reset` and `step` methods with `jax.vmap`:

```python

# Conceptual usage from evorl/envs/wrappers/training_wrapper.py

env = TrainingWrapper(env_factory(), num_envs=128, vmap_step=True)

# Internally, this executes as a single XLA kernel:

# jax.vmap(env.step)(actions, states)

```

This pattern eliminates the Python loop over environments and allows the XLA compiler to fuse environment physics computations across all parallel instances.

### Population-Level Evolutionary Operators

In evolutionary computation workflows, `jax.vmap` processes entire populations simultaneously. For example, in [`evorl/ec/operators/mutation/mlp_mutation.py`](https://github.com/emi-group/evorl/blob/main/evorl/ec/operators/mutation/mlp_mutation.py), mutation operators are vectorized over the population:

```python
import jax
from evorl.ec.operators.mutation.mlp_mutation import MLPMutation

def mutate_population(pop_params, rng):
    mut = MLPMutation(...)
    # pop_params has leading axis = population size

    # jax.vmap applies mutate_fn to each individual

    return jax.vmap(mut.mutate_fn, in_axes=(0, 0))(
        pop_params, 
        jax.random.split(rng, pop_params.shape[0])
    )

```

The [`evorl/networks/linear.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/linear.py) file provides `make_vmap_mlp`, a utility for creating vectorized multi-layer perceptrons that work seamlessly with these population-level operations.

## Combining `jit` and `vmap` for Maximum Performance

The optimal pattern in EvoRL is combining both transformations: `jax.jit(jax.vmap(fn))`. This compiles the entire vectorized computation into a single XLA kernel.

In [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py), the rollout pipeline demonstrates this pattern:

```python
from evorl.rollout import rollout

# Vectorize rollout over multiple random keys (batch of episodes)

vmapped_rollout = jax.vmap(
    lambda key: rollout(env, policy, rng=key, max_steps=1000),
    in_axes=(0,)
)

# JIT compile the entire batch operation

fast_rollout = jax.jit(vmapped_rollout)

# Execute: single kernel launch, no Python overhead

metrics, trajectories = fast_rollout(jax.random.split(rng, 32))

```

This approach is used throughout the codebase, from [`evorl/workflows/ec_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/ec_workflow.py) for population evaluation to [`evorl/rollout_ma.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout_ma.py) for multi-agent scenarios.

## Practical Code Examples

### JIT-Compiled Training Step

When implementing custom workflows, use the `jit_method` decorator from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to ensure your training step compiles efficiently:

```python
from evorl.utils.jax_utils import jit_method, rng_split
import jax.numpy as jnp

class CustomWorkflow:
    @jit_method(static_argnums=(0,))
    def step(self, state, batch, key):
        # Split key for pure functional randomness

        key, subkey = rng_split(key)
        
        # Compute loss (example: MSE)

        predictions = state.policy.apply(state.params, batch.observations)
        loss = jnp.mean((predictions - batch.targets) ** 2)
        
        # Compute gradients and update

        grads = jax.grad(lambda p: jnp.mean(
            (state.policy.apply(p, batch.observations) - batch.targets) ** 2
        ))(state.params)
        
        new_params = jax.tree_util.tree_map(
            lambda p, g: p - 0.001 * g, state.params, grads
        )
        
        return state.replace(params=new_params), loss, key

```

### Vectorized Environment Rollout

For parallel environment execution, configure the `TrainingWrapper` with `vmap_step=True` and use `jax.vmap` for batched rollouts:

```python
import jax
from evorl.envs.wrappers.training_wrapper import TrainingWrapper
from evorl.rollout import rollout

def run_batch_evaluation(env_factory, policy, num_envs=128, rng=None):
    # Initialize wrapper with vmap_step enabled

    env = TrainingWrapper(
        env_factory(), 
        num_envs=num_envs, 
        vmap_step=True
    )
    
    # Split RNG for each environment

    rngs = jax.random.split(rng, num_envs)
    
    # Vectorized rollout across all environments

    vmapped_rollout = jax.vmap(
        lambda key: rollout(env, policy, rng=key, max_steps=1000)
    )
    
    # JIT compile for maximum performance

    compiled_rollout = jax.jit(vmapped_rollout)
    
    metrics, trajectories = compiled_rollout(rngs)
    return metrics, trajectories

```

### Population Mutation with vmap

For evolutionary computation workflows, vectorize mutation operators across the entire population:

```python
import jax
from evorl.ec.operators.mutation.mlp_mutation import MLPMutation

def apply_mutation(population_params, mutation_config, rng):
    """
    Vectorized mutation across population.
    
    Args:
        population_params: Pytree with leading axis [pop_size, ...]
        mutation_config: Configuration for MLPMutation
        rng: JAX random key
    """
    mutator = MLPMutation(mutation_config)
    
    # Split RNG for each individual

    pop_size = jax.tree_util.tree_leaves(population_params)[0].shape[0]
    rngs = jax.random.split(rng, pop_size)
    
    # vmap over population axis (axis 0 for both params and rngs)

    mutated_pop = jax.vmap(
        mutator.mutate_fn,
        in_axes=(0, 0)
    )(population_params, rngs)
    
    return mutated_pop

```

## Summary

- **`jax.jit`** compiles EvoRL workflow methods (such as `RLWorkflow.step` and `ECWorkflow` operations) into XLA kernels, eliminating Python overhead and fusing operations. Use `static_argnums=(0,)` to treat class instances as compile-time constants.

- **`jax.vmap`** vectorizes computations across batch dimensions, enabling parallel environment simulation (via `TrainingWrapper` with `vmap_step=True`) and population-level evolutionary operators (as seen in [`mlp_mutation.py`](https://github.com/emi-group/evorl/blob/main/mlp_mutation.py)).

- **Combine both** by wrapping `jax.vmap` with `jax.jit` (e.g., `jax.jit(jax.vmap(rollout))`) to compile entire batched operations into single kernel launches, which is the pattern used in [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py) for maximum throughput.

- **Maintain purity** inside JIT-compiled regions by using `jax.random` for randomness (via `rng_split` from [`jax_utils.py`](https://github.com/emi-group/evorl/blob/main/jax_utils.py)) and avoiding Python side-effects like list appends or print statements.

## Frequently Asked Questions

### When should I use `jax.jit` versus `jax.vmap` in EvoRL?

Use `jax.jit` when you need to eliminate Python overhead and compile a function into optimized machine code, particularly for the main training loop (`step` methods) and loss computations. Use `jax.vmap` when you need to apply the same function across multiple inputs—such as parallel environment steps (`vmap_step=True` in `TrainingWrapper`) or mutation operators across a population. In practice, EvoRL uses `jit` for the outer loop and `vmap` for inner batch operations, often combining them as `jax.jit(jax.vmap(fn))`.

### How do I handle random number generation inside JIT-compiled functions?

Always use JAX’s functional random number generation via `jax.random`. Inside EvoRL workflows, use the `rng_split` utility from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to split keys without side effects. For example: `key, subkey = rng_split(key)`. This ensures reproducibility and compatibility with `jax.jit`, which requires pure functions without mutable state. Never use Python’s `random` module or NumPy random inside JIT-compiled regions.

### Can I use Python control flow inside `jax.jit` in EvoRL?

Python control flow (like `if` statements or `for` loops) works inside `jax.jit` only if it depends on static arguments (marked via `static_argnums` or `static_argnames`). In EvoRL, workflow classes often pass `self` as a static argument (`static_argnums=(0,)`), allowing Python control flow based on configuration attributes. However, control flow based on dynamic array values must use JAX control flow operators like `jax.lax.cond` or `jax.lax.scan`. The `TrainingWrapper` in [`evorl/envs/wrappers/training_wrapper.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/wrappers/training_wrapper.py) handles this by selecting between `vmap` and `lax.map` based on static configuration flags.

### What is the performance impact of `static_argnums` in EvoRL workflows?

Using `static_argnums=(0,)` to mark the class instance (`self`) or class (`cls`) as static allows JAX to treat object attributes as compile-time constants. This reduces the compiled graph size and enables more aggressive XLA optimizations because JAX knows these values won’t change between calls. In EvoRL’s `RLWorkflow` and `ECWorkflow`, this pattern ensures that configuration parameters (like network architecture or hyperparameters) are baked into the compiled kernel, significantly reducing Python overhead during the actual training loop. However, changing any static argument requires recompilation, so these should only be used for true configuration constants, not for dynamic data like observations or parameters.