Section 10.4: Vectorized environments; wrappers

"One environment teaches. Thirty-two environments running in lockstep teach thirty-two times faster, and every lesson is already batched."

A Throughput-Conscious Trainer

What This Section Builds

A single-arm policy that crawls toward convergence overnight can finish before lunch the moment you run 32 copies in lockstep and feed all their observations through one GPU forward pass; that 32x throughput, and the wrapper contract that makes it trustworthy, is what this section builds. Wrappers change what an environment exposes without rewriting the base simulator. Vector environments run several copies of a task behind one batched API. A single-environment CartPole training run typically needs around 500,000 steps to converge; a 32-environment vectorized run collects those same 500,000 steps in roughly 1/32 the wall-clock time, because every GPU forward pass consumes 32 transitions at once instead of one.

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.

In embodied AI, the wrapper contract is not just a software convenience. A real robot's sensors have fixed ranges, its actuators accept commands in specific units, and its safety system may cut power if a reward signal drives joint torques beyond hardware limits. A wrapper that rescales torque commands or clips force observations must match the physical actuator's operating envelope exactly. The wrong scaling can silently train the policy in a fantasy space, and the mismatch typically only surfaces when the policy runs on the physical arm and commands forces the joints cannot produce or cannot safely stop.

Figure 10.4B makes this two-part structure concrete. On the left, a wrapper stack sends actions down to the base simulator and passes observations back up through each transform. On the right, the same wrapped environment repeats N times behind a single batched vector API. The mechanism is direct. Every Gymnasium wrapper subclasses gym.Wrapper and overrides only the methods it transforms. When your code calls wrapped.step(action), Python calls the outermost wrapper's step. That wrapper applies its action transform, then delegates to the next wrapper's step, and so on until the innermost call reaches the base simulator. The returned observation travels back up through each wrapper's observation method and accumulates transformations in reverse. Gymnasium updates the spaces at construction time, so the policy always sees the final, post-wrapper space.

A wrapper that ships the wrong unit scaling to a physical joint does not produce a suboptimal policy; it produces a policy that was never trained for the hardware it will run on.

Worked Example

Code Fragment 10.4.1 applies a Gymnasium observation wrapper that 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 output shows the observation shape growing from four entries to five, and the second line confirms the wrapped environment still returns a legal step tuple. The policy now receives time information as part of its observation.

A single wrapper is easy to reason about; the harder skill this section promises is stacking several wrappers and knowing the resulting contract. Code Fragment 10.4.1b stacks a reward-clipping wrapper on top of the time-aware observation wrapper, in that order, and shows how to read off the combined effect.

# Stack two wrappers: first add time awareness, then clip rewards.
# Order matters: this clips the *base* reward, before any later
# reward-shaping wrapper would see it.
import gymnasium as gym
from gymnasium.wrappers import TimeAwareObservation, ClipReward

env = gym.make("CartPole-v1")
stacked = ClipReward(TimeAwareObservation(env), min_reward=0.0, max_reward=0.5)

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

print(stacked.observation_space.shape)
print(next_observation.shape, float(reward))
stacked.close()

The observation shape is still (5,) because TimeAwareObservation runs unchanged underneath; the reward is now clamped to 0.5 instead of the base environment's 1.0, because ClipReward is the outer layer. Reversing the stacking order would still clip to the same range here since CartPole's reward is a constant 1.0 per step, but with a reward that varies per step, clipping before versus after another reward transform can change the training signal, which is exactly why the wrapper order must be recorded alongside the result, not just the final wrapper class names.

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.

Why The Contract Must Be Logged

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.

Choosing the right wrapper and vectorization primitives is what makes that per-environment discipline practical, so the next table maps each transformation type to the role it plays and the cautions it carries.

One subtlety in vectorized rollouts is worth examining before touching code. When a sub-environment inside a SyncVectorEnv or AsyncVectorEnv reaches a terminal state, Gymnasium immediately resets it and places the first observation of the new episode into the returned batch. This behavior is called silent autoreset overwrite, and it is the most common source of corrupted value targets in vectorized training. Gymnasium moves the terminal observation into infos["final_observation"] for that sub-environment index. Consider a concrete case. Three CartPole environments run together, and environment index 1 falls after 12 steps. At step 12, envs.step(actions) returns the reset observation for index 1, not the observation where the pole fell. Code that reads observations[1] to compute a value-function bootstrap target (an estimate of future return computed from the state where an episode ended, used to fill in the reward-to-go beyond the collected rollout) then uses the wrong state. The correct value comes from infos["final_observation"][1]. In a 32-environment rollout of 1,000 steps, a CartPole episode ends roughly every 200 steps per sub-environment. The policy therefore receives around 160 corrupted bootstrap targets before any reward-curve blip grows large enough to notice, so the bug trains the value function on wrong states for hundreds of updates in complete silence.

Avoiding that autoreset trap is easier when the batched rollout is built on a foundation you have already verified one environment at a time.

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 a vectorized run fails, the wrapper stack is a reasonable first suspect, not automatically the policy. A concrete instance: an Isaac Lab ANYmal locomotion policy that trained cleanly in 4,096 parallel environments but collapsed on the physical ANYmal-C can fail because an action-rescaling wrapper mapped the policy output to a joint-torque range the real Series-Elastic Actuators, the spring-coupled motor units that let the robot's joints absorb impact instead of transmitting it directly, cannot deliver, rather than because Proximal Policy Optimization underfit. Isolate the cause before blaming the algorithm: pin one sub-environment, replay its logged action stream through the wrapper chain, and compare the post-wrapper torque command against the actuator's rated peak (ANYmal-C joints saturate near 80 Nm). If the wrapper emits commands above that ceiling, the bug is the contract, not the gradient. This single-environment replay turns a disappointing batched rollout into a reusable diagnostic asset.