Skip to content

Fix episode continuity across sampler batch slots - #8

Open
JLUbaijh2323 wants to merge 1 commit into
InternRobotics:mainfrom
JLUbaijh2323:fix/persistent-episode-batch-slots
Open

Fix episode continuity across sampler batch slots#8
JLUbaijh2323 wants to merge 1 commit into
InternRobotics:mainfrom
JLUbaijh2323:fix/persistent-episode-batch-slots

Conversation

@JLUbaijh2323

Copy link
Copy Markdown

Summary

Add an optional persistent batch-slot layout to SequentialEpisodeBatchSampler
and enable it for the RoboTwin-MeM training configuration.

Problem

EventVLA keeps predicted runtime keyframe state by batch slot:

  • _runtime_keyframe_image_bank[slot_idx]
  • _runtime_pending_keyframe_writes[slot_idx]
  • _runtime_slot_episode_ids[slot_idx]

The current sampler instead flattens every episode into one sample stream before
splitting that stream into batches. With batch size 4, one episode is arranged
approximately as:

batch 0: ep0/t0,   ep0/t50,  ep0/t100, ep0/t150
batch 1: ep0/t200, ep0/t250, ep0/t300, ep0/t350

Consequently, each runtime memory slot skips between sparse points while the
history of one episode is split across multiple independent slots.

Change

When preserve_episode_batch_slots is enabled, complete episodes are assigned
to approximately load-balanced persistent slots:

batch 0: ep0/t0,  ep1/t0,  ep2/t0,  ep3/t0
batch 1: ep0/t50, ep1/t50, ep2/t50, ep3/t50

The implementation also:

  • keeps all distributed ranks at the same number of batches;
  • cycles shorter slot streams only at complete episode boundaries;
  • handles ranks with fewer episodes than batch slots using a complete fallback
    stream rather than repeatedly copying one frame.

The option defaults to false at the sampler API boundary, preserving the
released behavior for existing users. It is enabled in the RoboTwin-MeM
training YAML, where predicted KEM state is slot-indexed.

Validation

  • 6 focused unit tests covering compatibility mode, slot identity, episode
    boundaries, sparse anchors, uneven distributed ranks, and empty slots.
  • 100 randomized configurations comparing compatibility mode sample-for-sample
    with the released sampler and checking persistent-mode causality/DDP lengths.
  • python3 -m pytest -q EventVLA/tests/test_sequence_sampler_slots.py
  • ruff check EventVLA/tests/test_sequence_sampler_slots.py
  • black --check EventVLA/tests/test_sequence_sampler_slots.py
  • git diff --check

Related to #7.

@JLUbaijh2323
JLUbaijh2323 marked this pull request as ready for review August 21, 2026 15:16
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