# How to Use PyTreeDict, PyTreeData, and PyTreeNode for JAX-Compatible Data Structures in EvoRL

> Leverage PyTreeDict, PyTreeData, and PyTreeNode in EvoRL for JAX-compatible data structures. Effectively manage configurations, immutable data, and mutable state with JAX's JIT compilation and vectorization.

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

---

**PyTreeDict, PyTreeData, and PyTreeNode are lightweight JAX pytree helpers in EvoRL that enable JIT-compiled, vectorized transformations of configuration objects, immutable data containers, and mutable state nodes by automatically registering with `jax.tree_util.register_pytree_node_class`.**

EvoRL is an open-source evolutionary reinforcement learning framework built on JAX. To ensure that configuration objects, trajectory data, and algorithmic states can seamlessly pass through JIT compilation and vectorization, the library provides three specialized pytree utilities in [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py). These helpers eliminate boilerplate code while maintaining the functional programming semantics required for high-performance GPU/TPU execution.

## Understanding the Three JAX-Compatible Helpers

EvoRL provides distinct abstractions for different mutability requirements, all registered with JAX through `jax.tree_util.register_pytree_node_class` at class definition time.

### PyTreeDict for Mutable Configuration

**`PyTreeDict`** is a dictionary subclass that registers itself as a pytree node and automatically converts nested `dict`, `list`, and `tuple` values into the same type. It implements `tree_flatten` and `tree_unflatten` to expose keys as auxiliary data and values as children.

Use `PyTreeDict` for configuration objects, environment-extra fields, and metrics containers—any mutable mapping that must travel through JAX transforms like `jax.vmap` or `jax.lax.scan`.

### PyTreeData for Immutable Containers

**`PyTreeData`** is a dataclass-style base class that registers its fields as a pytree via `register_pytree_with_keys`. All fields are required at construction and cannot be mutated afterwards—there is no `set_frozen_attr` method.

Use this for fixed data containers such as trajectories, replay-buffer entries, and algorithmic states that must remain immutable for JIT safety. The class provides a `replace` method for functional updates.

### PyTreeNode for Mutable State

**`PyTreeNode`** is similar to `PyTreeData` but allows mutable attributes after construction via `set_frozen_attr` and supports a mutable `__post_init__`. 

Use this for nodes that need post-initialization mutation, such as building auxiliary caches or counters during training loops.

## Core Implementation Details

The registration mechanism differs slightly between the dictionary and dataclass variants. In [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py), `PyTreeDict` implements `tree_flatten` to expose its internal dictionary structure to JAX, while `PyTreeData` and `PyTreeNode` are built on a custom `dataclass` wrapper that separates static (metadata) fields from data fields.

Both dataclass variants use `jax.tree_util.register_pytree_with_keys` to ensure that field names are preserved during flattening operations. The `replace` method—inherited by both `PyTreeData` and `PyTreeNode`—wraps `dataclasses.replace` to return new instances without mutating the original, preserving the functional semantics required for JIT compilation.

## Practical Code Examples

### Creating Configuration Objects with PyTreeDict

Instantiate a `PyTreeDict` to create JAX-compatible configuration objects that automatically convert nested dictionaries:

```python
from evorl.types import PyTreeDict

# A configuration that can be passed through JAX transforms

env_cfg = PyTreeDict(
    env_name="ant",
    env_type="brax",
    max_steps=1000,
)

# Nested structures are auto-converted

env_cfg = PyTreeDict(
    env=PyTreeDict(name="ant", type="brax"),
    training=PyTreeDict(seed=42, parallel=8),
)

```

All nested dicts become `PyTreeDict` automatically, enabling attribute-style access such as `env_cfg.env.name`.

### Defining Immutable Data with PyTreeData

Define trajectory containers and state objects using `PyTreeData` and the `pytree_field()` helper:

```python
from evorl.types import PyTreeData, pytree_field
import chex
import jax.numpy as jnp

class EpisodeStats(PyTreeData):
    """Aggregated statistics for a batch of episodes."""
    episode_returns: chex.Array = pytree_field()
    episode_lengths: chex.Array = pytree_field()

# Initialise (arrays must be JAX-compatible)

stats = EpisodeStats(
    episode_returns=jnp.zeros((8,)),   # 8 parallel envs

    episode_lengths=jnp.zeros((8,), dtype=jnp.int32),
)

# Immutable update – returns a new instance

new_stats = stats.replace(episode_returns=stats.episode_returns + 1.0)

```

The `SampleBatch` class in [`evorl/sample_batch.py`](https://github.com/emi-group/evorl/blob/main/evorl/sample_batch.py) inherits from `PyTreeData` and mixes in `PyTreeArrayMixin`, providing convenient arithmetic operators on whole batches while maintaining JIT compatibility.

### Managing Mutable State with PyTreeNode

For states requiring post-initialization mutation, use `PyTreeNode` with `set_frozen_attr`:

```python
from evorl.types import PyTreeNode, pytree_field

class CounterNode(PyTreeNode):
    """A simple mutable counter used during training."""
    count: int = pytree_field(static=True)  # static → does not affect JIT recompilation

    def incr(self):
        # Mutate safely after construction

        self.set_frozen_attr("count", self.count + 1)

counter = CounterNode(count=0)
counter.incr()
print(counter.count)  # 1

```

Marking a field as `static=True` via `pytree_field` tells JAX that changes to this field do not require recompilation, enabling efficient updates to metadata during training loops.

## Integration in EvoRL Workflows

EvoRL's core loops pass whole state objects through JAX transformations. In [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py), `PyTreeDict` collects extra information from environments, while `SampleBatch` (a `PyTreeData` subclass) stores trajectories:

```python
from evorl.rollout import rollout
from evorl.envs import create_env, AutoresetMode
from evorl.sample_batch import SampleBatch
from evorl.types import PyTreeDict

# Build env with PyTreeDict configuration

cfg = PyTreeDict(env_name="ant", env_type="brax")
env = create_env(cfg, parallel=4, autoreset_mode=AutoresetMode.NORMAL)

# Collect trajectory – returns SampleBatch (inherits from PyTreeData)

trajectory, final_state = rollout(
    env.step,
    agent.compute_actions,
    env_state,
    agent_state,
    key,
    rollout_length=128,
)

```

Because `SampleBatch` is registered as a pytree, `jax.vmap` automatically vectorizes over batched environments, with each field acting as a leaf array. The trajectory structure maintains type safety while supporting operations across the entire batch.

## Summary

- **PyTreeDict** provides mutable, dictionary-like configuration objects that auto-convert nested structures and pass through JIT compilation.
- **PyTreeData** offers immutable dataclass semantics with functional `replace()` updates, ideal for trajectories and algorithmic states in [`evorl/sample_batch.py`](https://github.com/emi-group/evorl/blob/main/evorl/sample_batch.py).
- **PyTreeNode** enables controlled mutation after construction via `set_frozen_attr`, supporting auxiliary caches and counters.
- All three classes are registered with `jax.tree_util.register_pytree_node_class` or `register_pytree_with_keys` in [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py), enabling seamless integration with `jax.vmap`, `jax.jit`, and `jax.lax.scan`.

## Frequently Asked Questions

### What is the difference between PyTreeData and PyTreeNode?

**`PyTreeData`** is strictly immutable after construction and uses the `replace()` method for updates, making it safe for JIT-compiled functional transformations. **`PyTreeNode`** allows mutation via `set_frozen_attr()` after initialization, supporting use cases that require post-construction modification such as building internal caches or updating counters.

### When should I use static=True in pytree_field?

Use **`static=True`** for fields that do not change during JIT-compiled execution or that should not trigger recompilation when modified. In [`evorl/ec/optimizers/vanilla_ga.py`](https://github.com/emi-group/evorl/blob/main/evorl/ec/optimizers/vanilla_ga.py), static fields often include metadata like population size or configuration flags, while array data remains dynamic.

### Can these data structures be nested arbitrarily?

Yes. **`PyTreeDict`** automatically converts nested dictionaries, and both **`PyTreeData`** and **`PyTreeNode`** can contain other pytree-registered objects as fields. The `tree_flatten` and `tree_unflatten` implementations in [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py) handle arbitrary nesting depth correctly.

### How does the replace() method work with JAX transformations?

The `replace()` method calls `dataclasses.replace()` under the hood to create a new instance with updated fields, leaving the original unchanged. This preserves the functional programming model required by JAX's JIT compiler, allowing values to flow through `jax.lax.scan` and gradient computations without side effects.