Trajectory store

StableTested and in active use. Breaking changes will be called out.

TrajectoryStore is the on-disk format for anything episodic that outlives a training run: offline datasets, demonstration capture, replay dumps, and the trajectory data behind zero-shot RL. It is one .h5 file with named columns and an episode index.

__manifest__  uint8   (M,)           UTF-8 key=value schema block
ep_len        int64   (N_ep,)        episode lengths
ep_offset     int64   (N_ep,)        start row of each episode
<column>      any     (N_rows, …)    one dataset per declared column

action, reward, done and truncated are reserved names. Everything else is declared by whoever writes the store — qpos/qvel for dm_control, pixels/proprio/state for PushT, ram for Atari.

The layout is deliberately a superset of the schema the PushT expert dataset already uses, so LeRobot-style datasets stay cheap to ingest — see importers below.

Writing

Declare the columns, append rows, mark episode boundaries, close. Every dataset is created chunked with an unlimited leading axis, so a collection run streams to disk instead of buffering in RAM:

from noeira.data import ColumnSpec, TrajectoryStoreWriter

var cols = List[ColumnSpec]()
cols.append(ColumnSpec(String("state"), DType.float32, DIM))

var w = TrajectoryStoreWriter(
    String(path), cols^, env_id=String("walker-walk"), seed=99,
)
w.append[DType.float32](String("state"), buf.unsafe_ptr().as_unsafe_any_origin(), N_ROWS)
w.end_episode()
w.close()

Two invariants are enforced rather than documented:

• dtype is checked against the declared spec on every append. A mismatch raises instead of reinterpreting the bytes, which is the failure mode that would otherwise corrupt a store silently.
• Every column must be appended in step. They share one row axis, so letting the counts drift would make the episode index meaningless. end_episode and close both check.

ep_len, ep_offset and the manifest are written at close().

A complete round trip — write, push, pull, reopen, verify — is in examples/data/dataset_remote_smoke.mojo.

Reading

from noeira.data import TrajectoryStore

var s = TrajectoryStore(String(path))
print(s.n_rows(), "rows,", s.n_episodes(), "episodes")

var state = s.load_column[DType.float32](String("state"))   # whole column, resident
s.read_range[DType.float32](String("state"), 0, 4096, buf)  # contiguous slab

The episode index is validated against the column length on open, not lazily: a store whose index disagrees with its columns produces samplers that read the wrong rows without ever raising.

Foreign files

A file we did not write — a HuggingFace dataset with no manifest of ours — opens the same way. The schema is recovered by enumerating the datasets and introspecting each one’s shape and dtype. This is gated against the real 44 GB pusht_expert_train.h5, not only against a fixture.

Importers

Demonstration datasets from elsewhere arrive as stores, through native importers with no Python in the path. Each one is gated against an independent implementation.

Source	Importer	Run it with	Columns
LeRobot v3.0 (SO-101 recordings, Hub datasets)	noeira/data/lerobot.mojo	examples/so101/act_so101_import_dataset.mojo	qpos, action, task_index, images
LIBERO’s demonstration HDF5	noeira/envs/libero/demos.mojo	pixi run libero-demo-import	state, joint_states, gripper_states, qpos, action, task_index, images
LAFAN1 motion for the Unitree G1	noeira/data/lafan.mojo	examples/g1/lafan_import.mojo	the reference’s motion states, column by column

LeRobot has its own page: LeRobot datasets.

LIBERO demonstrations

pixi run libero-demo-import                           # libero_goal, with images
pixi run libero-demo-import --no-images --demos 2     # a smoke run
pixi run libero-demo-import libero_goal --out build/demos/libero_goal.h5

It reads the suite’s demonstration files from references/libero_demos/<suite> and writes build/demos/<suite>.h5 (.lowdim.h5 with --no-images). --demos N caps the demonstrations per task. The columns:

• state is LIBERO’s recorded (time, qpos, qvel) rewritten into noeira’s joint addresses. A wrong remap would produce a file of the right shape, so the run checks the rows against the Python remap’s dumps (pixi run libero-demo-dump), to within 2 ULP, and prints how many values it compared.
• qpos is what the image policies read: joint_states (7) and gripper_states (2), then their one-step difference (9 more, zero on an episode’s first row), then a one-hot of the task. Its width is 18 + n_tasks.
• images holds the two recorded cameras, agent view and eye-in-hand.

The difference and the one-hot are there because a single pose does not identify the phase of a motion, and in libero_goal the picture does not identify the task: every task is the same scene at the same layout. See LIBERO for the environment and ACT for the policy trained on it.

.demo files are not stores

Transitions recorded for an off-policy replay — teleoperation and DAgger interventions on the SO-101 tabletop — go into a flat .demo file (noeira/deep_agents/demos/file.mojo), not a trajectory store. It holds (obs, action, reward, next_obs, done) rows with an intervention flag, and is loaded as the replay’s pinned demo prefix. See HIL-SERL.

Residency

Sampling reads from memory, never from disk. ResidentColumn loads a whole column and optionally mirrors it to the device:

from noeira.data import IndexBatch, ResidentColumn

var col = ResidentColumn[DType.float32].load(store, String("qpos"))
col.to_device(ctx)                              # optional — gather does it lazily
col.gather_device(ctx, idx, dst, dst_elems)     # dst[i, :] = column[idx[i], :]

The device view uses a runtime layout — Layout.row_major[N]() with the real extents carried in a RuntimeLayout — because n_rows and row_dim are store metadata, not compile-time constants.

WindowedColumn covers the case a column does not fit: it keeps a sliding slab resident and gathers within it.

var w = WindowedColumn[DType.uint8].make(store, String("pixels"), window_rows)
w.seek(store, start)        # position the window
_ = w.advance(store)        # next non-overlapping window; False when exhausted

A gather outside the resident window raises. There is deliberately no refill-on-miss: a silent refill turns one stray index into a disk read per minibatch element, which is precisely the 4500× case the measurements ruled out.

Index policies

A sampler produces an IndexBatch and knows nothing about columns. That is what makes them composable — the dual sampler a forward-backward representation needs is two calls to the same policy, not a new buffer.

Sampler	Draws	Notes
UniformSampler	host uniform rows	the baseline
UniformDeviceSampler	device uniform rows	one Philox stream per lane, seed + lane
PrioritizedSampler	host sum-tree	note_added, update_priorities, IS weights
SequenceWindowSampler	window starts	draw_starts + expand_window
WindowSampler	rows inside a resident window	pairs with WindowedColumn

Sequence windows span episode boundaries by default

SequenceWindowSampler defaults to within_episode=False. That is not an oversight — DreamerV3 relies on it, resetting its recurrent carry at the seam from the per-frame is_first flags. Making non-spanning the default would silently change its training distribution, so a test pins the default. Consumers that treat a window as one trajectory opt in via draw_starts_in_episodes.

The HDF5 layer

noeira/io/hdf5/ is a native FFI over libhdf5 — no Python in the read or write path. Each H5 sub-API lives in its own file, with reader.mojo and writer.mojo providing the H5File / H5Dataset / H5Writer surface the store consumes.

Compression filters shipped by the hdf5plugin package (Blosc, LZ4, ZSTD) are registered at first open via H5PLprepend, so compressed third-party files decode without HDF5_PLUGIN_PATH being set in the environment.

Compression policy: off for float32 state columns. Measured, deflate=4 reads at 545 MiB/s against 4377 MiB/s raw — 8× slower on data that barely compresses. Shuffle+deflate is worth it for uint8 pixel columns and nowhere else without measuring first.

Two silent traps in the C API, both now handled

H5T_NATIVE_FLOAT is not a constant. It is a macro expanding to a global hid_t variable plus an H5open() call the preprocessor inserts — and an FFI caller gets neither. Before H5open() runs, every such global reads -1, and the failure surfaces much later as “not a datatype”.

Worse, get_symbol returns a pointer to the global, so the value is p[][]. Writing Int(p[]) compiles and yields the symbol’s address — a plausible large integer that is not the id. The first probe printed 4659917064 and looked entirely reasonable.

h5native.mojo latches H5open() and resolves the ids properly. The write gate cross-checks against h5py rather than against our own reader: a round trip through one implementation proves nothing about the format.

Gates

Gate	What it pins
tests/io/hdf5/test_writer_roundtrip.mojo	Mojo and h5py agree on shape, dtype, chunking, maxshape and values, across chunk boundaries and with shuffle+deflate
tests/data/test_trajectory_store.mojo	Manifest round trip, exact values, dtype mismatch raises, boundary logic, and ingest of a real foreign dataset
tests/data/test_resident_gather.mojo	Device gather bit-exact against host — adversarial indices, duplicates, out-of-range raises
tests/data/test_windowed.mojo	Windowed gather equals full residency across every window; out-of-window raises; advance() exhaustion
tests/data/test_sampler_golden.mojo	Each index policy’s sequence under a fixed seed