vision_modules

vision_modules

Discovery of camera-addressable visual targets for policy explanations.

get_policy_encoders

get_policy_encoders(policy)

Return encoders registered in the policy encoding pipeline.

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

def get_policy_encoders(policy: Policy) -> dict[str, EncodingMixin]:
    """Return encoders registered in the policy encoding pipeline."""
    return {
        **policy.encoding_pipeline.encoders,
        **policy.encoding_pipeline.conditional_encoders,
    }

get_vision_explainable_modules

get_vision_explainable_modules(policy)

Return visual modules that can produce camera heatmaps.

Parameters:

Name	Type	Description	Default
policy	Policy	Policy whose encoding pipeline and decoder should be inspected.	required

Returns:

Type	Description
list[VisionExplainableModule]	Visual modules from the encoding pipeline and decoder-owned VLM vision
list[VisionExplainableModule]	towers. Each entry includes the camera keys that can be attributed
list[VisionExplainableModule]	through that module and the hook routing mode needed to isolate a
list[VisionExplainableModule]	camera when the same module is reused.

Raises:

Type	Description
RuntimeError	If no camera-addressable visual target is exposed.

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

def get_vision_explainable_modules(policy: Policy) -> list[VisionExplainableModule]:
    """Return visual modules that can produce camera heatmaps.

    Args:
        policy: Policy whose encoding pipeline and decoder should be inspected.

    Returns:
        Visual modules from the encoding pipeline and decoder-owned VLM vision
        towers. Each entry includes the camera keys that can be attributed
        through that module and the hook routing mode needed to isolate a
        camera when the same module is reused.

    Raises:
        RuntimeError: If no camera-addressable visual target is exposed.
    """
    modules = [
        *get_encoding_pipeline_vision_modules(policy=policy),
        *get_decoder_vision_modules(policy=policy),
    ]
    if modules:
        return modules
    raise RuntimeError(
        "No compatible vision explainability modules found. "
        "Explainability requires visual modules that expose target metadata "
        "through get_explainability_targets()."
    )

get_encoding_pipeline_vision_modules

get_encoding_pipeline_vision_modules(policy)

Return explainable visual modules from the encoding pipeline.

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

def get_encoding_pipeline_vision_modules(
    policy: Policy,
) -> list[VisionExplainableModule]:
    """Return explainable visual modules from the encoding pipeline."""
    modules = []
    for encoder_name, encoder in get_policy_encoders(policy=policy).items():
        camera_keys = _camera_keys_for_module(policy=policy, module=encoder)
        if not camera_keys:
            continue
        target = select_explainability_target(
            targets=_get_explainability_targets(module=encoder),
            module_name=encoder_name,
        )
        if target is None:
            continue
        modules.append(
            VisionExplainableModule(
                name=encoder_name,
                module=encoder,
                target=target,
                camera_keys=camera_keys,
                capture_mode=_capture_mode_for_module(
                    module=encoder,
                    camera_keys=camera_keys,
                ),
            )
        )
    return modules

get_decoder_vision_modules

get_decoder_vision_modules(policy)

Return explainable visual modules owned by decoder VLM backbones.

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

def get_decoder_vision_modules(policy: Policy) -> list[VisionExplainableModule]:
    """Return explainable visual modules owned by decoder VLM backbones."""
    vlm_backbone = _decoder_vlm_backbone(policy=policy)
    if vlm_backbone is None:
        return []

    modules: list[VisionExplainableModule] = []
    camera_keys = _camera_keys_for_module(policy=policy, module=vlm_backbone)
    direct_target = select_explainability_target(
        targets=_get_explainability_targets(module=vlm_backbone),
        module_name="decoder.vlm_backbone",
    )
    if direct_target is not None and camera_keys:
        modules.append(
            VisionExplainableModule(
                name="decoder.vlm_backbone",
                module=vlm_backbone,
                target=direct_target,
                camera_keys=camera_keys,
                capture_mode=_capture_mode_for_module(
                    module=vlm_backbone,
                    camera_keys=camera_keys,
                ),
            )
        )

    vision_encoders = getattr(vlm_backbone, "vision_encoders", None)
    if not isinstance(vision_encoders, torch.nn.ModuleList):
        return modules

    for index, vision_encoder in enumerate(vision_encoders):
        target = select_explainability_target(
            targets=_get_explainability_targets(module=vision_encoder),
            module_name=f"decoder.vlm_backbone.vision_encoders.{index}",
        )
        if target is None or not camera_keys:
            continue
        modules.append(
            VisionExplainableModule(
                name=f"decoder.vlm_backbone.vision_encoders.{index}",
                module=vision_encoder,
                target=target,
                camera_keys=camera_keys,
                capture_mode=_capture_mode_for_decoder_vision_tower(
                    camera_keys=camera_keys
                ),
            )
        )
    return modules

resolve_camera_explanation_targets

resolve_camera_explanation_targets(policy, target_camera=None, target_vision_module_names=None)

Resolve runner filters into concrete camera-level targets.

Parameters:

Name	Type	Description	Default
policy	Policy	Policy whose visual modules should be explained.	required
target_camera	str | None	Optional camera key selected by the runner.	None
target_vision_module_names	list[str] | None	Optional visual module allowlist. Names are the values returned by get_vision_explainable_modules().	None

Returns:

Type	Description
list[CameraExplanationTarget]	Concrete camera-target bindings for attribution methods.

Raises:

Type	Description
ValueError	If a configured camera or module allowlist matches nothing.
RuntimeError	If the policy exposes no visual explainability targets.

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

def resolve_camera_explanation_targets(
    policy: Policy,
    target_camera: str | None = None,
    target_vision_module_names: list[str] | None = None,
) -> list[CameraExplanationTarget]:
    """Resolve runner filters into concrete camera-level targets.

    Args:
        policy: Policy whose visual modules should be explained.
        target_camera: Optional camera key selected by the runner.
        target_vision_module_names: Optional visual module allowlist. Names are
            the values returned by ``get_vision_explainable_modules()``.

    Returns:
        Concrete camera-target bindings for attribution methods.

    Raises:
        ValueError: If a configured camera or module allowlist matches nothing.
        RuntimeError: If the policy exposes no visual explainability targets.
    """
    modules = get_vision_explainable_modules(policy=policy)
    if target_vision_module_names is not None:
        target_name_set = set(target_vision_module_names)
        modules = [module for module in modules if module.name in target_name_set]
        if not modules:
            available_names = [
                module.name for module in get_vision_explainable_modules(policy=policy)
            ]
            raise ValueError(
                f"target_vision_module_names={target_vision_module_names} did not "
                f"match visual modules: {available_names}"
            )

    targets = []
    for module in modules:
        camera_keys = (
            (target_camera,)
            if target_camera is not None and target_camera in module.camera_keys
            else module.camera_keys
        )
        if target_camera is not None and target_camera not in module.camera_keys:
            continue
        for camera_key in camera_keys:
            targets.append(
                CameraExplanationTarget(
                    camera_key=camera_key,
                    vision_module_name=module.name,
                    target=module.target,
                    capture_mode=module.capture_mode,
                    invocation_index=_invocation_index_for_camera(
                        camera_key=camera_key,
                        module=module,
                    ),
                    stacked_camera_index=_stacked_camera_index_for_camera(
                        camera_key=camera_key,
                        module=module,
                    ),
                    stacked_camera_count=_stacked_camera_count_for_module(
                        module=module
                    ),
                )
            )
    # TODO: Multiple visual modules can explain the same camera. Attribution
    #  methods aggregate those maps today; expose per-module output names later.
    if targets:
        return targets

    available_cameras = sorted(
        {camera_key for module in modules for camera_key in module.camera_keys}
    )
    raise ValueError(
        f"target_camera={target_camera!r} did not match visual module cameras: "
        f"{available_cameras}"
    )

select_explainability_target

select_explainability_target(targets, module_name)

Select the single image-map target exposed by a visual module.

Parameters:

Name	Type	Description	Default
targets	list[VisionExplanationTarget]	Target metadata exposed by a visual module.	required
module_name	str	Name used in error messages when target metadata is malformed.	required

Returns:

Type	Description
VisionExplanationTarget | None	The compatible target, or None when the module exposes no targets.

Raises:

Type	Description
RuntimeError	If targets are exposed but none are compatible with visual heatmap computation.
RuntimeError	If more than one compatible target is exposed. The runner currently has no config field to disambiguate multiple target layers inside the same visual module.

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

def select_explainability_target(
    targets: list[VisionExplanationTarget],
    module_name: str,
) -> VisionExplanationTarget | None:
    """Select the single image-map target exposed by a visual module.

    Args:
        targets: Target metadata exposed by a visual module.
        module_name: Name used in error messages when target metadata is
            malformed.

    Returns:
        The compatible target, or ``None`` when the module exposes no targets.

    Raises:
        RuntimeError: If targets are exposed but none are compatible with visual
            heatmap computation.
        RuntimeError: If more than one compatible target is exposed. The runner
            currently has no config field to disambiguate multiple target layers
            inside the same visual module.
    """
    if not targets:
        return None
    compatible_targets = [
        target
        for target in targets
        if target.target_kind in COMPATIBLE_EXPLANATION_TARGET_KINDS
    ]
    if len(compatible_targets) == 1:
        return compatible_targets[0]
    if len(compatible_targets) > 1:
        target_kinds = [target.target_kind for target in compatible_targets]
        raise RuntimeError(
            f"Visual module '{module_name}' exposes multiple compatible "
            f"explainability targets {target_kinds}. Configure the module to "
            "expose exactly one target until per-target selection is supported."
        )
    raise RuntimeError(
        f"Visual module '{module_name}' does not expose a compatible "
        "explainability target."
    )