""" Sequential Episode DataLoader implementation. This DataLoader initializes batches by selecting random episodes and then returning samples sequentially from each episode starting from the first step. When reset, it selects new random episodes and starts from the first step again. """ from collections.abc import Iterator import logging # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) import random from typing import Literal, Optional, List, Dict, Any, Callable, Tuple import jax import numpy as np import torch from torch.utils.data import DataLoader as TorchDataLoaderBase from torch.utils.data import Sampler import openpi_value.training.config as _config import openpi_value.models.model as _model class SequentialEpisodeSampler(Sampler): """ Sampler that samples batches where each element in the batch has its own independent episode sequence. Each element in the batch maintains its own episode state and can be reset independently. Args: dataset: The dataset to sample from batch_size: Number of elements (episodes) in each batch shuffle_episodes: Whether to shuffle the episodes when resetting seed: Random seed for reproducibility """ def __init__( self, dataset, batch_size: int, shuffle_episodes: bool = True, seed: int = 0 ): self.dataset = dataset._dataset # * self.dataset is correct # * CustomMultiLeRobotDataset self.batch_size = batch_size self.shuffle_episodes = shuffle_episodes self.seed = seed self.rng = random.Random(seed) self.episode_to_frames: Dict[int, List[int]] = {} # 1. Handle Multi/Concat Datasets if hasattr(self.dataset, '_datasets'): current_frame_offset = 0 global_ep_counter = 0 for sub_ds in self.dataset._datasets: # Extract frames for this sub-dataset (returns {local_ep_id: [local_frames]}) sub_mapping = self._extract_frames_from_single_ds(sub_ds) # Remap to global IDs and global Frame indices for local_ep_id, local_frames in sub_mapping.items(): # Create a new unique ID for the sampler to use unique_ep_id = global_ep_counter # Shift frame indices by the accumulated offset global_frames = [f + current_frame_offset for f in local_frames] if len(global_frames) > 0: self.episode_to_frames[unique_ep_id] = global_frames global_ep_counter += 1 # Increment offset for the next dataset current_frame_offset += len(sub_ds) # 2. Handle Single Dataset else: # Standard extraction sub_mapping = self._extract_frames_from_single_ds(self.dataset) # Just verify length > 0 self.episode_to_frames = {k: v for k, v in sub_mapping.items() if len(v) > 0} self.episode_indices = sorted(list(self.episode_to_frames.keys())) self.num_episodes = len(self.episode_indices) self.step_chunk_size = 1 # Check if we have enough episodes if self.num_episodes < 1: raise ValueError(f"Not enough episodes in dataset. Need at least 1, " f"but only found {self.num_episodes}") # Initialize state for each element in the batch # Each element has its own episode and step counter self.element_states = [] for i in range(batch_size): # Create a separate RNG for each element to ensure independent sampling element_rng = random.Random(seed + i) self.element_states.append({ 'current_episode': None, 'current_step': 0, 'rng': element_rng }) # Reset all elements to select initial episodes self.reset() logger.info(f"SequentialEpisodeSampler initialized with {self.num_episodes} episodes") logger.info(f"Batch size: {self.batch_size}, shuffle_episodes: {self.shuffle_episodes}") def _extract_frames_from_single_ds(self, dataset) -> Dict[int, List[int]]: """ Helper to extract {local_ep_id: [local_frame_indices]} from a single dataset. Does NOT handle global offsets. """ mapping = {} has_skipping = ( hasattr(dataset, 'index_map') and dataset.index_map is not None ) # Optimized path for LeRobot/Custom datasets if hasattr(dataset, 'episode_data_index') and hasattr(dataset, 'episodes'): # We must reconstruct the VIRTUAL indices. # The dataset's index_map is built sequentially based on the order of episodes. # So we maintain a running counter of virtual indices. virtual_idx_counter = 0 # dataset.episodes is the list of episode_indices used to init the dataset. # These correspond exactly to the order in episode_data_index['from'] / ['to'] num_episodes = len(dataset.episode_data_index['from']) for arr_idx in range(num_episodes): # 1. Get raw info start = dataset.episode_data_index['from'][arr_idx].item() end = dataset.episode_data_index['to'][arr_idx].item() raw_len = end - start # 2. Calculate valid length (matching Dataset logic EXACTLY) if has_skipping: # Get start skip ratio ratio_start = getattr(dataset, 'preceding_skipping_ratio', 0.0) skip_n_start = int(raw_len * ratio_start) # Get end skip ratio ratio_end = getattr(dataset, 'trailing_skipping_ratio', 0.0) skip_n_end = int(raw_len * ratio_end) # Replicate Safety Logic from Dataset to ensure counts align max_allowed_skip_end = raw_len - skip_n_start - 1 if max_allowed_skip_end < 0: # If start skip is already consuming the whole episode skip_n_start = raw_len - 1 skip_n_end = 0 else: # Cap end skip if it overlaps with start skip skip_n_end = min(skip_n_end, max_allowed_skip_end) valid_len = raw_len - skip_n_start - skip_n_end else: valid_len = raw_len # 3. Retrieve the actual Episode ID (the key for the dict) # dataset.episodes is the list [ep_idx_1, ep_idx_5, ...] # arr_idx is the index in that list. ep_idx = dataset.episodes[arr_idx] # 4. Generate VIRTUAL indices [current_counter, current_counter + len] # When passed to __getitem__, dataset will map these -> index_map -> Raw mapping[ep_idx] = list(range(virtual_idx_counter, virtual_idx_counter + valid_len)) # 5. Advance the virtual counter virtual_idx_counter += valid_len else: # Only use this if absolute necessary, it is slow # You might want to skip this or implement specific logic for your other dataset types # pass raise NotImplementedError("Dataset type not supported in multi-dataset extraction.") return mapping def reset(self, reset_episode_state: Optional[List[Tuple[int, int]]] = None, element_indices: Optional[List[int]] = None) -> None: """ Reset the sampler: select new random episodes for elements and start from the first step. Args: element_indices: If specified, only reset these particular elements. If None, reset all elements. """ # Calculate episode lengths if not already done if not hasattr(self, 'episode_lengths'): self.episode_lengths: Dict[int, int] = {ep: len(frames) for ep, frames in self.episode_to_frames.items()} # Determine which elements to reset elements_to_reset = element_indices if element_indices is not None else range(self.batch_size) for idx in elements_to_reset: if idx < 0 or idx >= self.batch_size: raise ValueError(f"Invalid element index: {idx}. Must be between 0 and {self.batch_size - 1}") element_state = self.element_states[idx] if reset_episode_state is not None: element_state['current_episode'] = reset_episode_state[idx][0] element_state['current_step'] = reset_episode_state[idx][1] else: # Select a random episode for this element if self.shuffle_episodes: # Select a random episode selected_ep = element_state['rng'].choice(self.episode_indices) else: # Select in order, cycling if necessary current_ep = element_state['current_episode'] if element_state['current_episode'] is not None else 0 current_idx = self.episode_indices.index(current_ep) if current_ep in self.episode_indices else 0 selected_ep = self.episode_indices[(current_idx + 1) % self.num_episodes] element_state['current_episode'] = selected_ep element_state['current_step'] = element_state['rng'].randint(0, self.episode_lengths[element_state['current_episode']] - 1) logger.debug(f"Element {idx} reset. Selected episode: {selected_ep}, starting at step 0") def __iter__(self): """ Iterate over the dataset, yielding batches where each element in the batch independently progresses through its own episode sequence. """ # Calculate episode lengths if not already done if not hasattr(self, 'episode_lengths'): self.episode_lengths: Dict[int, int] = {ep: len(frames) for ep, frames in self.episode_to_frames.items()} while True: batch_indices = [] for i in range(self.batch_size): element_state = self.element_states[i] # Check if we've reached the end of this element's current episode if element_state['current_step'] >= self.episode_lengths[element_state['current_episode']]: # This episode is exhausted, replace it with a new one new_ep_idx = self._get_new_episode(element_state['current_episode'], element_state['rng']) element_state['current_episode'] = new_ep_idx element_state['current_step'] = element_state['rng'].randint(0, self.episode_lengths[element_state['current_episode']] - 1) # Get the frame index for this element's current step frame_idx = self.episode_to_frames[element_state['current_episode']][element_state['current_step']] batch_indices.append(frame_idx) # Move to next step for this element element_state['current_step'] += self.step_chunk_size # We should always have exactly batch_size indices now assert len(batch_indices) == self.batch_size, \ f"Batch size mismatch: expected {self.batch_size}, got {len(batch_indices)}" yield batch_indices def _get_new_episode(self, exclude_ep_idx: int, rng: Optional[random.Random] = None) -> int: """Get a new episode index that is not the excluded one""" # Use the provided RNG or fall back to the main RNG rng = rng or self.rng # Create a list of available episodes excluding the exhausted one available_episodes = [ep for ep in self.episode_indices if ep != exclude_ep_idx] # If all episodes are excluded (shouldn't happen with proper batch_size), # just return the excluded one if not available_episodes: return exclude_ep_idx # Select a random episode from available ones return rng.choice(available_episodes) def __len__(self): return len(self.dataset) class SequentialEpisodeDataLoader: """ DataLoader that initializes batches by selecting episodes and then returning samples sequentially from each episode starting from the first step. Args: dataset: The dataset to load data from batch_size: Number of episodes to include in each batch shuffle_episodes: Whether to shuffle the episodes when resetting num_workers: Number of worker processes for data loading pin_memory: Whether to pin memory in DataLoader seed: Random seed for reproducibility framework: The framework to use ("jax" or "pytorch") sharding: JAX sharding configuration (only for JAX framework) """ def __init__( self, dataset, batch_size: int, shuffle_episodes: bool = True, num_workers: int = 0, pin_memory: bool = False, seed: int = 0, framework: Literal["jax", "pytorch"] = "pytorch", sharding: Optional[jax.sharding.Sharding] = None, ): self.dataset = dataset self.batch_size = batch_size self.shuffle_episodes = shuffle_episodes self.num_workers = num_workers self.pin_memory = pin_memory self.seed = seed self.framework = framework self.sharding = sharding # Create the sequential episode sampler self.sampler = SequentialEpisodeSampler( dataset=dataset, batch_size=batch_size, shuffle_episodes=shuffle_episodes, seed=seed ) # Create the underlying PyTorch DataLoader with our custom sampler self.torch_loader = TorchDataLoaderBase( dataset=dataset, batch_sampler=self.sampler, num_workers=num_workers, pin_memory=pin_memory, collate_fn=self._collate_fn, worker_init_fn=self._worker_init_fn, ) # Set up sharding for JAX if needed if self.sharding is None and framework == "jax": # Use data parallel sharding by default for JAX self.sharding = jax.sharding.NamedSharding( jax.sharding.Mesh(jax.devices(), ("B",)), jax.sharding.PartitionSpec("B"), ) def _collate_fn(self, items): """Collate function to handle dictionary items""" # Make sure to convert to numpy arrays before stacking since some elements # may be JAX arrays return jax.tree.map(lambda *xs: np.stack([np.asarray(x) for x in xs], axis=0), *items) def _worker_init_fn(self, worker_id: int) -> None: """Worker initialization function""" # Set random seed for reproducibility worker_seed = self.seed + worker_id random.seed(worker_seed) np.random.seed(worker_seed) # Configure JAX in worker processes if self.framework == "jax": import os os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" def reset(self, reset_episode_state: Optional[List[Tuple[int, int]]] = None, element_indices: Optional[List[int]] = None) -> None: """ Reset the data loader: select new random episodes for elements and start from the first step. Args: element_indices: If specified, only reset these particular elements. If None, reset all elements. """ self.sampler.reset(reset_episode_state, element_indices) if element_indices is None: logger.info("DataLoader reset: all elements have new episodes") else: logger.info(f"DataLoader reset: elements {element_indices} have new episodes") def __iter__(self): """Iterate over batches""" self._iterator = iter(self.torch_loader) return self def __next__(self): """Get the next batch""" if self._iterator is None: self._iterator = iter(self.torch_loader) batch = next(self._iterator) if self.framework == "jax" and self.sharding is not None: # Convert to sharded arrays for JAX return jax.tree.map(lambda x: jax.make_array_from_process_local_data(self.sharding, x), batch) else: # Return as torch tensors for PyTorch return jax.tree.map(torch.as_tensor, batch) def __len__(self): """Number of batches""" return len(self.sampler) class SequentialEpisodeDataLoaderImpl: """Implementation of the DataLoader interface using SequentialEpisodeDataLoader""" def __init__(self, data_config: _config.DataConfig, data_loader: SequentialEpisodeDataLoader): self._data_config = data_config self._data_loader = data_loader def data_config(self) -> _config.DataConfig: return self._data_config def reset(self) -> None: """Reset the data loader""" self._data_loader.reset() def __iter__(self): for batch in self._data_loader: yield _model.Observation.from_dict(batch), batch["actions"] def create_sequential_episode_data_loader( data_config: _config.DataConfig, model_config: _model.BaseModelConfig, action_horizon: int, batch_size: int, *, sharding: Optional[jax.sharding.Sharding] = None, skip_norm_stats: bool = False, shuffle_episodes: bool = True, num_workers: int = 0, seed: int = 0, framework: Literal["jax", "pytorch"] = "jax", config: Optional[Any] = None, ) -> SequentialEpisodeDataLoaderImpl: """ Create a sequential episode-based data loader for training. Args: data_config: The data configuration model_config: The model configuration action_horizon: The action horizon batch_size: Number of episodes per batch sharding: JAX sharding configuration skip_norm_stats: Whether to skip data normalization shuffle_episodes: Whether to shuffle episodes when resetting num_workers: Number of worker processes seed: Random seed framework: Framework to use ("jax" or "pytorch") config: Additional configuration """ # Import locally to avoid circular imports from openpi_value.training.data_loader import create_torch_dataset, transform_dataset dataset = create_torch_dataset( data_config=data_config, action_horizon=action_horizon, model_config=model_config, config=config ) # Apply transformations dataset = transform_dataset( dataset=dataset, data_config=data_config, skip_norm_stats=skip_norm_stats ) # Handle distributed training local_batch_size = batch_size if framework == "pytorch" and torch.distributed.is_initialized(): local_batch_size = batch_size // torch.distributed.get_world_size() elif framework == "jax": local_batch_size = batch_size // jax.process_count() logging.info(f"Creating SequentialEpisodeDataLoader with batch_size: {local_batch_size}") # Create the sequential episode-based data loader data_loader = SequentialEpisodeDataLoader( dataset=dataset, batch_size=local_batch_size, shuffle_episodes=shuffle_episodes, num_workers=num_workers, pin_memory=True if framework == "pytorch" else False, seed=seed, framework=framework, sharding=sharding, ) return SequentialEpisodeDataLoaderImpl(data_config, data_loader) def modify_create_data_loader_for_sequential_episodes(): """ Modify the original create_data_loader function to support sequential episode-based batching. This function should be called after importing the original create_data_loader. """ import openpi_value.training.data_loader as original_dataloader def modified_create_torch_data_loader( data_config: _config.DataConfig, model_config: _model.BaseModelConfig, action_horizon: int, batch_size: int, *, sharding: Optional[jax.sharding.Sharding] = None, skip_norm_stats: bool = False, shuffle: bool = False, num_batches: Optional[int] = None, num_workers: int = 0, seed: int = 0, framework: str = "jax", config: Optional[Any] = None, ): """Modified create_torch_data_loader that supports sequential episode-based batching""" # Use sequential episode-based data loader logging.info("Using sequential episode-based data loader") return create_sequential_episode_data_loader( data_config=data_config, model_config=model_config, action_horizon=action_horizon, batch_size=batch_size, sharding=sharding, skip_norm_stats=skip_norm_stats, shuffle_episodes=shuffle, num_workers=num_workers, seed=seed, framework=framework, config=config, ) # Replace the original function original_dataloader.create_torch_data_loader = modified_create_torch_data_loader logging.info("Modified create_torch_data_loader to support sequential episode-based batching") def create_dataloader_with_sequential_episode(config: _config.TrainConfig, framework = "pytorch", shuffle=True): """ Example function to create a data loader with sequential episode-based batching. Args: config: Training configuration """ # First, modify the original create_data_loader to support sequential episode-based batching modify_create_data_loader_for_sequential_episodes() # Then create the data loader as usual from openpi_value.training.data_loader import create_data_loader # Set the flag to use sequential episode-based batching data_loader = create_data_loader(config, framework=framework, shuffle=shuffle) return data_loader