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
¶
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
¶
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
|
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.
RescaleObservation
¶
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.
DtypeObservation
¶
Bases: ObservationWrapper[NDArray[Any], ActType, NDArray[Any]]
Casts Box observations to dtype, e.g. float64 simulators feeding a float32 policy.
FilterObservation
¶
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
|
'tail'
|
ActionMaskObservation
¶
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
¶
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
¶
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
¶
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
¶
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
|
Rewards¶
NormalizeReward
¶
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.