nemo_automodel.components.datasets.wam.episode_windows

Training windows from recorded single-arm episodes (.npz files).

An episode file holds canvas (frames, height, width, 3) uint8, obs8 and cmd8 (frames, 8) float32 flange poses [xyz, quat_wxyz, grip] and prompt. A window is window_frames consecutive frames; its actions are the commanded poses relative to the observed pose at the start of each chunk, and its proprio the observed pose at each chunk start, the same layout as the reference DROID windows. The LIBERO tutorial’s convert_libero_demos.py writes such files.

Module Contents

Classes

Name	Description
`EpisodeWindowDataset`	Endless random windows over a directory of recorded episodes.

Functions

Name	Description
`_mapping`	-

API

class nemo_automodel.components.datasets.wam.episode_windows.EpisodeWindowDataset(
    root: str | pathlib.Path,
    model_config: collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue],
    seed: int = 0,
    rank: int = 0,
    world_size: int = 1
)

Bases: IterableDataset

Endless random windows over a directory of recorded episodes.

Parameters:

root

str | Path

Directory with *.npz episodes.

model_config

Mapping[str, JsonValue]

The checkpoint’s config.json; its chunking section gives the chunk geometry and its video section the canvas size the episodes must have been rendered at.

seed

intDefaults to 0

Base seed; the draws also depend on the rank and the loader worker.

rank

intDefaults to 0

This data-parallel rank; episodes are split round-robin across ranks.

world_size

intDefaults to 1

Number of data-parallel ranks.

canvas_hw

= (int(video['height']), int(video['width']))

files

= sorted(Path(root).glob('*.npz'))[rank::world_size]

geometry

window_frames

= sum(self.geometry[-1])

nemo_automodel.components.datasets.wam.episode_windows.EpisodeWindowDataset.__iter__() -> collections.abc.Iterator[praxis.train.data.TrainingWindow]

nemo_automodel.components.datasets.wam.episode_windows.EpisodeWindowDataset.window(
    file: pathlib.Path,
    start: int
) -> praxis.train.data.TrainingWindow

nemo_automodel.components.datasets.wam.episode_windows._mapping(
    config: collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue],
    name: str
) -> collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue]