dataset

dataset

Dataset-based explanation source.

DatasetExplanationSource

DatasetExplanationSource(config, policy, split, batch_size, sample_stride, max_samples, data_path_override=None, zarr_cache_directory=None)

Yield explanation batches from versatil schema-based zarr data.

Note

Missing zarr data is created from raw data described by the dataset schema. Override inputs must use the same schema type as the checkpoint configuration.

Initialize the dataset-backed source.

Parameters:

Name	Type	Description	Default
config	DatasetRunnerConfig	Instantiated training configuration loaded from the checkpoint.	required
policy	Policy	Loaded policy whose normalizer should be reused by the dataset sample builder.	required
split	str	Dataset split to explain: train, val, or all.	required
batch_size	int	Number of sampled windows per explanation batch.	required
sample_stride	int	Keep every Nth sample in deterministic dataset order.	required
max_samples	int | None	Optional cap on the number of sampled windows.	required
data_path_override	str | list[str] | None	Optional offline input location to explain instead of the training data path stored in the checkpoint task config. - None uses the training data as-is. - A single .zarrpath is used directly - A non-zarr path has to be raw data in the same dataset schema as the checkpoint. - A list is only for raw schemas that already accept multiple inputs, such as CSV folders or HDF5 files.	None
zarr_cache_directory	Path | None	Directory used to store zarr data created from raw data_path_override inputs. When None and data_path_override is set, the zarr is written beside the override path. When both are None, the checkpoint schema zarr directory is used.	None

Raises:

Type	Description
ValueError	If split, batch_size, sample_stride, or max_samples is invalid.

Source code in src/versatil/explainability/sources/dataset.py

def __init__(
    self,
    config: DatasetRunnerConfig,
    policy: Policy,
    split: str,
    batch_size: int,
    sample_stride: int,
    max_samples: int | None,
    data_path_override: str | list[str] | None = None,
    zarr_cache_directory: Path | None = None,
) -> None:
    """Initialize the dataset-backed source.

    Args:
        config: Instantiated training configuration loaded from the
            checkpoint.
        policy: Loaded policy whose normalizer should be reused by the
            dataset sample builder.
        split: Dataset split to explain: ``train``, ``val``, or ``all``.
        batch_size: Number of sampled windows per explanation batch.
        sample_stride: Keep every Nth sample in deterministic dataset order.
        max_samples: Optional cap on the number of sampled windows.
        data_path_override: Optional offline input location to explain
            instead of the training data path stored in the checkpoint task config.
            - ``None`` uses the training data as-is.
            - A single ``.zarr``path is used directly
            - A non-zarr path has to be raw data in the same dataset schema as the checkpoint.
            - A list is only for raw schemas that already accept multiple inputs, such as CSV
              folders or HDF5 files.
        zarr_cache_directory: Directory used to store zarr data created from
            raw ``data_path_override`` inputs. When ``None`` and
            ``data_path_override`` is set, the zarr is written beside the
            override path. When both are ``None``, the checkpoint schema
            zarr directory is used.

    Raises:
        ValueError: If ``split``, ``batch_size``, ``sample_stride``, or
            ``max_samples`` is invalid.
    """
    valid_splits = [member.value for member in ExplanationDatasetSplit]
    if split not in valid_splits:
        raise ValueError(f"split must be one of {valid_splits}. Got: {split}")
    if batch_size <= 0:
        raise ValueError(f"batch_size must be positive. Got: {batch_size}")
    if sample_stride <= 0:
        raise ValueError(f"sample_stride must be positive. Got: {sample_stride}")
    if max_samples is not None and max_samples <= 0:
        raise ValueError(
            f"max_samples must be positive when set. Got: {max_samples}"
        )

    self.config = config
    self.policy = policy
    self.split = split
    self.batch_size = batch_size
    self.sample_stride = sample_stride
    self.max_samples = max_samples
    self.data_path_override = data_path_override
    self.zarr_cache_directory = self._resolve_zarr_cache_directory(
        data_path_override=data_path_override,
        zarr_cache_directory=zarr_cache_directory,
    )
    self.dataset_schema = resolve_dataset_schema_for_explanation(
        schema=self.config.task.dataset_schema,
        data_path_override=data_path_override,
        zarr_cache_directory=self.zarr_cache_directory,
    )
    self.dataset = self._build_dataset()

__iter__

__iter__()

Yield deterministic explanation batches.

Returns:

Type	Description
Iterator[ExplanationBatch]	Iterator of model-ready batches. Dataset samples are already
Iterator[ExplanationBatch]	normalized/tokenized, so preprocess_observation is False.

Source code in src/versatil/explainability/sources/dataset.py

def __iter__(self) -> Iterator[ExplanationBatch]:
    """Yield deterministic explanation batches.

    Returns:
        Iterator of model-ready batches. Dataset samples are already
        normalized/tokenized, so ``preprocess_observation`` is ``False``.
    """
    selected_indices = list(range(0, len(self.dataset), self.sample_stride))
    if self.max_samples is not None:
        selected_indices = selected_indices[: self.max_samples]

    for batch_start in range(0, len(selected_indices), self.batch_size):
        sample_indices = selected_indices[
            batch_start : batch_start + self.batch_size
        ]
        samples = [self.dataset[index] for index in sample_indices]
        sample_batch = default_collate(samples)
        observation = sample_batch[SampleKey.OBSERVATION.value]
        actions = sample_batch.get(SampleKey.ACTION.value)
        yield ExplanationBatch(
            observation=observation,
            actions=actions,
            display_observation=self._extract_display_observation(
                observation=observation
            ),
            metadata={
                "source": ExplanationSourceType.DATASET.value,
                "split": self.split,
                "sample_indices": sample_indices,
                "zarr_path": self.dataset_schema.zarr_path,
            },
            preprocess_observation=False,
        )