Skip to content

[python] Extend ContiguousWindowDataset for frame-level training - #9875

Open
YannByron wants to merge 2 commits into
apache:masterfrom
YannByron:codex/cwd-frame-training
Open

YannByron wants to merge 2 commits into
apache:masterfrom
YannByron:codex/cwd-frame-training

Conversation

@YannByron

Copy link
Copy Markdown
Contributor

Purpose

Extend ContiguousWindowDataset (CWD) into a training reader for a materialized, frame-level multimodal table. Previously, CWD used one forward window for all non-anchor columns. A sample containing state history, a current camera image, and future actions needs independent windows for those fields.

The two Dataset APIs serve different data layouts:

  • CWD reads one Paimon frame table. Each row is one time step, with task text, labels, and other training context already materialized. It constructs samples without discovering or joining companion tables.
  • PaimonLeRobotDataset (PLD) serves LeRobot table groups, combining frame data with companion metadata such as episodes, tasks, info, and stats to provide LeRobot-compatible samples.

They should converge on training capabilities while retaining these distinct data-source responsibilities. This PR focuses on CWD; PLD only reuses the image conversion helpers here.

This PR adds:

  • Per-column windows: frame_offsets supports history, future, sparse, and repeated positions, preserving their order. Unspecified selected columns default to [0]. For example, frame_offsets={"state": [-2, -1, 0], "action": range(16)} combines state history and future actions while a selected image column stays at the anchor.
  • Offsets in seconds: delta_timestamps converts offsets using an explicit fps and validates frame-grid alignment with tolerance_s.
  • Boundaries and masks: boundary="drop"|"pad"|"error" applies at both ends of each group. Padding repeats the nearest endpoint or uses per-column raw pad_values, with a <column>_is_pad mask. Padding precedes transforms.
  • Multimodal conversion: reusable numeric tensor and image-to-TCHW transforms, including EXIF orientation, uint8 handling, and preservation of higher-bit-depth pixel units. Existing callbacks support explicit normalization and model-specific sample adaptation.
  • Sample access: slices and batched reads preserve order and duplicates. Reads are coalesced by column window; constant-padded positions avoid payload reads. Snapshot/tag pinning and lazy BLOB loading are retained.
  • Training integration: the ACT benchmark uses field offsets while preserving its model inputs and loss. Documentation covers normalization, multi-worker loading, and standard PyTorch distributed samplers.

Existing window_size / anchor_columns / tail calls retain their output shapes and global is_pad mask. New offset-based calls use per-column masks. Pillow is added to the optional torch extra for image conversion.

This does not yet unify PLD's window API or add CWD video decoding. CWD continues to require contiguous integer frame positions; irregular-timestamp interpolation and automatic cross-table metadata/statistics discovery are outside this change.

Tests

152 tests passed across the affected CWD, transform, PLD, and ACT suites:

cd paimon-python
python -m pytest -q \
    pypaimon/tests/contiguous_window_dataset_test.py \
    pypaimon/tests/window_transforms_test.py \
    pypaimon/tests/multimodal_lerobot_test.py \
    pypaimon/tests/act_runner_test.py \
    pypaimon/tests/act_benchmark_test.py

Coverage includes offset conversion and validation, both group boundaries, per-column masks and projections, batch/slice semantics, mutable-value isolation, worker serialization, distributed sampling, legacy compatibility, and HDF5/Paimon ACT tensor and loss parity.

Flake8 with the repository configuration and git diff --check also passed.

YannByron and others added 2 commits September 16, 2026 14:25
Build training samples from one materialized frame table with per-column
frame or time offsets, two-sided padding and per-column masks. Add tensor
and image transforms, adapt the ACT benchmark, and preserve legacy windows.

Co-Authored-By: Codex <noreply@anthropic.com>
AI-Model: gpt-6
Co-Authored-By: Codex <noreply@ai-tool.com>
AI-Contributed/Feature: 643/643
AI-Contributed/UT: 365/365
Use importorskip for torch and Pillow, matching the other optional training
suites. CI environments without Pillow can collect the remaining tests.

Verified 44 related tests, collection without Pillow or torch, and flake8.

Co-Authored-By: Codex <noreply@anthropic.com>
AI-Model: gpt-6
Co-Authored-By: Codex <noreply@ai-tool.com>
Co-Authored-By: Codex <noreply@openai.com>
AI-Contributed/Feature: 0/0
AI-Contributed/UT: 6/6
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