Skip to content

Wrappers

All wrappers are importable from sekai.wrappers and from the top-level sekai package. RecordFrames lives in sekai.rendering. See the Wrappers guide for an overview table.

Episode control

TimeLimit

TimeLimit(env: Env[ObsType, ActType], max_episode_steps: int)

Bases: Wrapper[ObsType, ActType, ObsType, ActType]

Truncates episodes after max_episode_steps steps.

Adds info["elapsed_steps"] to every step result.

elapsed_steps property

elapsed_steps: int

Steps taken in the current episode.

AutoReset

AutoReset(env: Env[ObsType, ActType])

Bases: Wrapper[ObsType, ActType, ObsType, ActType]

Automatically resets the environment when an episode ends.

On terminal steps, the observation in StepResult is the first observation of the new episode, not the terminal observation. The terminal observation and info are stored in info["final_observation"] and info["final_info"], and the reset's info in info["reset_info"].

Make AutoReset the outermost wrapper (AutoReset(TimeLimit(env, 100)), not TimeLimit(AutoReset(env), 100)). Its reset() calls go to the layers inside it only, so a wrapper outside it that keeps per-episode state (TimeLimit, TimeAwareObservation, FrameStackObservation, DelayObservation, DelayAction, StickyAction) is never reset and carries one episode's state into the next: TimeLimit would truncate every step after the first limit. Those wrappers warn when they are built around an AutoReset. RecordEpisodeStatistics works on either side of it.

OrderEnforcing

OrderEnforcing(env: Env[ObsType, ActType])

Bases: Wrapper[ObsType, ActType, ObsType, ActType]

Raises ResetRequired if step() is called before the first reset().

Enforces proper API usage — essential for catching subtle bugs during development, especially in distributed RL setups.

RecordEpisodeStatistics

RecordEpisodeStatistics(env: Env[ObsType, ActType], tracker: StatsTracker | None = None, deque_size: int = 100)

Bases: Wrapper[ObsType, ActType, ObsType, ActType]

Records per-episode statistics via a pluggable StatsTracker.

On episode termination, info["episode"] is set to an EpisodeStats dataclass (not an untyped dict), so type checkers and IDEs can verify access.

episode_stats property

episode_stats: EpisodeStats | None

The most recently completed episode's stats.

Observations

FlattenObservation

FlattenObservation(env: Env[Any, ActType])

Bases: ObservationWrapper[NDArray[float32], ActType, Any]

Flattens the observation from a nested space into a 1D float32 array.

NormalizeObservation

NormalizeObservation(env: Env[Any, ActType], epsilon: float = 1e-08, *, stats: StatsKind = 'running', halflife: float = 1000.0, clip: float | None = None)

Bases: ObservationWrapper[Any, ActType, Any]

Normalises observations to have approximately zero mean and unit variance.

Works on any observation space: every Box gets its own running statistics, a MaskedBox gets per-feature statistics over its occupied slots only (so padding does not bias them), and other parts of a Dict/Tuple observation (Discrete, MultiBinary, ...) pass through unchanged. Normalised parts are float64.

stats="running" (the default) weights every observation ever seen equally. On non-stationary data, such as prices whose volatility changes by regime, use stats="ewm" with a halflife (in steps), which tracks the recent distribution instead.

Set update_running_mean = False to freeze the statistics, e.g. for evaluation, and use state_dict() / load_state_dict() to save and restore them. For a Box observation obs_rms is the statistics object itself; for nested spaces it mirrors the nesting (a dict for a Dict, ...).

Parameters:

Name Type Description Default
env Env[Any, ActType]

The environment to wrap.

required
epsilon float

Added to the variance before dividing.

1e-08
stats StatsKind

"running" or "ewm".

'running'
halflife float

Halflife in steps for stats="ewm".

1000.0
clip float | None

If given, normalised values are clipped to [-clip, clip].

None

obs_rms property

obs_rms: Any

The running statistics (nested like the observation space).

state_dict

state_dict() -> Any

The normalisation statistics as plain data.

load_state_dict

load_state_dict(state: Any) -> None

Restore statistics saved by state_dict().

FrameStackObservation

FrameStackObservation(env: Env[Any, ActType], num_stack: int, *, padding: Literal['reset', 'zeros'] = 'reset')

Bases: Wrapper[Any, ActType, Any, ActType]

Stacks the last N observations along a new leading axis.

Gives the agent temporal context: recent prices for a policy without memory, or motion for pixel observations. Any observation space that can be batched works, and nested spaces stack leaf by leaf, so a Dict observation of {"returns": (L, n), "weights": (n + 1,)} becomes {"returns": (N, L, n), "weights": (N, n + 1)}; a MaskedBox keeps its values/mask layout with a leading stack axis on both.

padding: what fills the stack right after reset(). "reset" (the default) repeats the first observation; "zeros" uses the space's zero sample (sekai.spaces.zero_sample), which makes the start of an episode explicit to the agent.

Input obs shape: (d,) or (h, w, c) Output obs shape: (N, d) or (N, h, w, c)

TransformObservation

TransformObservation(env: Env[Any, ActType], func: Callable[[Any], Any] | None = None, new_observation_space: Space[Any] | None = None, *, fn: Callable[[Any], Any] | None = None)

Bases: ObservationWrapper[Any, ActType, Any]

Applies an arbitrary callable transformation to every observation.

func was called fn before sekai 0.2; fn= still works but is deprecated.

Example
env = TransformObservation(env, lambda obs: obs / 255.0)

RescaleObservation

RescaleObservation(env: Env[NDArray[Any], ActType], min_obs: ArrayLike, max_obs: ArrayLike)

Bases: ObservationWrapper[NDArray[Any], ActType, NDArray[Any]]

Affinely maps a bounded Box observation onto [min_obs, max_obs].

Integer observations (e.g. uint8 pixels) are rescaled to float32; floating observations keep their dtype.

Example
env = RescaleObservation(env, -1.0, 1.0)

DtypeObservation

DtypeObservation(env: Env[NDArray[Any], ActType], dtype: DTypeLike)

Bases: ObservationWrapper[NDArray[Any], ActType, NDArray[Any]]

Casts Box observations to dtype, e.g. float64 simulators feeding a float32 policy.

FilterObservation

FilterObservation(env: Env[dict[str, Any], ActType], keys: Sequence[str])

Bases: ObservationWrapper[dict[str, Any], ActType, dict[str, Any]]

Keeps only the given keys of a Dict observation.

Example
env = FilterObservation(env, ["prices", "holdings"])

TimeAwareObservation

TimeAwareObservation(env: Env[Any, ActType], max_episode_steps: int | None = None, *, normalize_time: bool = True, dict_key: str = 'time')

Bases: ObservationWrapper[Any, ActType, Any]

Adds the episode's elapsed time to the observation.

Without it, a time-limited episode is not Markov: the agent cannot tell that it is about to be truncated (e.g. the end of a trading session).

For a 1-D Box observation the time feature is appended as one extra element; for a Dict observation it is added under the key "time". With normalize_time=True (the default) the feature is t / max_episode_steps in [0, 1]; max_episode_steps is read from an inner TimeLimit or the env's spec when not given.

DelayObservation

DelayObservation(env: Env[Any, ActType], delay: int, *, padding: Literal['zeros', 'reset'] = 'zeros')

Bases: ObservationWrapper[Any, ActType, Any]

Delays observations by delay steps, modelling feed or decision latency.

Until delay observations have accumulated, the wrapper returns padding: "zeros" (the space's zero point, clipped into bounds) or "reset" (the episode's first observation, often more natural for price series).

PadSequenceObservation

PadSequenceObservation(env: Env[Any, ActType], max_length: int | None = None, *, truncate: Truncate = 'tail')

Bases: ObservationWrapper[Any, ActType, Any]

Turns variable-length Sequence observations into fixed-size MaskedBox ones.

Every Sequence of Box elements in the observation space (at the top level or inside a Dict/Tuple) becomes a MaskedBox of max_length slots, with the elements in order in the first slots. The result has fixed shapes, so it can be vectorised, flattened, normalised and fed to a standard network.

Parameters:

Name Type Description Default
env Env[Any, ActType]

The environment to wrap.

required
max_length int | None

Slots per sequence. Defaults to each Sequence's own max_length.

None
truncate Truncate

What to do with a longer sequence: keep the most recent max_length elements ("tail", the default, right for event streams), keep the first ones ("head"), or raise ValueError ("error").

'tail'

ActionMaskObservation

ActionMaskObservation(env: Env[Any, ActType], *, key: str = 'action_mask')

Bases: ObservationWrapper[dict[str, Any], ActType, Any]

Adds the env's current action mask to the observation.

For policies that apply the mask themselves (masked logits in PPO, RLlib's action-masking models). The mask is flattened to an int8 MultiBinary (see sekai.utils.flatten_action_mask). A Dict observation gains a key entry; any other observation becomes {"observation": obs, key: mask}.

The wrapped env must implement action_mask() (see sekai.utils.masking).

Actions

ClipAction

ClipAction(env: Env[ObsType, NDArray[float32]])

Bases: ActionWrapper[ObsType, NDArray[float32], NDArray[float32]]

Clips continuous actions to the env's action space bounds.

Useful when a policy can output values outside the valid range.

RescaleAction

RescaleAction(env: Env[ObsType, NDArray[float32]], min_action: float | NDArray[float32], max_action: float | NDArray[float32])

Bases: ActionWrapper[ObsType, NDArray[float32], NDArray[float32]]

Rescales a continuous action from [min_action, max_action] to the env's range.

Useful when your policy outputs actions in a fixed range (e.g. [-1, 1]) but the environment expects a different range. Rescaled actions are clipped to the env's bounds (and reverse_action() results to [min_action, max_action]).

If the env has an action mask (see sekai.utils.masking), the components it masks reach the env as the env's zero (action_space.zero() of the env), as the mask convention requires, whatever the policy put there: rescaled, the wrapper's zero (what action_space.sample(mask=...) puts in a masked component) would be the middle of the env's range instead. reverse_action() does not apply the mask.

StickyAction

StickyAction(env: Env[ObsType, ActType], repeat_action_probability: float)

Bases: ActionWrapper[ObsType, ActType, ActType]

With probability repeat_action_probability, repeats the previous action.

Models execution that sometimes lags the agent's decision (Machado et al., 2018, for Atari; order-routing delay in trading). The wrapper keeps its own RNG, derived from the reset seed (an int, or a Generator's SeedSequence), so it does not perturb the env's stream.

ActionRepeat

ActionRepeat(env: Env[ObsType, ActType], repeat: int)

Bases: Wrapper[ObsType, ActType, ObsType, ActType]

Repeats each action repeat times, summing the rewards.

Stops early if the episode ends. The returned observation and info are from the last inner step; info["action_repeats"] says how many ran. Useful to decouple the agent's decision interval from the simulator's tick (e.g. decide every 5 minutes on a 1-minute bar simulator).

DelayAction

DelayAction(env: Env[ObsType, ActType], delay: int, *, noop: Any = None)

Bases: ActionWrapper[ObsType, ActType, ActType]

Executes each action delay steps after the agent chooses it.

Models order latency: the decision made on bar t reaches the market on bar t + delay, so a policy cannot trade on the very price it just observed. For the first delay steps of an episode the env receives noop.

Parameters:

Name Type Description Default
env Env[ObsType, ActType]

The environment to wrap.

required
delay int

Steps between choosing and executing an action (0 = no delay).

required
noop Any

Action executed while the pipeline fills. Defaults to the zero action for Box action spaces (e.g. all-cash or no-trade, depending on the env); required for other spaces, because index 0 of a Discrete space is rarely "do nothing".

None

pending_actions property

pending_actions: list[ActType]

Actions chosen but not yet executed, oldest first.

Rewards

NormalizeReward

NormalizeReward(env: Env[ObsType, ActType], gamma: float = 0.99, epsilon: float = 1e-08)

Bases: Wrapper[ObsType, ActType, ObsType, ActType]

Scales rewards by a running estimate of the discounted return's std.

As in OpenAI baselines' VecNormalize, the mean is NOT subtracted, only the scale is adjusted. The discounted return is reset at the end of every episode. Set update_running_mean = False to freeze the statistics.

TransformReward

TransformReward(env: Env[ObsType, ActType], func: Callable[[float], float] | None = None, *, fn: Callable[[float], float] | None = None)

Bases: RewardWrapper[ObsType, ActType]

Applies an arbitrary callable transformation to every reward.

func was called fn before sekai 0.2; fn= still works but is deprecated.

Example
env = TransformReward(env, lambda r: np.sign(r))  # reward clipping
env = TransformReward(env, lambda r: r * 0.01)    # reward scaling

ClipReward

ClipReward(env: Env[ObsType, ActType], min_reward: float | None = -inf, max_reward: float | None = inf)

Bases: RewardWrapper[ObsType, ActType]

Clips every reward to [min_reward, max_reward].

A bound of None (as in gymnasium) or infinity leaves that side unbounded.

Example
env = ClipReward(env, -1.0, 1.0)
env = ClipReward(env, None, 0.5)  # cap gains only