"""Vectorized environments.
`VecEnv` runs $B$ copies of a gymnasium environment in parallel and talks
**torch tensors**: `reset()` returns a batch of observations
(`[B, obs_dim]`), and `step(actions)` returns an `EnvStep` — one entry per
environment — where `obs` is the observation to act upon at the next step
and `next_obs` is the true successor state $s_{t+1}$ (they differ at
episode boundaries, see below).
**Structured observations.** If the environment's observation space is a
gymnasium `Dict` (e.g. named sub-observations, as in SuperTuxKart), the
observations are returned as a `TensorStruct` with one tensor field per key
(recursively for nested `Dict`s). Because every `TensorStruct` operation
recurses into nested structs, such observations flow through
`Transitions` / `Episode` / `Rollout`, the `ReplayBuffer` and the collectors
unchanged — you only need to write a policy/critic that *encodes* the
structured observation (that part is environment-specific). Ragged fields
(a variable number of entities per step) must be padded to a fixed size
first, since tensors have fixed shapes.
Environments **auto-reset** when an episode ends (`terminated` or
`truncated` is `True`), in one of two modes:
- *next-step reset* (default): the step that ends the episode returns its
final observation; the **following** call to `step` ignores the action
and returns the first observation of a fresh episode;
- *same-step reset* (`same_step_reset=True`): the step that ends the
episode directly returns the first observation of a fresh episode in
`obs`, while `next_obs` holds the final observation of the ended
episode. Every step is then a valid transition (used by
`RolloutCollector`).
The collectors (`rl_mind.collectors`) handle this bookkeeping for you.
The distinction between the two ending flags matters for learning:
- `terminated`: the episode reached a **terminal state** (e.g. the pole
fell). The value of the future is 0 — do *not* bootstrap.
- `truncated`: the episode was **cut short** (e.g. time limit). The agent
could have continued — *do* bootstrap with the value of the next state.
"""
from dataclasses import dataclass, make_dataclass
from typing import Callable, Sequence, Union
import gymnasium as gym
import numpy as np
import torch
from torch import Tensor
from rl_mind.core import TensorStruct
__all__ = ["EnvStep", "VecEnv"]
#: An observation is either a plain tensor (flat `Box` space) or a
#: `TensorStruct` of named tensors (`Dict` space)
Observation = Union[Tensor, TensorStruct]
def _make_obs_converter(space: gym.Space) -> Callable[[object], Observation]:
"""Build a function converting a gymnasium observation to a tensor or,
for a `Dict` space, to a `TensorStruct` of named tensors (recursively).
The same converter works for a batched observation (`[B, ...]`) and for a
single one (`[...]`): only the leading dimension differs.
A `Discrete` space (tabular environments) gives **integer** observations:
they are state *indices*, not features, so they keep an integer dtype and
can be used directly to index a Q-table.
"""
if isinstance(space, gym.spaces.Discrete):
return lambda obs: torch.as_tensor(obs, dtype=torch.long)
if isinstance(space, gym.spaces.Dict):
subs = {key: _make_obs_converter(sub) for key, sub in space.spaces.items()}
# A frozen TensorStruct dataclass with one field per key
cls = make_dataclass(
"Observation",
[(key, Tensor) for key in subs],
bases=(TensorStruct,),
frozen=True,
)
def convert(obs, subs=subs, cls=cls):
return cls(**{key: sub(obs[key]) for key, sub in subs.items()})
return convert
return lambda obs: torch.as_tensor(obs, dtype=torch.float32)
def _clone_obs(obs: Observation) -> Observation:
if isinstance(obs, TensorStruct):
return obs._map(lambda tensor: tensor.clone())
return obs.clone()
def _set_obs_at(obs: Observation, index, value: Observation) -> None:
if isinstance(obs, TensorStruct):
obs.set_(index, value)
else:
obs[index] = value
[docs]
@dataclass(frozen=True)
class EnvStep:
"""The result of one step in a vectorized environment (all shapes `[B, ...]`)
:param obs: The observation to act upon at the next step (with same-step
reset, this is the first observation of a fresh episode when the
episode just ended). A tensor, or a `TensorStruct` for `Dict` spaces.
:param next_obs: The successor state $s_{t+1}$ of this step (the final
observation of the episode when it just ended)
:param reward: The reward
:param terminated: True if a terminal state was reached (no bootstrap)
:param truncated: True if the episode was cut short, e.g. by a time
limit (bootstrap)
"""
obs: Observation
next_obs: Observation
reward: Tensor
terminated: Tensor
truncated: Tensor
@property
def done(self) -> Tensor:
"""True when the episode just ended (terminated or truncated)"""
return self.terminated | self.truncated
[docs]
class VecEnv:
"""`num_envs` copies of a gymnasium environment, with torch tensors in/out
The environments are stepped **synchronously** (in the same process): one
call to `step` advances every copy by one step, and the returned tensors
have the number of environments as their leading dimension.
"""
def __init__(
self,
env_name: str,
num_envs: int,
seed: int | None = None,
same_step_reset: bool = False,
wrappers: Sequence[Callable[[gym.Env], gym.Env]] = (),
**kwargs,
):
"""Create the vectorized environment
:param env_name: The gymnasium environment name, e.g. `"CartPole-v1"`
:param num_envs: The number of parallel environments
:param seed: The random seed (used on the first `reset`)
:param same_step_reset: Reset ended episodes in the same step (see
the module documentation)
:param wrappers: gymnasium wrappers applied to *each* sub-environment
(e.g. to hide or extend observations); the same wrappers must be
used for training, evaluation and video recording
:param kwargs: Extra arguments passed to `gymnasium.make`
"""
#: The gymnasium environment name
self.env_name = env_name
#: The number of parallel environments (the batch size `B`)
self.num_envs = num_envs
#: Whether ended episodes are reset within the ending step
self.same_step_reset = same_step_reset
#: The wrappers applied to each sub-environment
self.wrappers = tuple(wrappers)
#: The extra arguments passed to `gymnasium.make`
self.env_kwargs = kwargs
vector_kwargs = {}
if same_step_reset:
vector_kwargs["autoreset_mode"] = gym.vector.AutoresetMode.SAME_STEP
#: The underlying gymnasium vector environment
self.gym_env = gym.make_vec(
env_name,
num_envs=num_envs,
vectorization_mode="sync",
vector_kwargs=vector_kwargs,
wrappers=self.wrappers,
**kwargs,
)
# How to turn a gymnasium observation into a tensor / TensorStruct
self._to_obs = _make_obs_converter(self.gym_env.single_observation_space)
self._seed = seed
@property
def observation_space(self) -> gym.Space:
"""The (single-environment) gymnasium observation space.
Use this to inspect the structure of `Dict` observations when writing
the encoder.
"""
return self.gym_env.single_observation_space
@property
def gym_unwrapped(self) -> tuple[gym.Env, ...]:
"""The underlying gymnasium sub-environments, unwrapped
The escape hatch to environment-specific behavior that `VecEnv` does
not (and should not) know about — typically rendering:
`env.gym_unwrapped[0].draw_v_pi(v, policy)`.
"""
return tuple(sub_env.unwrapped for sub_env in self.gym_env.envs)
@property
def observation_dim(self) -> int:
"""The dimension of the (flat) observation space.
Only defined for a flat `Box` observation space: use `n_states` for a
tabular (`Discrete`) space, and inspect `observation_space` for a
structured (`Dict`) one.
"""
space = self.gym_env.single_observation_space
assert isinstance(space, gym.spaces.Box), (
f"{self.env_name} does not have a flat observation space "
f"({type(space).__name__}); use "
f"{'`n_states`' if isinstance(space, gym.spaces.Discrete) else '`observation_space`'}"
" instead"
)
return int(np.prod(space.shape))
@property
def n_states(self) -> int:
"""The number of states of a tabular (`Discrete`) observation space
Observations are then `[B]` tensors of state indices (`torch.long`),
usable as-is to index a Q-table of shape `[n_states, n_actions]`.
"""
space = self.gym_env.single_observation_space
assert isinstance(
space, gym.spaces.Discrete
), f"{self.env_name} does not have a tabular (Discrete) observation space"
return int(space.n)
@property
def is_continuous(self) -> bool:
"""True if actions are continuous (a `Box` space)"""
return isinstance(self.gym_env.single_action_space, gym.spaces.Box)
@property
def action_dim(self) -> int:
"""The dimension of a continuous action space"""
space = self.gym_env.single_action_space
assert isinstance(space, gym.spaces.Box), f"{self.env_name} is not continuous"
return int(np.prod(space.shape))
@property
def n_actions(self) -> int:
"""The number of actions of a discrete action space"""
space = self.gym_env.single_action_space
assert isinstance(
space, gym.spaces.Discrete
), f"{self.env_name} is not discrete"
return int(space.n)
[docs]
def reset(self) -> Observation:
"""Reset all the environments, starting a fresh episode in each
Only needed once, at the very beginning: afterwards the environments
reset themselves when an episode ends (see the module documentation).
The collectors call it for you.
:returns: The initial observations, shape `[B, obs_dim]` for a flat
space, `[B]` (state indices) for a tabular one, or a
`TensorStruct` of `[B, ...]` tensors for a `Dict` space
"""
obs, _ = self.gym_env.reset(seed=self._seed)
self._seed = None # seed only the first reset
return self._to_obs(obs)
[docs]
def step(self, actions: Tensor) -> EnvStep:
"""Perform one step in every environment
:param actions: One action per environment: `[B, action_dim]`
(continuous) or `[B]` (discrete). The tensor is detached and moved
to the CPU, so gradients never flow into the environment.
:returns: The `EnvStep` of this step (all fields `[B, ...]`). Ended
episodes are reset automatically, in the mode chosen at
construction time.
"""
np_actions = actions.detach().cpu().numpy()
np_actions = np_actions.astype(self.gym_env.action_space.dtype, copy=False)
obs, reward, terminated, truncated, infos = self.gym_env.step(np_actions)
obs = self._to_obs(obs)
# With same-step reset, the final observation of an ended episode is
# reported in the infos; elsewhere the successor is `obs` itself
next_obs = obs
if "final_obs" in infos:
done = np.flatnonzero(infos["_final_obs"])
if len(done):
next_obs = _clone_obs(obs)
for i in done:
_set_obs_at(next_obs, i, self._to_obs(infos["final_obs"][i]))
return EnvStep(
obs=obs,
next_obs=next_obs,
reward=torch.as_tensor(reward, dtype=torch.float32),
terminated=torch.as_tensor(terminated, dtype=torch.bool),
truncated=torch.as_tensor(truncated, dtype=torch.bool),
)