dataloader

dataloader

get_dataloaders

get_dataloaders(config)

Create train and validation dataloaders with normalizer and optional tokenizer.

Parameters:

Name	Type	Description	Default
config	DictConfig	Main configuration object instantiated by Hydra	required

Returns:

Type	Description
DataLoader	Tuple of (train_loader, val_loader, normalizer, tokenizer, gripper_class_weights).
DataLoader | None	val_loader is None when val_ratio is 0.

The type hint for config indicates DictConfig, but at runtime hydra instantiates a MainConfig with

all target fields resolved into python objects.

Source code in src/versatil/data/dataloader.py

def get_dataloaders(
    config: DictConfig,
) -> tuple[
    data.DataLoader,
    data.DataLoader | None,
    LinearNormalizer,
    Tokenizer | None,
    float | None,
]:
    """Create train and validation dataloaders with normalizer and optional tokenizer.

    Args:
        config: Main configuration object instantiated by Hydra

    Returns:
        Tuple of (train_loader, val_loader, normalizer, tokenizer, gripper_class_weights).
        val_loader is None when val_ratio is 0.

    Note: The type hint for `config` indicates `DictConfig`, but at runtime hydra instantiates a `MainConfig` with
      all target fields resolved into python objects.
    """
    schema: DatasetSchema = config.task.dataset_schema
    action_space: ActionSpace = config.task.action_space
    observation_space: ObservationSpace = config.task.observation_space
    dataloader_config: DataLoaderConfig = config.task.dataloader
    tokenization_config: TokenizationConfig = dataloader_config.tokenization

    validate_dataloader_config(dataloader_config)
    validate_tokenizer_config(tokenization_config)

    logging.info(f"Using dataset schema: {schema.__class__.__name__}")
    _ensure_zarr_exists(
        schema=schema,
        recreate_on_missing_keys=dataloader_config.recreate_zarr_on_missing_keys,
    )
    skip_validation = dataloader_config.val_ratio == 0

    train_dataset = EpisodicDataset(
        zarr_path=schema.zarr_path,
        pred_horizon=config.task.prediction_horizon,
        obs_horizon=config.task.observation_horizon,
        dataloader_config=dataloader_config,
        train=True,
        seed=config.experiment.seed,
        action_space=action_space,
        observation_space=observation_space,
    )

    normalizer, tokenizer = train_dataset.get_normalizer_and_tokenizer(
        winsorize_depth=dataloader_config.winsorize_depth,
        depth_winsorize_quantiles=dataloader_config.depth_winsorize_quantiles,
        winsorize_kinematics=dataloader_config.winsorize_kinematics,
        kinematics_winsorize_quantiles=dataloader_config.kinematics_winsorize_quantiles,
        tokenization_config=tokenization_config,
        clamp_kinematics_range=dataloader_config.clamp_kinematics_range,
        min_kinematics_std=dataloader_config.min_kinematics_std,
        min_kinematics_range=dataloader_config.min_kinematics_range,
        action_sample_size=dataloader_config.action_sample_size,
        device=torch.device("cpu"),  # Keep on CPU for DataLoader workers
    )
    train_dataset.set_normalizer(normalizer)
    train_dataset.set_tokenizer(tokenizer)

    num_workers = config.task.dataloader.num_workers
    use_multiprocessing = num_workers > 0
    train_loader = data.DataLoader(
        train_dataset,
        batch_size=config.task.dataloader.batch_size,
        shuffle=config.task.dataloader.shuffle,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=use_multiprocessing,
        prefetch_factor=2 if use_multiprocessing else None,
    )

    val_loader: data.DataLoader | None = None
    if not skip_validation:
        val_dataset = EpisodicDataset(
            zarr_path=schema.zarr_path,
            pred_horizon=config.task.prediction_horizon,
            obs_horizon=config.task.observation_horizon,
            dataloader_config=dataloader_config,
            train=False,
            seed=config.experiment.seed,
            action_space=action_space,
            observation_space=observation_space,
            # Downsampling replaces the training buffer with a copy holding
            # only train-selected episodes: sharing it would leak training
            # episodes into the validation split.
            replay_buffer=train_dataset.replay_buffer
            if dataloader_config.downsample_factor <= 1
            else None,
        )
        val_dataset.set_normalizer(normalizer)
        val_dataset.set_tokenizer(tokenizer)

        if action_space.denoise_actions:
            val_dataset.action_processor.denoising_thresholds = (
                train_dataset.action_processor.denoising_thresholds.copy()
            )
            val_dataset.action_processor._denoising_thresholds_computed = True

        val_num_workers = min(4, config.task.dataloader.num_workers)
        val_use_multiprocessing = val_num_workers > 0
        val_loader = data.DataLoader(
            val_dataset,
            batch_size=config.task.dataloader.batch_size,
            shuffle=False,
            num_workers=val_num_workers,
            pin_memory=True,
            persistent_workers=val_use_multiprocessing,
            prefetch_factor=2 if val_use_multiprocessing else None,
        )
    else:
        logging.info("Validation disabled (val_ratio=0). Training without validation.")

    gripper_positive_class_weights = None
    if (
        config.task.action_space.has_gripper_actions
        and config.task.action_space.use_gripper_class_weights
    ):
        gripper_positive_class_weights = (
            train_dataset.get_gripper_positive_class_imbalance_weight()
        )

    return (
        train_loader,
        val_loader,
        normalizer,
        tokenizer,
        gripper_positive_class_weights,
    )

validate_dataloader_config

validate_dataloader_config(config)

Validate Dataloader configuration.

Source code in src/versatil/data/dataloader.py

def validate_dataloader_config(config: DataLoaderConfig) -> None:
    """Validate Dataloader configuration."""
    if config.batch_size <= 0:
        raise ValueError(f"batch_size must be positive, got {config.batch_size}")
    if config.num_workers < 0:
        raise ValueError(f"num_workers cannot be negative, got {config.num_workers}")
    if not 0 <= config.val_ratio < 1:
        raise ValueError(f"val_ratio must be in range [0, 1), got {config.val_ratio}")
    if not 0 < config.total_ratio <= 1:
        raise ValueError(
            f"total_ratio must be in range (0, 1], got {config.total_ratio}"
        )
    if config.skip_initial_episode_steps < 0:
        raise ValueError(
            f"skip_initial_episode_steps cannot be negative, "
            f"got {config.skip_initial_episode_steps}"
        )
    if config.downsample_factor < 1:
        raise ValueError(
            f"downsample_factor must be >= 1, got {config.downsample_factor}"
        )
    for name, quantiles in (
        ("depth_winsorize_quantiles", config.depth_winsorize_quantiles),
        ("kinematics_winsorize_quantiles", config.kinematics_winsorize_quantiles),
    ):
        if quantiles is None:
            continue
        lower, upper = quantiles
        if not 0.0 <= lower <= upper <= 1.0:
            raise ValueError(
                f"{name} must satisfy 0 <= lower <= upper <= 1, got {quantiles}"
            )
    if config.action_backward_shift < 0:
        raise ValueError(
            f"action_backward_shift cannot be negative, "
            f"got {config.action_backward_shift}"
        )
    _validate_uniform_binning_normalization(config=config)