Buffer

Online history buffer for policies that need past observations

HistoryBuffer is a per-env ring buffer over batched info dicts. It is used internally by WorldModelPolicy to feed strided history into the planner, but it can also be used standalone for any code path that consumes (n_envs, ...)-shaped step data and needs a sliding window over time.

[ How it works ]

Each call to append takes a dict whose values have a leading env dim of size n_envs (e.g. EnvPool's stacked infos with shape (n_envs, 1, ...)). The buffer slices the env dim and pushes one entry per env onto its corresponding deque.

get(n) returns the last n entries per env, strided by action_block env steps so history is surfaced at planning cadence regardless of frameskip. Output is in chronological order (oldest → newest). Pass env_ids to build history only for a subset of envs (e.g. the ones being re-planned).

Warm-up

Each env is padded independently, so the output shape is always the same and envs never constrain each other (a freshly-reset env does not shrink its neighbors' history):

Strided (non-block) keys are left-padded by repeating the env's oldest available entry.
Block keys are left-padded with zero blocks.

Warm-up output is partially synthetic

Padded entries are fake: until an env has lived (n - 1) * action_block + 1 steps, the consumer — in planning, the world model — receives duplicated copies of the episode's first frame with zero actions between them, as if the env had been stationary before the episode began. This is a deliberate trade-off: training clips are never padded, so any warm-up scheme is off-distribution, and repeating the first frame (the repo's forward-fill precedent, cf. PreJEPA._encode_video) with zero actions is the self-consistent choice that keeps histories stackable across desynchronized envs.

get returns None only when some selected env is empty (or n <= 0).

Output shapes

Per-step value shape	Output shape
`(n_envs,)`	`(n_envs, n)`
`(n_envs, T, ...)`	`(n_envs, n * T, ...)`
`(n_envs, T, ...)` and key in `block_keys`	`(n_envs, n - 1, action_block * D)` where `D` is the flattened per-env per-step size

block_keys is intended for actions when the planner operates at macro-action cadence. Block i concatenates (flat, chronological) the action_block raw entries recorded between strided frames i and i + 1 — "the block leaving frame i" — matching the training convention that pairs action[t] with frame[t] (an env's info['action'] at step j echoes the action executed entering step j, so the entries strictly after frame i up to and including frame i + 1 are exactly the actions executed between the two frames). There are therefore n - 1 blocks for n frames; block keys are omitted entirely when n == 1. NaN entries (the reset frame's action echo) are zeroed.

[ Example ]

import numpy as np
from stable_worldmodel.buffer import HistoryBuffer

# 2 envs, hold up to 5 env steps, return history at frameskip-2 cadence
buf = HistoryBuffer(n_envs=2, max_len=5, action_block=2, block_keys=('action',))

for t in range(5):
    buf.append({
        'pixels': np.full((2, 1, 64, 64, 3), t, dtype=np.uint8),  # (n_envs, T, H, W, C)
        'action': np.full((2, 1, 2), t, dtype=np.float32),        # (n_envs, T, D)
    })

out = buf.get(3)
# out['pixels'].shape == (2, 3, 64, 64, 3)   # strided frames: t=0, t=2, t=4
# out['action'].shape == (2, 2, 4)           # blocks between them:
#                                            #   [a_1,a_2] (leaving t=0) and
#                                            #   [a_3,a_4] (leaving t=2)

buf.reset(env_ids=[0])  # clear env 0 only (e.g. on episode reset)

[ Reference ]

HistoryBuffer

HistoryBuffer(
    n_envs: int,
    max_len: int,
    action_block: int = 1,
    block_keys: Iterable[str] = (),
)

Per-env ring buffer over batched info dicts.

Inputs are dicts whose tensor/ndarray values have a leading env dim of size n_envs (e.g. EnvPool's stacked infos with shape (n_envs, 1, ...)). Each call to :meth:append stores one slice per env. :meth:get returns the last n entries strided by action_block env steps so history is surfaced at planning cadence regardless of frameskip.

.. warning:: Warm-up returns synthetic entries. While an env holds fewer than (n - 1) * action_block + 1 entries (the first steps after a reset), its history is left-padded with copies of its oldest real entry — for block keys, with zero blocks. The consumer (e.g. the world model) therefore sees fake repeated frames with zero actions between them, as if the env had been stationary before the episode began. Padding keeps the time dim at n for every env so histories stack across desynchronized envs without one env's reset truncating the others. All entries are real once the env has lived (n - 1) * action_block + 1 steps. Callers that would rather shrink n than receive padding can size their request with :meth:num_strided.

Output for a key with per-step shape (n_envs, T, ...) is (n_envs, n*T, ...) (concatenated along the time dim, oldest → newest). For a key with per-step shape (n_envs,) the output is (n_envs, n).

Keys listed in block_keys are instead returned as n - 1 macro-blocks: block i concatenates (flat, chronological) the action_block raw entries recorded after strided frame i up to and including frame i + 1's entry — the actions executed between frames i and i + 1 ("the block leaving frame i"), matching the training convention that pairs action[t] with frame[t]. With per-env per-step size D, the output is (n_envs, n - 1, action_block * D). Missing blocks are zero-padded on the left and NaNs (e.g. the reset entry's action) are zeroed. Block keys are omitted when n == 1 (there is no "between" with a single frame). Use this for actions when the planner operates at macro-action cadence.

Parameters:

n_envs (int) –

Number of parallel environments.
max_len (int) –

Maximum entries retained per env (in env steps).
action_block (int, default: 1 ) –

Frameskip — stride used when returning history.
block_keys (Iterable[str], default: () ) –

Keys to aggregate as macro-blocks (typically ('action',)).

append

append(info_dict: dict[str, Any]) -> None

Append one entry per env from a batched info dict.

Splits each value along dim 0 (the env dim) and pushes the per-env slice onto the corresponding deque.

Parameters:

info_dict (dict[str, Any]) –

Dict of values each with leading dim n_envs.

get

get(
    n: int, env_ids: list[int] | None = None
) -> dict[str, Tensor | ndarray] | None

Return the last n strided entries per env, padded in warm-up.

Entries are returned in chronological order (oldest → newest) along the time dim. An env holding fewer than n strided samples is left-padded with synthetic entries — copies of its oldest real entry (regular keys) or zero blocks (block_keys) — so the time dim is always n (n - 1 for block keys) regardless of how full each env's buffer is. See the class-level warning: during warm-up the output is partially fake.

Parameters:

n (int) –

Number of strided entries per env to return.
env_ids (list[int] | None, default: None ) –

Env indices to build history for (all if None); output rows follow this order.

Returns:

dict[str, Tensor | ndarray] | None –

A dict of stacked tensors/arrays with time dim n * T
dict[str, Tensor | ndarray] | None –

(T = per-step time dim) for regular keys and n - 1
dict[str, Tensor | ndarray] | None –

for block keys (omitted when n == 1). None if any
dict[str, Tensor | ndarray] | None –

selected env's buffer is empty or n <= 0.

reset

reset(env_ids: list[int] | None = None) -> None

Clear buffers for the given envs (all if None).