Skip to content

feat(lerobot_v3): add image-mode output to the v3 writer - #6

Merged
griswaldbrooks merged 1 commit into
mainfrom
fm/forge-image-mode
Aug 22, 2026
Merged

feat(lerobot_v3): add image-mode output to the v3 writer#6
griswaldbrooks merged 1 commit into
mainfrom
fm/forge-image-mode

Conversation

@griswaldbrooks

@griswaldbrooks griswaldbrooks commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Adds an image-mode output option to the LeRobot v3 writer: camera frames can be stored as
PNG images inline in the data parquet instead of being encoded to MP4. Video mode remains
the default and its output is unchanged.

What changed

Writer (forge/formats/lerobot_v3/writer.py) — new LeRobotV3WriterConfig.video flag
(default True), following upstream LeRobot's own video=False convention. With
video=False:

  • each camera frame is PNG-encoded into the data parquet as a struct<bytes, path> column,
    matching the HF datasets.Image() storage LeRobot uses for image datasets;
  • the feature is declared in info.json as dtype: "image";
  • the video encoder is never called and no videos/ directory is created;
  • video_path and video_files_size_in_mb are null, with the keys retained because v3
    loaders require them to exist.

Pillow is required for image mode (added to the [lerobot] extra) and checked up front.

Reader (forge/formats/lerobot_v3/reader.py) — camera features with dtype: "image" are
decoded from the inline PNG bytes (or a referenced path) through a lazy PIL loader, so
image-mode datasets round-trip through Forge.

Converter (forge/convert/converter.py, forge/config/models.py) — writer_config is
carried through ConversionConfig.to_dict() and applied in the parallel worker, so
--workers > 1 honours image mode. The sequential and parallel paths now share one
_configure_writer helper.

Usage:

fps: 10
robot_type: pusht
writer_config:
  video: false

Tests

  • tests/test_lerobot_v3_writer.py — image-mode structure (info.json dtype, no videos/,
    frame counts, null video_path), plus a video-mode-unchanged regression class.
  • tests/test_lerobot_v3_roundtrip.py — image-mode frames survive write → read through the
    real v3 reader.
  • tests/test_converter.pywriter_config round-trips through to_dict/from_dict, and
    the parallel worker actually applies video=False.
  • tests/test_lerobot_v3_upstream_validity.py — extended with a written_v3_image fixture and
    a TestUpstreamLoaderImageMode class, gated on lerobot being installed like the existing
    upstream-validity tests.

Verification against stock LeRobot

lerobot is not a test dependency, so the spec check was run manually. In a scratch venv with
lerobot 0.4.4, the real 206-episode pusht zarr sample was converted through
forge convert --config in image mode and loaded with stock LeRobotDataset:

lerobot metadata: robot_type=pusht fps=10 episodes=206 frames=25650
LeRobotDataset len: 25650
item 120: observation.images.img tensor (3, 96, 96) dtype=torch.float32 range=[0.255,1.000]

image_keys is populated and video_keys is empty. The gated upstream-validity suite passes
23/23 against real lerobot. Frames decoded by Forge's reader, by stock HuggingFace datasets
via Image(), and by LeRobotDataset are all pixel-identical to the source zarr.

Video mode was compared base vs branch on the same sample: every parquet and JSON artifact is
byte-identical. The MP4 files differ byte-wise, but decoded pixels match, and two runs of the
same build differ by the same margin — x264 run-to-run nondeterminism, not a behaviour change.

Known limitations

Filed as follow-ups, all pre-existing or narrow enough to ride a patch release:

Add `LeRobotV3WriterConfig.video` (default True). With `video=False` the
writer PNG-encodes each camera frame into the data parquet as a
`struct<bytes, path>` cell and declares the feature as `dtype: "image"` in
info.json, matching the LeRobot v3 image-dataset layout and upstream
LeRobot's own `video=False` convention. The video encoder is bypassed
entirely in that mode and no videos/ directory is created; `video_path`
and `video_files_size_in_mb` are null while the keys are retained.

Video mode remains the default and its output is unchanged.

Also in this change:
- carry `writer_config` through `ConversionConfig.to_dict()` and apply it in
  the parallel conversion worker, so `--workers > 1` honours image mode;
- read inline image cells back in the v3 reader, tolerating null and
  path-only cells;
- keep late-appearing parquet columns by filling the union of row keys
  before building a chunk table;
- skip null rows when computing episode stats and when inferring features;
- bound finalize's feature sampling to one batch instead of reading whole
  PNG payloads;
- share one `_configure_writer` helper between the sequential and parallel
  converter paths.

Tests cover image-mode dataset structure, a read-back through the
LeRobot v3 reader, and a video-mode regression class.
@griswaldbrooks griswaldbrooks changed the title feat(lerobot_v3): add PNG image-mode output to the v3 writer, reader, and parallel path feat(lerobot_v3): add image-mode output to the v3 writer Aug 22, 2026
@griswaldbrooks
griswaldbrooks merged commit 18886a1 into main Aug 22, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant