# EnvPool Parallel Environment Execution and Masking in Stable-WorldModel: A Technical Deep Dive

> Master EnvPool parallel environment execution and masking in Stable-WorldModel. Achieve high throughput for your world models by understanding selective masking and pre-allocated buffers.

- Repository: [GalilAI-group/stable-worldmodel](https://github.com/galilai-group/stable-worldmodel)
- Tags: deep-dive
- Published: 2026-05-30

---

**The `EnvPool` class in Stable-WorldModel enables high-throughput parallel environment execution through selective masking and pre-allocated info buffers, solving the challenge of heterogeneous episode lengths while maintaining strict `(batch, time, …)` tensor shapes required by world model networks.**

The `stable-worldmodel` repository by Galilai Group provides a specialized framework for training and evaluating world models that demands efficient handling of thousands of parallel environment interactions. Unlike standard vectorized environments, the custom **`EnvPool`** implementation adds critical features for world-model training: selective environment masking to handle episodes of varying lengths and pre-allocated stacked info dictionaries that preserve temporal batching conventions.

## How Selective Masking Handles Heterogeneous Episode Lengths

World model training frequently encounters scenarios where environments terminate at different timesteps. The **`mask`** argument in `EnvPool` solves this by allowing selective stepping and resetting of individual environments.

When you pass a boolean mask to `reset()` or `step()`, only environments where the mask value is **True** execute the operation. In [`stable_worldmodel/world/env_pool.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/env_pool.py), the `step` method implements this logic at lines 34-41:

```python
def step(self, actions: np.ndarray, mask: np.ndarray | None = None) -> ...:
    for i, env in enumerate(self.envs):
        if mask is not None and not mask[i]:
            continue
        _, rewards[i], terminateds[i], truncateds[i], info = env.step(actions[i])
        _write_env_info(self._stacked_infos, i, info)

```

For skipped environments, the method preserves the previous info state and forces reward and termination fields to zero/False. This mechanism is essential for evaluation modes where finished environments must remain frozen while others continue running.

The `reset` method follows identical masking semantics at lines 82-115:

```python
def reset(..., mask: np.ndarray | None = None) -> tuple[None, dict]:
    # Broadcasting and reset logic...

    if self._stacked_infos is None or mask is None:
        self._stacked_infos = _stack_fresh(per_env_infos)
    else:
        for i, info in enumerate(per_env_infos):
            if info is not None:
                _write_env_info(self._stacked_infos, i, info)

```

When a mask is provided, only selected environments receive fresh resets; others retain their existing state, enabling precise control over trajectory synchronization.

## Pre-Allocated Stacked Info Buffers

World models require auxiliary per-step data such as pixel observations, proprioception vectors, and goal images. Re-allocating these structures on every step creates overhead and breaks the `(batch, time, …)` tensor convention that downstream networks expect.

`EnvPool` solves this through **`_stacked_infos`**, a pre-allocated dictionary created during the first reset. As implemented in [`env_pool.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/env_pool.py) lines 109-115, subsequent resets write into existing buffers via `_write_env_info` rather than creating new dictionaries.

All array values in the stacked info obtain a leading time dimension of size 1, resulting in canonical shapes of `(num_envs, 1, …)`. This convention allows seamless concatenation along the time dimension for training batches without expensive memory reallocations.

## Reset Modes: Auto vs. Wait

The **`World`** class orchestrates `EnvPool` through two distinct reset modes controlled via `World.evaluate()` and the `reset_mode` parameter:

**`auto` mode** (default for episodic evaluation) immediately resets terminated environments with optional new seeds. This keeps the pool fully populated, maximizing throughput for standard reinforcement learning scenarios.

**`wait` mode** freezes environments that finish early, continuing the loop until all environments complete. This is critical for dataset-driven evaluation where each environment must follow pre-recorded trajectory lengths exactly, preserving the temporal structure of demonstration data.

In [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py), the `_run_iter` method implements this logic by constructing an `alive` mask after each step:

```python
mask = alive if not alive.all() else None
_, self.rewards, self.terminateds, self.truncateds, self.infos = self.envs.step(actions, mask=mask)

```

When any environment finishes in `wait` mode, the mask excludes it from future steps. In `auto` mode, the finished environments receive immediate resets via `self.envs.reset(..., mask=done)`.

## Data Flow and Implementation Details

`EnvPool` construction receives a list of factory functions (`env_fns`) that initialize individual environments, stored in `self.envs` (lines 34-35 of [`env_pool.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/env_pool.py)). During initialization, the wrapper batches action and observation spaces using Gymnasium's `batch_space`:

```python
self._action_space = batch_space(self._single_env.action_space, len(self.envs))
self._observation_space = batch_space(self._single_env.observation_space, len(self.envs))

```

This occurs once at construction (lines 40-45), avoiding the overhead of rebuilding space objects on every call. The environment then maintains consistent batched spaces regardless of internal masking state.

## Practical Code Examples

### Basic EnvPool Usage with Masking

```python
from stable_worldmodel.world.env_pool import EnvPool
import gymnasium as gym
import numpy as np

def make_cartpole():
    return gym.make("CartPole-v1")

# Create pool of 4 environments

pool = EnvPool([make_cartpole] * 4)

# Reset all environments

_, infos = pool.reset()

# Step with random actions

actions = np.random.randint(0, 2, size=(4,))
_, rewards, terminated, truncated, infos = pool.step(actions)

# Mask out first two environments (freeze them)

mask = np.array([False, False, True, True])
_, rewards, terminated, truncated, infos = pool.step(actions, mask=mask)

# Environments 0-1 maintain previous state; 2-3 step forward

```

### Dataset-Driven Evaluation with Wait Mode

```python
import stable_worldmodel as swm

world = swm.World(
    env_name="swm/PushT-v1",
    num_envs=3,
    image_shape=(64, 64),
)

world.set_policy(swm.RandomPolicy(world.envs.action_space))

# Use 'wait' mode to preserve exact trajectory lengths

results = world.evaluate(
    dataset=my_dataset,
    episodes_idx=[0, 1, 2],
    start_steps=[0, 10, 20],
    goal_offset=30,
    eval_budget=100,
    reset_mode="wait",
)

```

### Custom Training Loop with Manual Masking

```python
world.reset(seed=42)
while True:
    actions = world.policy.get_action(world.infos)
    # Only step environments still alive

    mask = ~world.terminateds
    _, _, terminated, _, _ = world.envs.step(actions, mask=mask)
    world.terminateds |= terminated
    if world.terminateds.all():
        break

```

## Summary

- **`EnvPool`** provides a lightweight alternative to Gymnasium's `SyncVectorEnv` with specialized masking capabilities for world model training.
- The **`mask`** argument in `reset()` and `step()` allows selective environment execution, freezing finished episodes while others continue in parallel.
- **Pre-allocated info buffers** maintain `(num_envs, 1, …)` tensor shapes, eliminating per-step memory allocation and preserving batch conventions.
- **`auto` reset mode** maximizes throughput by immediately resetting finished environments, while **`wait` mode** preserves exact trajectory lengths for dataset evaluation.
- Implementation resides primarily in [`stable_worldmodel/world/env_pool.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/env_pool.py) and [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py), with the `World` class handling high-level mask orchestration.

## Frequently Asked Questions

### How does EnvPool handle environments that terminate at different times?

When environments terminate asynchronously, `EnvPool` uses the **`mask`** argument to exclude finished environments from subsequent `step()` calls. In `wait` mode, the `World` class constructs an `alive` mask that becomes `False` for terminated environments, causing `EnvPool` to skip them while preserving their final state. This allows heterogeneous episode lengths without disrupting the batch structure of active environments.

### What is the difference between _stack_fresh and _write_env_info in the EnvPool implementation?

**`_stack_fresh`** creates new pre-allocated info dictionaries with the proper `(num_envs, 1, …)` shape during the initial reset or full resets without masking. **`_write_env_info`** performs in-place updates to existing buffers during masked resets, writing new info data only for environments where the mask is True. This distinction eliminates memory reallocation during training loops while maintaining consistent tensor shapes.

### Why does EnvPool add a time dimension of size 1 to info tensors?

World model networks typically expect input tensors with a time dimension for temporal modeling (e.g., transformers or RNNs). By prepending a dimension of size 1 to create shapes like `(num_envs, 1, height, width)` for pixels, `EnvPool` ensures that downstream components can concatenate observations along the time axis using standard array operations without costly reshaping or data copying.

### When should I use wait mode versus auto mode for environment resetting?

Use **wait mode** when evaluating on fixed datasets where each environment must complete exactly the number of steps specified in the dataset, regardless of early termination signals. This preserves the temporal alignment between demonstration trajectories. Use **auto mode** for standard RL training or episodic evaluation where maximizing throughput is priority, as it immediately resets finished environments to keep the parallel batch fully utilized.