Section 10.4: Vectorized environments; wrappers

A Careful Control Loop

What This Section Builds

Wrappers and vectorized environments change the contract around a base task. Wrappers change what an environment exposes without rewriting the base simulator. Vector environments run several copies of a task behind one batched API.

The goal is to keep throughput improvements honest. A wrapper stack should be declared, ordered, and logged, and a vectorized rollout should preserve per-environment termination, truncation, reward, and info fields.

Theory

A wrapper is an environment transformation. It can alter observations, actions, rewards, metadata, or the behavior of reset and step. The base task remains underneath, but the policy sees the wrapped contract.

A vector environment batches several environment copies so the learner collects experience faster. Gymnasium vector envs return arrays for rewards, terminations, and truncations with one element per sub-environment. Observations are batched according to the observation space, which is why space design from Section 10.2 matters before vectorization.

Worked Example

Code Fragment 10.4.1 applies a simple Gymnasium observation wrapper. The wrapper adds time awareness to the observation, so the policy sees a five-value observation instead of the base four-value CartPole state.

# Wrap an environment so the observation includes elapsed time.
# The policy sees the wrapped observation space, not the base one.
import gymnasium as gym
from gymnasium.wrappers import TimeAwareObservation

env = gym.make("CartPole-v1")
wrapped = TimeAwareObservation(env)

observation, info = wrapped.reset(seed=13)
wrapped.action_space.seed(13)
action = wrapped.action_space.sample()
next_observation, reward, terminated, truncated, info = wrapped.step(action)

print(wrapped.observation_space.shape)
print(next_observation.shape, float(reward), terminated, truncated)
wrapped.close()

The expected output shows the observation shape grow from the usual CartPole size to five entries after the wrapper is applied. The second line confirms that the wrapped environment still returns a legal step tuple, but the policy is now receiving time information as part of its observation.

Practical Recipe

1. Write the base environment contract before adding wrappers.
2. Add one wrapper at a time and record how it changes spaces, rewards, or info.
3. Use vector environments when rollout throughput, not environment semantics, is the bottleneck.
4. Interpret vector outputs per sub-environment, not as one scalar episode.
5. Log autoreset mode and final observations when using vector rollouts that reset sub-environments automatically.

Wrappers are powerful because they separate task dynamics from interface transformations. That separation also creates a risk: the experiment may claim to evaluate a base environment while the policy actually saw normalized observations, clipped rewards, time features, action rescaling, and a time-limit wrapper.

Vector environments add another layer. They make the rollout batch look like one object, but each sub-environment still has its own episode boundary. Evaluation code should preserve that identity so one unstable instance does not disappear inside an average.

A robust vector implementation proves the single environment first, then creates a batched version with the same wrappers. The first artifact explains semantics; the second artifact explains throughput.

1. Run one unwrapped environment step and save the return shapes.
2. Add wrappers and rerun the same seed to show how the contract changed.
3. Create the vector environment only after the wrapper stack is fixed.
4. Store per-sub-environment rewards, termination flags, truncation flags, and final info.
5. Use aggregate plots only after preserving the per-environment traces.

# Step three CartPole environments through one vectorized API call.
# Rewards and ending flags retain one element per sub-environment.
import gymnasium as gym

envs = gym.make_vec("CartPole-v1", num_envs=3, vectorization_mode="sync")
observations, infos = envs.reset(seed=11)
envs.action_space.seed(11)
actions = envs.action_space.sample()
observations, rewards, terminations, truncations, infos = envs.step(actions)

print(observations.shape)
print(actions.tolist())
print(rewards.tolist())
print(terminations.tolist(), truncations.tolist())
envs.close()

The expected output should be read row-wise across three sub-environments: a batched observation tensor of shape (3, 4), three sampled actions, three rewards, and one termination and truncation flag per environment. Nothing is aggregated yet, which is exactly the right interpretation for vectorized evidence.

When an experiment about vectorized environments; wrappers fails, avoid labeling the whole method as weak. First assign the failure to perception, state estimation, planning, control, timing, data coverage, or evaluation. Then rerun one controlled perturbation that isolates the suspected cause. This pattern turns a disappointing rollout into a reusable diagnostic asset.