Fix episode continuity across sampler batch slots

EventVLA/eventvla/dataloader/__init__.py

@@ -92,6 +92,9 @@ def build_dataloader(cfg, dataset_py="lerobot_datasets_oxe"):  # TODO now here o
                 sampling_interval=sampling_interval,
                 action_horizon=action_horizon,
                 balance_dataset_step_counts=balance_task_step_counts,
+                preserve_episode_batch_slots=_cfg_bool(
+                    vla_dataset_cfg.get("preserve_episode_batch_slots", False)
+                ),
                 rank=rank,
                 num_replicas=world_size,
             )

EventVLA/eventvla/dataloader/sequence_sampler.py

@@ -42,6 +42,7 @@ def __init__(
         sampling_interval: int = 1,
         action_horizon: int = 1,
         balance_dataset_step_counts: bool = False,
+        preserve_episode_batch_slots: bool = False,
         rank: int | None = None,
         num_replicas: int | None = None,
     ) -> None:
@@ -52,6 +53,7 @@ def __init__(
         self.sampling_interval = max(1, int(sampling_interval))
         self.action_horizon = max(1, int(action_horizon))
         self.balance_dataset_step_counts = bool(balance_dataset_step_counts)
+        self.preserve_episode_batch_slots = bool(preserve_episode_batch_slots)
 
         self.epoch = 0
         if num_replicas is None:
@@ -366,6 +368,14 @@ def _build_flat_step_stream(
     def _compute_target_num_batches(self, trajectories: List[Tuple[int, object, int]]) -> int:
         # Deterministic across all ranks (no collectives): all ranks have the same
         # trajectory list + same shuffling seed, and rank assignment is via slicing.
+        if self.preserve_episode_batch_slots:
+            batches_per_rank = []
+            for rank in range(self.num_replicas):
+                rank_pool = trajectories[rank :: self.num_replicas]
+                _, slot_step_counts = self._assign_trajectories_to_slots(rank_pool)
+                batches_per_rank.append(max(slot_step_counts, default=0))
+            return int(max(batches_per_rank, default=0))
+
         steps_per_rank: List[int] = []
         for r in range(self.num_replicas):
             rank_pool = trajectories[r :: self.num_replicas]
@@ -382,6 +392,39 @@ def _compute_target_num_batches(self, trajectories: List[Tuple[int, object, int]
         # Use max + local cycling to keep all ranks at identical batch count.
         return int(np.ceil(max(steps_per_rank) / self.batch_size))
 
+    def _assign_trajectories_to_slots(
+        self,
+        trajectories: List[Tuple[int, object, int]],
+    ) -> Tuple[List[List[Tuple[int, object, int]]], List[int]]:
+        """Assign complete episodes to persistent, approximately balanced slots."""
+        slot_trajectories: List[List[Tuple[int, object, int]]] = [
+            [] for _ in range(self.batch_size)
+        ]
+        slot_step_counts = [0 for _ in range(self.batch_size)]
+
+        for trajectory in trajectories:
+            dataset_index, trajectory_id, trajectory_length = trajectory
+            sampled_steps = self._count_sampled_steps(
+                dataset_index=int(dataset_index),
+                trajectory_id=trajectory_id,
+                trajectory_length=int(trajectory_length),
+            )
+            if sampled_steps <= 0:
+                continue
+
+            slot_index = min(range(self.batch_size), key=slot_step_counts.__getitem__)
+            slot_trajectories[slot_index].append(trajectory)
+            slot_step_counts[slot_index] += int(sampled_steps)
+
+        return slot_trajectories, slot_step_counts
+
+    def _build_slot_streams(
+        self,
+        trajectories: List[Tuple[int, object, int]],
+    ) -> List[List[EpisodeSampleIndex]]:
+        slot_trajectories, _ = self._assign_trajectories_to_slots(trajectories)
+        return [self._build_flat_step_stream(slot_pool) for slot_pool in slot_trajectories]
+
     def __len__(self) -> int:
         trajectories = self._build_epoch_trajectory_pool()
         return self._compute_target_num_batches(trajectories)
@@ -421,6 +464,10 @@ def _build_fallback_stream_sample(
     def __iter__(self) -> Iterator[List[EpisodeSampleIndex]]:
         all_trajectories = self._build_epoch_trajectory_pool()
         trajectories = all_trajectories[self.rank :: self.num_replicas]
+        if self.preserve_episode_batch_slots:
+            yield from self._iter_persistent_slot_batches(all_trajectories, trajectories)
+            return
+
         local_stream = self._build_flat_step_stream(trajectories)
 
         target_num_batches = self._compute_target_num_batches(all_trajectories)
@@ -444,3 +491,29 @@ def __iter__(self) -> Iterator[List[EpisodeSampleIndex]]:
 
         for i in range(0, total_needed, self.batch_size):
             yield local_stream[i : i + self.batch_size]
+
+    def _iter_persistent_slot_batches(
+        self,
+        all_trajectories: List[Tuple[int, object, int]],
+        trajectories: List[Tuple[int, object, int]],
+    ) -> Iterator[List[EpisodeSampleIndex]]:
+        slot_streams = self._build_slot_streams(trajectories)
+        target_num_batches = self._compute_target_num_batches(all_trajectories)
+        if target_num_batches <= 0:
+            return
+
+        fallback_stream = next((stream for stream in slot_streams if stream), None)
+        if fallback_stream is None:
+            fallback_stream = next(
+                (stream for stream in self._build_slot_streams(all_trajectories) if stream),
+                None,
+            )
+        if fallback_stream is None:
+            return
+
+        resolved_streams = [stream if stream else fallback_stream for stream in slot_streams]
+        for batch_index in range(target_num_batches):
+            yield [
+                stream[batch_index % len(stream)]
+                for stream in resolved_streams
+            ]

EventVLA/examples/RoboTwin-Mem/train_files/eventvla_robotwin_mem.yaml

@@ -102,6 +102,9 @@ datasets:
     action_mode: abs
     sequential_step_sampling: False
     use_sequential_episode_sampler: true
+    # Runtime predicted keyframe memory is indexed by batch slot, so keep each
+    # episode in the same slot across consecutive batches.
+    preserve_episode_batch_slots: true
     shuffle_trajectories: true
     balance_task_step_counts: true
     num_workers: 4

EventVLA/tests/test_sequence_sampler_slots.py

@@ -0,0 +1,124 @@
+import importlib.util
+from itertools import pairwise
+from pathlib import Path
+from types import SimpleNamespace
+
+SAMPLER_PATH = Path(__file__).parents[1] / "eventvla" / "dataloader" / "sequence_sampler.py"
+SPEC = importlib.util.spec_from_file_location("eventvla_sequence_sampler_test", SAMPLER_PATH)
+MODULE = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+SPEC.loader.exec_module(MODULE)
+SequentialEpisodeBatchSampler = MODULE.SequentialEpisodeBatchSampler
+
+
+class Dataset:
+    def __init__(self, lengths):
+        self.trajectory_ids = [f"episode-{index}" for index in range(len(lengths))]
+        self.trajectory_lengths = lengths
+
+
+def make_sampler(
+    lengths,
+    *,
+    batch_size=4,
+    preserve_episode_batch_slots=True,
+    rank=0,
+    num_replicas=1,
+    sampling_interval=1,
+    action_horizon=1,
+):
+    dataset = SimpleNamespace(datasets=[Dataset(lengths)])
+    return SequentialEpisodeBatchSampler(
+        dataset=dataset,
+        batch_size=batch_size,
+        sampling_interval=sampling_interval,
+        action_horizon=action_horizon,
+        preserve_episode_batch_slots=preserve_episode_batch_slots,
+        rank=rank,
+        num_replicas=num_replicas,
+    )
+
+
+def test_default_mode_preserves_official_flat_batch_layout():
+    sampler = make_sampler(
+        [8, 8, 8, 8],
+        preserve_episode_batch_slots=False,
+    )
+
+    first_batch = next(iter(sampler))
+
+    assert [sample[1] for sample in first_batch] == ["episode-0"] * 4
+    assert [sample[2] for sample in first_batch] == [0, 1, 2, 3]
+
+
+def test_persistent_mode_keeps_episode_identity_per_slot():
+    sampler = make_sampler([8, 8, 8, 8])
+    batches = list(sampler)
+
+    assert [sample[1] for sample in batches[0]] == [
+        "episode-0",
+        "episode-1",
+        "episode-2",
+        "episode-3",
+    ]
+    assert [sample[1] for sample in batches[1]] == [sample[1] for sample in batches[0]]
+    assert [sample[2] for sample in batches[0]] == [0, 0, 0, 0]
+    assert [sample[2] for sample in batches[1]] == [1, 1, 1, 1]
+
+
+def test_persistent_mode_preserves_episode_boundaries_within_each_slot():
+    sampler = make_sampler([3, 5, 4, 6, 2, 7], batch_size=2)
+    batches = list(sampler)
+
+    for slot_index in range(2):
+        slot_samples = [batch[slot_index] for batch in batches]
+        for previous, current in pairwise(slot_samples):
+            same_episode = previous[1] == current[1]
+            if same_episode and not previous[4]:
+                assert current[2] > previous[2]
+                assert not current[3]
+            elif not same_episode:
+                assert previous[4]
+                assert current[3]
+
+
+def test_persistent_mode_keeps_all_ranks_at_the_same_length():
+    samplers = [
+        make_sampler(
+            [2, 9, 3, 8, 4, 7, 5],
+            batch_size=2,
+            rank=rank,
+            num_replicas=3,
+        )
+        for rank in range(3)
+    ]
+
+    lengths = [len(list(sampler)) for sampler in samplers]
+
+    assert lengths[0] == lengths[1] == lengths[2]
+    assert lengths == [len(sampler) for sampler in samplers]
+
+
+def test_persistent_mode_handles_fewer_rank_episodes_than_batch_slots():
+    sampler = make_sampler([5], batch_size=4)
+    batches = list(sampler)
+
+    assert len(batches) == len(sampler) == 5
+    assert all(len(batch) == 4 for batch in batches)
+    for slot_index in range(4):
+        assert [batch[slot_index][2] for batch in batches] == [0, 1, 2, 3, 4]
+
+
+def test_persistent_mode_retains_sparse_anchor_progression():
+    sampler = make_sampler(
+        [260, 260],
+        batch_size=2,
+        sampling_interval=50,
+        action_horizon=50,
+    )
+    batches = list(sampler)
+
+    for slot_index in range(2):
+        steps = [batch[slot_index][2] for batch in batches]
+        assert steps[0] == 0
+        assert all(current > previous for previous, current in pairwise(steps))