Skip to content

fix: placeholder stats for deferred video keys respect channel axis#164

Merged
shuheng-liu merged 1 commit into
mainfrom
fix/deferred-video-placeholder-channel-axis
Apr 21, 2026
Merged

fix: placeholder stats for deferred video keys respect channel axis#164
shuheng-liu merged 1 commit into
mainfrom
fix/deferred-video-placeholder-channel-axis

Conversation

@shuheng-liu
Copy link
Copy Markdown
Member

@shuheng-liu shuheng-liu commented Apr 20, 2026

What this does

LeRobotDataset.save_episode synthesizes placeholder statistics for deferred video keys so downstream stat-aggregation code never sees a missing entry. The fallback pulled the channel count from shape[0], which is only correct for (C, H, W) callers — LeRobot v2.1's default image convention is (H, W, C) with names ["height", "width", "channel"], where shape[0] is the height.

Concretely, a 480×640×3 deferred video key produced placeholder stats with shape (480, 1, 1). On the second save_episode call, aggregate_stats → _assert_type_and_shape rejected it because the assertion requires (3, 1, 1) for any feature with "image" in the key. Datasets with only one episode never hit this because aggregate_stats needs two entries to merge.

Fix: look up the channel axis via the feature's names list (handling both "channel" and "channels"), falling back to the existing shape[0] behavior when names are absent or don't mention a channel axis, and to 3 as a last resort. Preserves current behavior for every (C, H, W) caller, including the existing tests at tests/datasets/test_datasets.py that declare shape (3, 96, 128) with names ["channels", ...].

Adds a regression test that declares a deferred video key in (H, W, C) order and saves two episodes — this fails on main and passes with the fix.

How it was tested

Ran a test that declares a deferred video key in (H, W, C) order and saves two episodes — this fails on main and passes with the fix.

How to checkout & try? (for the reviewer)

Run a test that declares a deferred video key in (H, W, C) order and saves two episodes — this fails on main and passes with the fix.

Checklist

  • I have added Google-style docstrings to important functions and ensured function parameters are typed.
  • My PR includes policy-related changes.
    • If the above is checked: I have run the GPU pytests (pytest -m "gpu") and regression tests.

LeRobotDataset.save_episode synthesizes placeholder statistics for
deferred video keys so downstream stat-aggregation code never sees a
missing entry. The fallback pulled the channel count from shape[0],
which is only correct for (C, H, W) callers — LeRobot v2.1's default
image convention is (H, W, C) with names ["height", "width", "channel"],
where shape[0] is the height.

Concretely, a 480×640×3 deferred video key produced placeholder stats
with shape (480, 1, 1). On the second save_episode call, aggregate_stats
→ _assert_type_and_shape rejected it because the assertion requires
(3, 1, 1) for any feature with "image" in the key. Datasets with only
one episode never hit this because aggregate_stats needs two entries
to merge.

Fix: look up the channel axis via the feature's ``names`` list (handling
both "channel" and "channels"), falling back to the existing shape[0]
behavior when names are absent or don't mention a channel axis, and to
3 as a last resort. Preserves current behavior for every (C, H, W)
caller, including the existing tests at tests/datasets/test_datasets.py
that declare shape (3, 96, 128) with names ["channels", ...].

Adds a regression test that declares a deferred video key in (H, W, C)
order and saves two episodes — this fails on main and passes with the
fix.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@shuheng-liu shuheng-liu self-assigned this Apr 20, 2026
@shuheng-liu shuheng-liu merged commit 1491b38 into main Apr 21, 2026
6 checks passed
@shuheng-liu shuheng-liu deleted the fix/deferred-video-placeholder-channel-axis branch April 21, 2026 02:46
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.

2 participants