# How to Create Custom Neural Network Architectures for Policies in EvoRL

> Learn to build custom neural network architectures for policies in EvoRL. Use the flexible make_policy_network factory to easily assemble MLP, CNN, and normalization components for your AI agents.

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

---

**EvoRL provides a flexible `make_policy_network` factory in [`evorl/networks/linear.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/linear.py) that lets you assemble custom policy architectures from modular MLP, CNN, and normalization components.**

Creating custom neural network architectures for policies in EvoRL allows you to tailor agent behavior to specific observation spaces and action requirements. The repository provides reusable building blocks—including standard MLPs, spectral-normalized layers, CNN backbones, and vectorized ensemble networks—that can be composed through a unified factory interface. This guide walks you through the core components, the factory function signature, and practical implementations for common scenarios.

## Core Building Blocks for Custom Policy Networks in EvoRL

EvoRL organizes its neural network primitives into focused modules under `evorl/networks/`. Understanding these components is essential before assembling your custom architecture.

### MLP and Spectral-Normalized Layers

The standard multilayer perceptron is implemented in [`evorl/networks/linear.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/linear.py). It supports configurable hidden sizes, activation functions, bias terms, and optional normalization. For improved training stability, you can substitute the standard dense layers with spectral-normalized variants defined in [`evorl/networks/spectral_norm.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/spectral_norm.py). The `SNDense` layer wraps each linear transformation with spectral normalization, which constrains the Lipschitz constant of the network.

### Normalization Options: LayerNorm and StaticLayerNorm

Normalization layers reside in [`evorl/networks/layer_norm.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/layer_norm.py). EvoRL offers two variants:

- **LayerNorm**: Standard trainable scale and bias parameters.
- **StaticLayerNorm**: Fixed scale and bias, useful for deterministic policies where you want consistent normalization behavior during inference.

The `get_norm_layer` function selects the appropriate implementation based on the `norm_layer_type` string parameter passed to the factory.

### CNN Backbones for Image-Based Policies

For vision-based reinforcement learning, [`evorl/networks/cnn/atari_cnn.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/cnn/atari_cnn.py) provides the `CNN_AgentStem` class. This convolutional feature extractor processes image observations (e.g., stacked Atari frames) and outputs a flat feature vector suitable for the subsequent MLP policy head. The `make_cnn_agent` helper function facilitates creating a complete agent with both CNN stem and linear heads.

### V-Map MLP for Ensemble Policies

Ensemble methods require multiple independent network heads. EvoRL supports this through `make_vmap_mlp` in [`evorl/networks/linear.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/linear.py). This factory creates a vectorized batch of MLPs using JAX's `vmap` transformation, enabling parallel forward passes through multiple policy heads with shared parameter structure but independent weights.

## Using the make_policy_network Factory

The `make_policy_network` function in [`evorl/networks/linear.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/linear.py) serves as the primary entry point for creating custom policy architectures. It constructs a `PolicyModule` that orchestrates three operations:

1. **Observation Selection**: If `obs_key` is provided, extracts a specific sub-observation from dictionary-style observation spaces.
2. **Feature Extraction**: Instantiates an MLP (or spectral-normalized variant) via `make_mlp` with user-specified hidden sizes and normalization.
3. **Action Projection**: Applies a final dense layer projecting to `action_size` units, optionally followed by a custom final activation function.

```python
def make_policy_network(
    action_size: int,
    hidden_layer_sizes: Sequence[int] = (256, 256),
    use_bias: bool = True,
    activation: ActivationFn = nn.relu,
    activation_final: ActivationFn | None = None,
    norm_layer_type: str = "none",
    obs_key: str = "",
) -> nn.Module:
    ...

```

## Practical Examples: Creating Custom Policy Architectures

### Example: Deep MLP Policy with Layer Normalization

This example demonstrates creating a deeper architecture with layer normalization for improved training stability in continuous control tasks.

```python
import jax
import jax.numpy as jnp
from evorl.networks.linear import make_policy_network

def create_deep_norm_policy(action_dim: int, obs_key: str = ""):
    """Creates a 3-layer MLP policy with layer normalization."""
    return make_policy_network(
        action_size=action_dim,
        hidden_layer_sizes=(512, 256, 128),
        activation=jax.nn.relu,
        activation_final=jax.nn.tanh,
        norm_layer_type="layer_norm",
        obs_key=obs_key,
    )

# Initialize parameters

policy = create_deep_norm_policy(action_dim=4)
rng = jax.random.PRNGKey(0)
dummy_obs = jnp.zeros((1, 24))  # Replace with actual observation shape

params = policy.init(rng, dummy_obs)

```

The `norm_layer_type="layer_norm"` parameter triggers `get_norm_layer` in [`evorl/networks/layer_norm.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/layer_norm.py), inserting normalization after each hidden layer.

### Example: CNN-Backed Policy for Atari Environments

For image-based observations, combine the CNN stem from [`evorl/networks/cnn/atari_cnn.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/cnn/atari_cnn.py) with a standard MLP head.

```python
import jax
import jax.numpy as jnp
from evorl.networks.cnn.atari_cnn import make_cnn_agent

def create_atari_policy(action_dim: int, obs_shape: tuple = (84, 84, 4)):
    """Creates a CNN policy for Atari games."""
    (stem, actor_head, critic_head), init_fn = make_cnn_agent(
        obs_shape=obs_shape,
        action_size=action_dim,
        hidden_size=512,
    )
    
    # Initialize parameters

    rng = jax.random.PRNGKey(0)
    stem_params, actor_params, critic_params = init_fn(rng)
    
    def forward(rng, observation):
        # Extract features using CNN stem

        hidden, _ = stem.apply(stem_params, observation, mutable=["intermediates"])
        # Generate action logits

        action_logits = actor_head.apply(actor_params, hidden)
        return action_logits
    
    return forward, (stem_params, actor_params, critic_params)

# Usage

policy_fn, params = create_atari_policy(action_dim=6)
dummy_obs = jnp.zeros((1, 84, 84, 4))
logits = policy_fn(jax.random.PRNGKey(1), dummy_obs)

```

The `make_cnn_agent` function returns a tuple containing the `CNN_AgentStem` and linear heads, allowing flexible composition of convolutional and dense components.

### Example: Ensemble Policy Using V-Map

For ensemble methods or multi-head policies, use `make_vmap_mlp` to create vectorized independent networks.

```python
import jax
import jax.numpy as jnp
from evorl.networks.linear import make_vmap_mlp

def create_ensemble_policy(action_dim: int, n_ensembles: int = 4):
    """Creates an ensemble of independent policy networks."""
    # Create vectorized MLP with out_axes=0 for ensemble dimension

    mlp = make_vmap_mlp(
        layer_sizes=(256, 256, action_dim),
        norm_layer_type="none",
        out_axes=0,
    )
    
    def init(rng):
        dummy_obs = jnp.zeros((1, 24))  # Example observation shape

        return mlp.init(rng, dummy_obs)
    
    def apply(params, obs):
        # Broadcast observation to match ensemble dimension

        obs_batched = jnp.broadcast_to(obs, (n_ensembles,) + obs.shape)
        # Forward pass through all ensemble members

        return mlp.apply(params, obs_batched)  # Shape: (n_ensembles, batch, action_dim)

    
    return init, apply

# Initialize ensemble

init_fn, apply_fn = create_ensemble_policy(action_dim=4, n_ensembles=5)
rng = jax.random.PRNGKey(0)
params = init_fn(rng)

```

The `make_vmap_mlp` function leverages JAX's `vmap` transformation to create independent network instances that can be evaluated in parallel.

## Summary

- **EvoRL** provides modular network components in `evorl/networks/` that can be composed into custom policy architectures.
- The **`make_policy_network`** factory in [`evorl/networks/linear.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/linear.py) is the primary interface for creating MLP-based policies, supporting custom hidden sizes, activations, and normalization.
- **Layer normalization** (`layer_norm`) and **spectral normalization** (`spectral_norm`) can be inserted via string parameters to improve training stability.
- For **image-based observations**, combine `make_cnn_agent` from [`evorl/networks/cnn/atari_cnn.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/cnn/atari_cnn.py) with standard MLP heads.
- For **ensemble policies**, use `make_vmap_mlp` to create vectorized batches of independent networks with shared initialization logic.

## Frequently Asked Questions

### How do I add spectral normalization to my EvoRL policy network?

To enable spectral normalization, you need to modify the network factory to use the spectral-normalized MLP variant. While `make_policy_network` currently uses standard MLPs by default, you can import `make_mlp` with spectral normalization support from [`evorl/networks/spectral_norm.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/spectral_norm.py). The `SNDense` layer wraps standard linear transformations with spectral normalization, constraining the Lipschitz constant for more stable training dynamics.

### Can I use a custom CNN architecture different from the Atari CNN?

Yes, while EvoRL provides `CNN_AgentStem` in [`evorl/networks/cnn/atari_cnn.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/cnn/atari_cnn.py) for standard Atari preprocessing, you can define custom convolutional layers using standard Flax modules. Create your own stem module with `nn.Conv` layers, then pass the flattened output to `make_policy_network` or a custom linear head. The key is ensuring your CNN outputs a feature vector compatible with the subsequent MLP's input dimension.

### What is the difference between LayerNorm and StaticLayerNorm in EvoRL?

**LayerNorm** (accessible via `norm_layer_type="layer_norm"`) includes trainable scale and bias parameters that update during gradient descent, allowing the network to learn optimal normalization statistics. **StaticLayerNorm** (`norm_layer_type="static_layer_norm"`) uses fixed scale and bias values, making it ideal for deterministic policies or evolutionary strategies where you want consistent normalization behavior without additional trainable parameters. Both are implemented in [`evorl/networks/layer_norm.py`](https://github.com/emi-group/evorl/blob/main/evorl/networks/layer_norm.py) and selected via the `get_norm_layer` factory.

### How do I initialize parameters for a custom policy network in EvoRL?

All EvoRL network factories return standard Flax `nn.Module` instances. Initialize parameters by calling `module.init(rng, dummy_observation)` where `rng` is a JAX random key and `dummy_observation` has the correct shape matching your environment's observation space. For ensemble networks created with `make_vmap_mlp`, use the same initialization pattern—the `vmap` transformation handles batching across ensemble members automatically. The initialization function returns a parameter PyTree compatible with EvoRL's agent training loops.