fix: placeholder stats for deferred video keys respect channel axis#164
Merged
Merged
Conversation
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>
akshay18iitg
approved these changes
Apr 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
nameslist (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