runner

runner

Hydra-instantiated runner for policy explainability workflows.

ExplainabilityRunner

ExplainabilityRunner(checkpoint_path, checkpoint_name=value, output_directory=None, device='auto', source=value, split='all', sample_stride=50, max_samples=None, data_path_override=None, batch_size=1, channel_batch_size=32, explanation_types=None, target_camera_keys=None, target_vision_module_names=None, online=None, writer=None)

Generate xAI insights for the predictions of a trained policy.

Initialize the runner.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	Directory containing the training config.yaml and the model checkpoint.	required
checkpoint_name	str	Checkpoint filename to restore trained policy.	value
output_directory	str | None	Directory for written explanation files. None writes under checkpoint_path/explainability/<timestamp>.	None
device	str	Device name, or auto to prefer CUDA when available.	'auto'
source	str	Explanation source type. Standalone run() supports dataset and online_inference.	value
split	str	Dataset split for offline explanations.	'all'
sample_stride	int	Explanation interval. In offline dataset mode, keep every Nth episodic dataset sample. In online inference mode, explain every Nth inference timestep.	50
max_samples	int | None	Optional cap on the number of observation windows to explain. Offline mode applies this after sample_stride. Online mode applies it to ready inference windows and derives the inference-loop step budget from sample_stride.	None
data_path_override	str | list[str] | None	Optional offline input location to explain instead of the data path stored in the checkpoint task config. None uses the checkpoint schema as-is. A single .zarr path is sampled directly. A non-zarr path is 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; multiple zarr paths are rejected. Raw inputs are converted to offline_dataset.zarr beside the first override path unless zarr_cache_directory is set by the dataset source.	None
batch_size	int	Number of samples per attribution call.	1
channel_batch_size	int	Number of channels per Ablation-CAM forward.	32
explanation_types	list[str] | None	Explanation methods to run. If None, it runs all supported methods.	None
target_camera_keys	list[str] | None	Optional camera-key allowlist for generated heatmaps.	None
target_vision_module_names	list[str] | None	Optional visual module allowlist. Names include encoding-pipeline entries and decoder-owned VLM vision tower paths.	None
online	InferenceClientConfig | None	Socket inference client settings used by source=online_inference. None uses defaults.	None
writer	ExplanationWriterConfig | None	Explanation writer settings controlling raw heatmaps, overlays, blending, and image format. None uses defaults.	None

Raises:

Type	Description
ValueError	If explanation types or source data are invalid.

Source code in src/versatil/explainability/runner.py

def __init__(
    self,
    checkpoint_path: str,
    checkpoint_name: str = CheckpointFilename.DEFAULT_CHECKPOINT.value,
    output_directory: str | None = None,
    device: str = "auto",
    source: str = ExplanationSourceType.DATASET.value,
    split: str = "all",
    sample_stride: int = 50,
    max_samples: int | None = None,
    data_path_override: str | list[str] | None = None,
    batch_size: int = 1,
    channel_batch_size: int = 32,
    explanation_types: list[str] | None = None,
    target_camera_keys: list[str] | None = None,
    target_vision_module_names: list[str] | None = None,
    online: InferenceClientConfig | None = None,
    writer: ExplanationWriterConfig | None = None,
) -> None:
    """Initialize the runner.

    Args:
        checkpoint_path: Directory containing the training ``config.yaml``
            and the model checkpoint.
        checkpoint_name: Checkpoint filename to restore trained policy.
        output_directory: Directory for written explanation files. ``None``
            writes under ``checkpoint_path/explainability/<timestamp>``.
        device: Device name, or ``auto`` to prefer CUDA when available.
        source: Explanation source type. Standalone ``run()`` supports
            ``dataset`` and ``online_inference``.
        split: Dataset split for offline explanations.
        sample_stride: Explanation interval. In offline dataset mode, keep
            every Nth episodic dataset sample. In online inference mode,
            explain every Nth inference timestep.
        max_samples: Optional cap on the number of observation windows to
            explain. Offline mode applies this after ``sample_stride``. Online
            mode applies it to ready inference windows and derives the
            inference-loop step budget from ``sample_stride``.
        data_path_override: Optional offline input location to explain
            instead of the data path stored in the checkpoint task config.
            ``None`` uses the checkpoint schema as-is. A single ``.zarr``
            path is sampled directly. A non-zarr path is 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; multiple zarr paths are rejected. Raw
            inputs are converted to ``offline_dataset.zarr`` beside the
            first override path unless ``zarr_cache_directory`` is set by
            the dataset source.
        batch_size: Number of samples per attribution call.
        channel_batch_size: Number of channels per Ablation-CAM forward.
        explanation_types: Explanation methods to run. If None, it runs all
            supported methods.
        target_camera_keys: Optional camera-key allowlist for generated
            heatmaps.
        target_vision_module_names: Optional visual module allowlist. Names
            include encoding-pipeline entries and decoder-owned VLM vision
            tower paths.
        online: Socket inference client settings used by
            ``source=online_inference``. ``None`` uses defaults.
        writer: Explanation writer settings controlling raw heatmaps,
            overlays, blending, and image format. ``None`` uses defaults.

    Raises:
        ValueError: If explanation types or source data are invalid.
    """
    self.checkpoint_path = checkpoint_path
    self.checkpoint_name = checkpoint_name
    self.output_directory = self._resolve_output_directory(
        checkpoint_path=checkpoint_path,
        output_directory=output_directory,
    )
    self.device = self._resolve_device(device=device)
    self.source = source
    self.split = split
    self.sample_stride = sample_stride
    self.max_samples = max_samples
    self.data_path_override = data_path_override
    self.batch_size = batch_size
    self.online = online if online is not None else InferenceClientConfig()
    self.explanation_types = self._resolve_explanation_types(
        explanation_types=explanation_types
    )
    self.target_camera_keys = target_camera_keys
    self.target_vision_module_names = target_vision_module_names
    self.writer_config = writer if writer is not None else ExplanationWriterConfig()
    self.channel_batch_size = channel_batch_size
    self.explanation_heatmaps = to_explanation_heatmaps(
        channel_batch_size=channel_batch_size
    )
    self._validate_source(source=source)
    self._validate_sampling_configuration(
        sample_stride=sample_stride,
        max_samples=max_samples,
    )
    if source == ExplanationSourceType.ONLINE_INFERENCE.value:
        self._validate_online_configuration(
            model_server_port=self.online.model_server_port,
            action_execution_horizon=self.online.action_execution_horizon,
            update_rate_hz=self.online.update_rate_hz,
            temporal_max_timesteps=self.online.temporal_max_timesteps,
            compression_type=self.online.compression_type,
        )
    self.output_directory.mkdir(parents=True, exist_ok=True)
    self.writer = ExplanationWriter(
        output_directory=self.output_directory,
        image_weight=self.writer_config.image_weight,
        overlay_image_format=self.writer_config.overlay_image_format,
    )

    checkpoint_loader = FloatCheckpointLoader(
        device=self.device,
        checkpoint_path=checkpoint_path,
        checkpoint_name=checkpoint_name,
    )
    self.checkpoint_loader = checkpoint_loader
    self.config = checkpoint_loader.config
    self.policy: Policy = checkpoint_loader.policy
    self.policy.eval()
    self._batch_counter = 0

run

run()

Run the configured explainability workflow.

Raises:

Type	Description
ValueError	If source is not a supported explainability mode.

Source code in src/versatil/explainability/runner.py

def run(self) -> None:
    """Run the configured explainability workflow.

    Raises:
        ValueError: If ``source`` is not a supported explainability mode.
    """
    match self.source:
        case ExplanationSourceType.ONLINE_INFERENCE.value:
            self._run_online_inference()
        case ExplanationSourceType.DATASET.value:
            self._run_dataset_source()
        case _:
            self._validate_source(source=self.source)
            raise ValueError(f"Unhandled explainability source: {self.source}")

build_online_source

build_online_source()

Build an online inference adapter for InferenceClient.

Returns:

Type	Description
OnlineInferenceExplanationSource	Online source that delegates accepted inference windows to this
OnlineInferenceExplanationSource	runner.

Source code in src/versatil/explainability/runner.py

def build_online_source(self) -> OnlineInferenceExplanationSource:
    """Build an online inference adapter for ``InferenceClient``.

    Returns:
        Online source that delegates accepted inference windows to this
        runner.
    """
    return OnlineInferenceExplanationSource(
        consumer=self,
        sample_stride=self.sample_stride,
        max_samples=self.max_samples,
    )

explain_batch

explain_batch(batch)

Generate configured explanations for one batch.

Parameters:

Name	Type	Description	Default
batch	ExplanationBatch	Observation window and metadata from a supported source.	required

Source code in src/versatil/explainability/runner.py

def explain_batch(self, batch: ExplanationBatch) -> None:
    """Generate configured explanations for one batch.

    Args:
        batch: Observation window and metadata from a supported source.
    """
    moved_observation = to_device(data=batch.observation, device=self.device)
    if not isinstance(moved_observation, dict):
        raise RuntimeError(
            f"Expected observation dictionary, got {type(moved_observation)}."
        )
    observation = dict(moved_observation)
    actions = None
    if batch.actions is not None:
        moved_actions = to_device(data=batch.actions, device=self.device)
        if not isinstance(moved_actions, dict):
            raise RuntimeError(
                f"Expected action dictionary, got {type(moved_actions)}."
            )
        actions = dict(moved_actions)
    for explanation_type in self.explanation_types:
        heatmaps = self._compute_heatmaps(
            observation=observation,
            actions=actions,
            explanation_type=explanation_type,
            preprocess_observation=batch.preprocess_observation,
        )
        if self.writer_config.save_raw_heatmaps:
            self.writer.save_raw_heatmaps(
                heatmaps=heatmaps,
                explanation_type=explanation_type,
                metadata=batch.metadata,
                batch_counter=self._batch_counter,
            )
        if self.writer_config.save_overlays:
            self.writer.save_overlays(
                heatmaps=heatmaps,
                explanation_type=explanation_type,
                batch=batch,
                batch_counter=self._batch_counter,
            )
    self._batch_counter += 1