diff --git a/EventVLA/eventvla/dataloader/__init__.py b/EventVLA/eventvla/dataloader/__init__.py index 93194cd..2498314 100644 --- a/EventVLA/eventvla/dataloader/__init__.py +++ b/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, ) diff --git a/EventVLA/eventvla/dataloader/sequence_sampler.py b/EventVLA/eventvla/dataloader/sequence_sampler.py index 4525e0d..6e5d6d6 100644 --- a/EventVLA/eventvla/dataloader/sequence_sampler.py +++ b/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 + ] diff --git a/EventVLA/examples/RoboTwin-Mem/train_files/eventvla_robotwin_mem.yaml b/EventVLA/examples/RoboTwin-Mem/train_files/eventvla_robotwin_mem.yaml index bbad783..2d9b723 100644 --- a/EventVLA/examples/RoboTwin-Mem/train_files/eventvla_robotwin_mem.yaml +++ b/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 diff --git a/EventVLA/tests/test_sequence_sampler_slots.py b/EventVLA/tests/test_sequence_sampler_slots.py new file mode 100644 index 0000000..e147eca --- /dev/null +++ b/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))