Environments — rl_mind.env

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.

rl_mind.env.Observation

An observation is either a plain tensor (flat Box space) or a TensorStruct of named tensors (Dict space)

alias of Tensor | TensorStruct

class rl_mind.env.EnvStep(obs: Tensor | TensorStruct, next_obs: Tensor | TensorStruct, reward: Tensor, terminated: Tensor, truncated: Tensor)

Bases: object

The result of one step in a vectorized environment (all shapes [B, ...])

Parameters:

• 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.

• next_obs – The successor state \(s_{t+1}\) of this step (the final observation of the episode when it just ended)

• reward – The reward

• terminated – True if a terminal state was reached (no bootstrap)

• truncated – True if the episode was cut short, e.g. by a time limit (bootstrap)

property done: Tensor

True when the episode just ended (terminated or truncated)

class rl_mind.env.VecEnv(env_name: str, num_envs: int, seed: int | None = None, same_step_reset: bool = False, wrappers: Sequence[Callable[[Env], Env]] = (), **kwargs)

Bases: object

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.

env_name

The gymnasium environment name

num_envs

The number of parallel environments (the batch size B)

same_step_reset

Whether ended episodes are reset within the ending step

wrappers

The wrappers applied to each sub-environment

env_kwargs

The extra arguments passed to gymnasium.make

gym_env

The underlying gymnasium vector environment

property observation_space: Space

The (single-environment) gymnasium observation space.

Use this to inspect the structure of Dict observations when writing the encoder.

property gym_unwrapped: tuple[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).

property observation_dim: 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.

property n_states: 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].

property is_continuous: bool

True if actions are continuous (a Box space)

property action_dim: int

The dimension of a continuous action space

property n_actions: int

The number of actions of a discrete action space

reset() → Tensor | TensorStruct

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

step(actions: Tensor) → EnvStep

Perform one step in every environment

Parameters:

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.