from collections.abc import Iterator, Sequence import json import logging import multiprocessing import os from pathlib import Path import typing from typing import Literal, Protocol, SupportsIndex, TypeVar import jax import jax.numpy as jnp import lerobot import numpy as np import torch import random try: # LeRobot v3 (0.5.x): consolidated parquet/video files. import lerobot.datasets.lerobot_dataset as lerobot_dataset_v3 except ImportError: lerobot_dataset_v3 = None try: # Legacy LeRobot v2.1 path retained for the original CFGRL configs. import lerobot.common.datasets.lerobot_dataset as lerobot_dataset_v2 except ImportError: lerobot_dataset_v2 = None import openpi_value.models.model as _model import openpi_value.training.config as _config import openpi_value.transforms as _transforms from inspect import signature T_co = TypeVar("T_co", covariant=True) def _local_lerobot_version(repo_id: str) -> str | None: """Return the codebase version for a local LeRobot dataset.""" if not os.path.isdir(repo_id): return None info_path = Path(repo_id) / "meta" / "info.json" if not info_path.is_file(): return None with info_path.open() as f: return json.load(f).get("codebase_version") def _is_local_v3_dataset(repo_id: str) -> bool: version = _local_lerobot_version(repo_id) return version is not None and version.startswith("v3") def _get_local_v3_episodes(repo_id: str, split: str) -> list[int]: """Return deterministic episode splits without assuming v2 metadata objects.""" with (Path(repo_id) / "meta" / "info.json").open() as f: total_episodes = int(json.load(f)["total_episodes"]) split_index = int(0.9 * total_episodes) if split == "all": return list(range(total_episodes)) if split == "train_tasks": return list(range(split_index)) if split == "val_tasks": return list(range(split_index, total_episodes)) raise ValueError(f"Unsupported split for LeRobot v3 dataset: {split}") class Dataset(Protocol[T_co]): """Interface for a dataset with random access.""" def __getitem__(self, index: SupportsIndex) -> T_co: raise NotImplementedError("Subclasses of Dataset should implement __getitem__.") def __len__(self) -> int: raise NotImplementedError("Subclasses of Dataset should implement __len__.") class IterableDataset(Protocol[T_co]): """Interface for an iterable dataset.""" def __iter__(self) -> Iterator[T_co]: raise NotImplementedError("Subclasses of IterableDataset should implement __iter__.") def __len__(self) -> int: raise NotImplementedError("Subclasses of Dataset should implement __len__.") class DataLoader(Protocol[T_co]): """Interface for a data loader.""" def data_config(self) -> _config.DataConfig: """Get the data config for this data loader.""" raise NotImplementedError("Subclasses of DataLoader should implement data_config.") def __iter__(self) -> Iterator[T_co]: raise NotImplementedError("Subclasses of DataLoader should implement __iter__.") class TransformedDataset(Dataset[T_co]): def __init__(self, dataset: Dataset, transforms: Sequence[_transforms.DataTransformFn]): self._dataset = dataset self._transform = _transforms.compose(transforms) def __getitem__(self, index: SupportsIndex) -> T_co: return self._transform(self._dataset[index]) def __len__(self) -> int: return len(self._dataset) class IterableTransformedDataset(IterableDataset[T_co]): def __init__( self, dataset: IterableDataset, transforms: Sequence[_transforms.DataTransformFn], *, is_batched: bool = False, ): self._dataset = dataset self._transform = _transforms.compose(transforms) self._is_batched = is_batched def __iter__(self): for sample in self._dataset: if self._is_batched: # Transforms are designed to be applied to individual samples. So we need to split the batch into # individual samples and apply the transform to each sample individually. batch_size = next(v.shape[0] for v in sample.values()) # Split batch into individual samples using tree_map individual_samples = [jax.tree.map(lambda x: x[i], sample) for i in range(batch_size)] # noqa: B023 # Transform each sample transformed = [self._transform(s) for s in individual_samples] # Recombine batch with tree_map yield jax.tree.map(lambda *x: np.stack(x, axis=0), *transformed) else: yield self._transform(sample) def __len__(self) -> int: return len(self._dataset) class FakeDataset(Dataset): def __init__(self, model_config: _model.BaseModelConfig, num_samples: int): self._num_samples = num_samples self._observation_spec, self._action_spec = model_config.inputs_spec() def __getitem__(self, index: SupportsIndex) -> dict: rng = jax.random.key(index.__index__()) def make_from_spec(spec: jax.ShapeDtypeStruct): nonlocal rng rng, data_rng = jax.random.split(rng) # Remove the batch dimension. shape = spec.shape[1:] if spec.dtype == jnp.float32: return jax.random.uniform(data_rng, shape=shape, minval=-1.0, maxval=1.0) if spec.dtype == jnp.int32: return jax.random.randint(data_rng, shape=shape, minval=0, maxval=2048) return jnp.zeros(shape=shape, dtype=spec.dtype) observation = jax.tree.map(make_from_spec, self._observation_spec) action = jax.tree.map(make_from_spec, self._action_spec) return { **observation.to_dict(), "actions": action, } def __len__(self) -> int: return self._num_samples # * Used for norm states. NOT used for training def create_torch_dataset_naive( data_config: _config.DataConfig, action_horizon: int, model_config: _model.BaseModelConfig ) -> Dataset: """Create a dataset for training.""" if lerobot_dataset_v2 is None: raise RuntimeError("The legacy LeRobot v2.1 API is not installed.") repo_id = data_config.repo_id if repo_id is None: raise ValueError("Repo ID is not set. Cannot create dataset.") if repo_id == "fake": return FakeDataset(model_config, num_samples=1024) if isinstance(repo_id, str) and os.path.exists(repo_id): if 'data' not in os.listdir(repo_id) and 'videos' not in os.listdir(repo_id): repo_id = [os.path.join(repo_id, d) for d in os.listdir(repo_id) if os.path.isdir(os.path.join(repo_id, d))] if not isinstance(repo_id, list): repo_id = [repo_id] repo_id = [ rp for rp in repo_id if 'Put_The_Items_Into_The_Storage_Box_20250929_002_007' not in rp ] # * Remove this problematic dataset if len(repo_id) > 1: dataset_meta = lerobot_dataset_v2.LeRobotDatasetMetadata(repo_id[0]) # * Just use the first repo to get fps dataset = lerobot_dataset_v2.MultiLeRobotDataset( repo_id, delta_timestamps={ key: [t / dataset_meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys }, # * tolerance tolerances_s = dict.fromkeys(repo_id, 0.1) ) else: repo_id = repo_id[0] dataset_meta = lerobot_dataset_v2.LeRobotDatasetMetadata(repo_id) dataset = lerobot_dataset_v2.LeRobotDataset( repo_id, delta_timestamps={ key: [t / dataset_meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys }, # * tolerance tolerance_s = 0.1, ) return dataset def create_torch_dataset( data_config: _config.DataConfig, action_horizon: int, model_config: _model.BaseModelConfig, config=None, ) -> Dataset: """Create a dataset for training (supports multiple repo_ids).""" split=config.split repo_ids = data_config.repo_id if not repo_ids: raise ValueError("Repo ID(s) not set in data_config. Cannot create dataset.") # * Support folder containing multiple lerobot datasetss if isinstance(repo_ids, str) and os.path.exists(repo_ids): if 'data' not in os.listdir(repo_ids) and 'videos' not in os.listdir(repo_ids): repo_ids = [os.path.join(repo_ids, d) for d in os.listdir(repo_ids) if os.path.isdir(os.path.join(repo_ids, d))] if not isinstance(repo_ids, list): repo_ids = [repo_ids] if repo_ids == ["fake"]: return FakeDataset(model_config, num_samples=1024) repo_ids = [ rp for rp in repo_ids if 'Put_The_Items_Into_The_Storage_Box_20250929_002_007' not in rp ] # * Remove this problematic dataset repo_to_episodes: dict[str, list[int]] = {} valid_repo_ids: list[str] = [] for rid in repo_ids: if os.path.isabs(rid): info_path = os.path.join(rid, "meta", "info.json") if not os.path.isfile(info_path): logging.warning("Skipping local dataset without metadata: %s", info_path) continue episodes = get_episodes(rid, split) if len(episodes) == 0: logging.warning("Skipping dataset with no episodes for split '%s': %s", split, rid) continue repo_to_episodes[rid] = episodes valid_repo_ids.append(rid) repo_ids = valid_repo_ids if len(repo_ids) == 0: raise ValueError("No valid datasets found after filtering unavailable local dataset paths.") # The new assemble-battery dataset is LeRobot v3. Use the maintained v3 # reader directly; the CFGRL custom reader is tied to the per-episode v2.1 # file layout and cannot decode consolidated v3 parquet/video files. v3_repo_ids = [rid for rid in repo_ids if _is_local_v3_dataset(rid)] if v3_repo_ids: if lerobot_dataset_v3 is None: raise RuntimeError("LeRobot v3 dataset detected, but lerobot>=0.5 is not installed.") if len(repo_ids) != 1: raise NotImplementedError("Mixing LeRobot v3 and legacy datasets is not supported in this loader.") rid = v3_repo_ids[0] with (Path(rid) / "meta" / "info.json").open() as f: fps = int(json.load(f)["fps"]) delta_timestamps = { key: [t / fps for t in range(action_horizon)] for key in data_config.action_sequence_keys } dataset = lerobot_dataset_v3.LeRobotDataset( repo_id=f"local/{Path(rid).name}", root=rid, episodes=repo_to_episodes[rid], delta_timestamps=delta_timestamps, tolerance_s=0.1, video_backend="pyav", ) if data_config.prompt_from_task: dataset = TransformedDataset(dataset, [_transforms.PromptFromLeRobotTask(tasks=None)]) return dataset if lerobot_dataset_v2 is None: raise RuntimeError("Legacy dataset detected, but the LeRobot v2.1 API is not installed.") # Import the old CFGRL reader only for legacy datasets. Importing it with # LeRobot 0.5 would fail before the v3 branch above can run. from openpi_value.training.custom_lerobot_dataset import CustomLeRobotDataset, CustomMultiLeRobotDataset dataset_kwargs = signature(CustomLeRobotDataset.__init__).parameters skip_args = {"self", "repo_id", "episodes", "image_transforms", "delta_timestamps", "tolerance_s", "download_videos", "video_backend"} valid_kwargs = [k for k in dataset_kwargs if k not in skip_args] data_kwargs = {k: getattr(config, k) for k in valid_kwargs if hasattr(config, k)} if len(repo_ids) > 1: all_delta_timestamps = [] for rid in repo_ids: meta = lerobot_dataset_v2.LeRobotDatasetMetadata(rid) delta_timestamps = { key: [t / meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys } all_delta_timestamps.append(delta_timestamps) dataset = CustomMultiLeRobotDataset( repo_ids=repo_ids, episodes=repo_to_episodes, # dict[repo_id] -> list[int] delta_timestamps=all_delta_timestamps, **data_kwargs ) else: rid = repo_ids[0] meta = lerobot_dataset_v2.LeRobotDatasetMetadata(rid) delta_timestamps = { key: [t / meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys } dataset = CustomLeRobotDataset( repo_id=rid, episodes=repo_to_episodes[rid], delta_timestamps=delta_timestamps, **data_kwargs ) if data_config.prompt_from_task: if len(repo_ids) == 1: dataset_meta = lerobot_dataset_v2.LeRobotDatasetMetadata(repo_ids[0]) dataset = TransformedDataset( dataset, [_transforms.PromptFromLeRobotTask(dataset_meta.tasks)] ) else: dataset = TransformedDataset( dataset, [_transforms.PromptFromLeRobotTask(tasks=None)] ) return dataset # * split tasks could be troublesome, # * suppose you have a dataset with two tasks: task1 and task2, when split task, # * it's possible to train on one task only, and val on another task only, which is not desired. # * One simple update is to train/val on all tasks without heldout tasks. def get_episodes(repo_id, split): assert split in ['all', 'train_tasks', 'val_tasks'], \ f"Invalid split option: {split}. Choose from 'all', 'train_tasks', 'val_tasks'. heldout_tasks is deprecated." if _is_local_v3_dataset(repo_id): return _get_local_v3_episodes(repo_id, split) if lerobot_dataset_v2 is None: raise RuntimeError("Legacy dataset detected, but the LeRobot v2.1 API is not installed.") if os.path.isabs(repo_id): info_path = os.path.join(repo_id, "meta", "info.json") if not os.path.isfile(info_path): raise FileNotFoundError(f"Local dataset metadata not found: {info_path}") dataset_meta = lerobot_dataset_v2.LeRobotDatasetMetadata(repo_id) episodes_meta = dataset_meta.episodes # Step 1: Group episodes by task (assuming episodes have a "tasks" field) task_to_episodes = {} for episode_index, episode_data in episodes_meta.items(): tasks = episode_data['tasks'] # This should be a list of tasks for this episode for task in tasks: if task not in task_to_episodes: task_to_episodes[task] = [] task_to_episodes[task].append(episode_index) # Step 2: Split the episodes total_episodes = set() train_episodes = set() val_episodes = set() all_tasks = list(task_to_episodes.keys()) train_val_ratio = 0.9 # train_val tasks could be the same, but different episodes for task in all_tasks: task_episodes = task_to_episodes[task] if len(task_episodes) <= 1: train_val_ratio = 1. else: train_val_ratio = 0.9 split_index = int(train_val_ratio * len(task_episodes)) # 90% for train, 20% for val total_episodes.update(task_episodes) train_episodes.update(task_episodes[:split_index]) val_episodes.update(task_episodes[split_index:]) if split == "all": out_eps = list(total_episodes) elif split == "train_tasks": out_eps = list(train_episodes - val_episodes) elif split == "val_tasks": out_eps = list(val_episodes - train_episodes) else: raise ValueError(f"Invalid split option: {split}. Choose from 'train', 'val', or 'heldout'.") # * Order the episodes out_eps = sorted(out_eps) return out_eps def transform_dataset(dataset: Dataset, data_config: _config.DataConfig, *, skip_norm_stats: bool = False) -> Dataset: """Transform the dataset by applying the data transforms.""" norm_stats = None if data_config.repo_id != "fake" and not skip_norm_stats: if data_config.norm_stats is None: raise ValueError( "Normalization stats not found. " "Make sure to run `scripts/compute_norm_stats.py --config-name=`." ) norm_stats = data_config.norm_stats return TransformedDataset( dataset, [ *data_config.repack_transforms.inputs, *data_config.data_transforms.inputs, _transforms.Normalize(norm_stats, use_quantiles=data_config.use_quantile_norm), *data_config.model_transforms.inputs, ], ) def transform_iterable_dataset( dataset: IterableDataset, data_config: _config.DataConfig, *, skip_norm_stats: bool = False, is_batched: bool = False, ) -> IterableDataset: """Transform the dataset by applying the data transforms.""" norm_stats = {} if data_config.repo_id != "fake" and not skip_norm_stats: if data_config.norm_stats is None: raise ValueError( "Normalization stats not found. " "Make sure to run `scripts/compute_norm_stats.py --config-name=`." ) norm_stats = data_config.norm_stats return IterableTransformedDataset( dataset, [ *data_config.repack_transforms.inputs, *data_config.data_transforms.inputs, _transforms.Normalize(norm_stats, use_quantiles=data_config.use_quantile_norm), *data_config.model_transforms.inputs, ], is_batched=is_batched, ) def create_data_loader( config: _config.TrainConfig, *, sharding: jax.sharding.Sharding | None = None, shuffle: bool = False, num_batches: int | None = None, skip_norm_stats: bool = False, framework: Literal["jax", "pytorch"] = "jax", ) -> DataLoader[tuple[_model.Observation, _model.Actions]]: """Create a data loader for training. Args: config: The training configuration. sharding: The sharding to use for the data loader (JAX only). shuffle: Whether to shuffle the data. num_batches: Determines the number of batches to return. skip_norm_stats: Whether to skip data normalization. framework: The framework to use ("jax" or "pytorch"). """ data_config = config.data.create(config.assets_dirs, config.model) logging.info(f"data_config: {data_config}") return create_torch_data_loader( data_config, model_config=config.model, action_horizon=config.model.action_horizon, batch_size=config.batch_size, sharding=sharding, shuffle=shuffle, num_batches=num_batches, num_workers=config.num_workers, seed=config.seed, skip_norm_stats=skip_norm_stats, framework=framework, config=config, ) def create_torch_data_loader( data_config: _config.DataConfig, model_config: _model.BaseModelConfig, action_horizon: int, batch_size: int, *, sharding: jax.sharding.Sharding | None = None, skip_norm_stats: bool = False, shuffle: bool = False, num_batches: int | None = None, num_workers: int = 0, seed: int = 0, framework: str = "jax", config: str = None, ) -> DataLoader[tuple[_model.Observation, _model.Actions]]: """Create a data loader for training. Args: data_config: The data configuration. action_horizon: The action horizon. batch_size: The batch size. sharding: The sharding to use for the data loader. If None, the data loader will use a single device sharding. skip_norm_stats: Whether to skip data normalization. shuffle: Whether to shuffle the data. num_batches: Determines the number of batches to return. If the number exceeds the number of batches in the dataset, the data loader will loop over the dataset. If not provided, will iterate over the dataset indefinitely. num_workers: The number of worker processes to use. If zero, the data loader will execute in the main process. seed: The seed to use for shuffling the data. """ dataset = create_torch_dataset(data_config, action_horizon, model_config, config=config) dataset = transform_dataset(dataset, data_config, skip_norm_stats=skip_norm_stats) # Use TorchDataLoader for both frameworks # For PyTorch DDP, create DistributedSampler and divide batch size by world size # For JAX, divide by process count sampler = None if framework == "pytorch": if torch.distributed.is_initialized(): sampler = torch.utils.data.distributed.DistributedSampler( dataset, num_replicas=torch.distributed.get_world_size(), rank=torch.distributed.get_rank(), shuffle=shuffle, drop_last=config.drop_last, ) local_batch_size = batch_size // torch.distributed.get_world_size() else: local_batch_size = batch_size else: local_batch_size = batch_size // jax.process_count() logging.info(f"local_batch_size: {local_batch_size}") data_loader = TorchDataLoader( dataset, local_batch_size=local_batch_size, sharding=None if framework == "pytorch" else sharding, shuffle=(sampler is None and shuffle), # Don't shuffle if using sampler sampler=sampler, num_batches=num_batches, num_workers=num_workers, seed=seed, framework=framework, drop_last=config.drop_last, ) return DataLoaderImpl(data_config, data_loader) class TorchDataLoader: """Torch data loader implementation.""" def __init__( self, dataset, local_batch_size: int, *, sharding: jax.sharding.Sharding | None = None, shuffle: bool = False, sampler: torch.utils.data.Sampler | None = None, num_batches: int | None = None, num_workers: int = 0, seed: int = 0, framework: str = "jax", drop_last: bool = True, ): """Create a PyTorch data loader. Args: dataset: The dataset to load. local_batch_size: The local batch size for each process. sharding: The sharding to use for the data loader. shuffle: Whether to shuffle the data. num_batches: If provided, determines the number of returned batches. If the number is larger than the number of batches in the dataset, the data loader will loop over the dataset. If not provided, will iterate over the dataset indefinitely. num_workers: The number of worker processes to use. If zero, the data loader will execute in the main process. seed: The seed to use for shuffling the data. """ if jax.process_count() > 1: raise NotImplementedError("Data loading with multiple processes is not supported.") if len(dataset) < local_batch_size: raise ValueError(f"Local batch size ({local_batch_size}) is larger than the dataset size ({len(dataset)}).") # Store sharding - None for PyTorch, JAX sharding for JAX self._sharding = sharding if sharding is None and framework == "jax": # Use data parallel sharding by default for JAX only. self._sharding = jax.sharding.NamedSharding( jax.sharding.Mesh(jax.devices(), ("B",)), jax.sharding.PartitionSpec("B"), ) self._num_batches = num_batches mp_context = None if num_workers > 0: mp_context = multiprocessing.get_context("spawn") generator = torch.Generator() generator.manual_seed(seed) self._data_loader = torch.utils.data.DataLoader( typing.cast(torch.utils.data.Dataset, dataset), batch_size=local_batch_size, shuffle=(sampler is None and shuffle), # Don't shuffle if using sampler sampler=sampler, num_workers=num_workers, multiprocessing_context=mp_context, persistent_workers=num_workers > 0, collate_fn=_collate_fn, worker_init_fn=_worker_init_fn, drop_last=drop_last, generator=generator, ) @property def torch_loader(self) -> torch.utils.data.DataLoader: return self._data_loader def __iter__(self): num_items = 0 while True: data_iter = iter(self._data_loader) while True: if self._num_batches is not None and num_items >= self._num_batches: return try: batch = next(data_iter) except StopIteration: break # We've exhausted the dataset. Create a new iterator and start over. num_items += 1 # For JAX, convert to sharded arrays; for PyTorch, return torch tensors if self._sharding is not None: yield jax.tree.map(lambda x: jax.make_array_from_process_local_data(self._sharding, x), batch) else: yield jax.tree.map(torch.as_tensor, batch) def _collate_fn(items): """Collate the batch elements into batched numpy arrays.""" # Make sure to convert to numpy arrays before stacking since some of the incoming 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(worker_id: int) -> None: """Tell JAX inside the worker process not to preallocate the GPU memory.""" # NOTE: This is called after jax is imported inside the worker process. This # means that this approach will not work for selecting the backend. os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" class DataLoaderImpl(DataLoader): def __init__(self, data_config: _config.DataConfig, data_loader: TorchDataLoader): self._data_config = data_config self._data_loader = data_loader def data_config(self) -> _config.DataConfig: return self._data_config def __iter__(self): for batch in self._data_loader: yield _model.Observation.from_dict(batch), batch["actions"]