objectives

objectives

Policy prediction objectives used by visual attribution methods.

resolve_actions_for_explanation

resolve_actions_for_explanation(policy, observation, actions, preprocess_observation)

Resolve action targets needed by the policy explanation objective.

Decoders with requires_tokenized_actions=True need action token IDs to compute a differentiable likelihood score. Offline dataset batches can provide true tokenized actions. Online inference batches have no labels, so this function runs an unhooked inference pass first and uses generated action tokens as pseudo-targets for the subsequent hooked teacher-forced pass.

Parameters:

Name	Type	Description	Default
policy	Policy	Policy being explained.	required
observation	ObservationBatch	Observation values keyed by observation-space names.	required
actions	ActionBatch | None	Optional action dictionary from the explanation source.	required
preprocess_observation	bool	Whether to normalize/tokenize observations before building policy features.	required

Returns:

Type	Description
ActionBatch | None	None or the original action batch for decoders that do not require
ActionBatch | None	tokenized actions, existing tokenized labels when available, or
ActionBatch | None	generated pseudo-target actions for tokenized-action decoders.

Raises:

Type	Description
RuntimeError	If a tokenized-action decoder cannot produce action tokens for an unlabeled explanation batch.

Source code in src/versatil/explainability/attribution/objectives.py

def resolve_actions_for_explanation(
    policy: Policy,
    observation: ObservationBatch,
    actions: ActionBatch | None,
    preprocess_observation: bool,
) -> ActionBatch | None:
    """Resolve action targets needed by the policy explanation objective.

    Decoders with ``requires_tokenized_actions=True`` need action token IDs to
    compute a differentiable likelihood score. Offline dataset batches can
    provide true tokenized actions. Online inference batches have no labels, so
    this function runs an unhooked inference pass first and uses generated
    action tokens as pseudo-targets for the subsequent hooked teacher-forced
    pass.

    Args:
        policy: Policy being explained.
        observation: Observation values keyed by observation-space names.
        actions: Optional action dictionary from the explanation source.
        preprocess_observation: Whether to normalize/tokenize observations
            before building policy features.

    Returns:
        ``None`` or the original action batch for decoders that do not require
        tokenized actions, existing tokenized labels when available, or
        generated pseudo-target actions for tokenized-action decoders.

    Raises:
        RuntimeError: If a tokenized-action decoder cannot produce action
            tokens for an unlabeled explanation batch.
    """
    if not _decoder_requires_tokenized_actions(policy=policy):
        return actions
    if actions is not None and SampleKey.TOKENIZED_ACTIONS.value in actions:
        return actions

    prepared_observation = prepare_policy_observation_for_explanation(
        policy=policy,
        observation=observation,
        preprocess_observation=preprocess_observation,
    )
    features = policy._build_algorithm_features(observation=prepared_observation)
    with torch.no_grad():
        predictions = policy.algorithm.predict(
            features=features, network=policy.decoder
        )
    if DecoderOutputKey.PREDICTED_ACTION_TOKENS.value not in predictions:
        raise RuntimeError(
            "Tokenized-action explanation without action labels requires "
            f"'{DecoderOutputKey.PREDICTED_ACTION_TOKENS.value}' from policy inference."
        )
    return {
        SampleKey.TOKENIZED_ACTIONS.value: predictions[
            DecoderOutputKey.PREDICTED_ACTION_TOKENS.value
        ].detach()
    }

compute_policy_explanation_objective

compute_policy_explanation_objective(policy, observation, actions, preprocess_observation, output_selector=None)

Compute the differentiable policy score used by attribution.

Parameters:

Name	Type	Description	Default
policy	Policy	Policy being explained.	required
observation	ObservationBatch	Observation values keyed by observation-space names.	required
actions	ActionBatch | None	Optional action dictionary. Decoders with requires_tokenized_actions=True require tokenized_actions; call resolve_actions_for_explanation before registering attribution hooks when labels are absent.	required
preprocess_observation	bool	Whether to normalize/tokenize observations before building policy features.	required
output_selector	PolicyPredictionSelector | None	Optional selector for continuous prediction tensors. None scores the norm of all normalized action components.	None

Returns:

Type	Description
Tensor	Tensor of per-sample or per-token scores. Attribution methods average
Tensor	this tensor for gradient backpropagation and compare score drops for
Tensor	ablation.

Raises:

Type	Description
ValueError	If a custom selector is passed for a decoder that requires tokenized actions.
RuntimeError	If tokenized-action logits or target tokens are missing.

Source code in src/versatil/explainability/attribution/objectives.py

def compute_policy_explanation_objective(
    policy: Policy,
    observation: ObservationBatch,
    actions: ActionBatch | None,
    preprocess_observation: bool,
    output_selector: PolicyPredictionSelector | None = None,
) -> torch.Tensor:
    """Compute the differentiable policy score used by attribution.

    Args:
        policy: Policy being explained.
        observation: Observation values keyed by observation-space names.
        actions: Optional action dictionary. Decoders with
            ``requires_tokenized_actions=True`` require ``tokenized_actions``;
            call ``resolve_actions_for_explanation`` before registering
            attribution hooks when labels are absent.
        preprocess_observation: Whether to normalize/tokenize observations
            before building policy features.
        output_selector: Optional selector for continuous prediction tensors.
            ``None`` scores the norm of all normalized action components.

    Returns:
        Tensor of per-sample or per-token scores. Attribution methods average
        this tensor for gradient backpropagation and compare score drops for
        ablation.

    Raises:
        ValueError: If a custom selector is passed for a decoder that requires
            tokenized actions.
        RuntimeError: If tokenized-action logits or target tokens are missing.
    """
    prepared_observation = prepare_policy_observation_for_explanation(
        policy=policy,
        observation=observation,
        preprocess_observation=preprocess_observation,
    )
    features = policy._build_algorithm_features(observation=prepared_observation)
    if _decoder_requires_tokenized_actions(policy=policy):
        if output_selector is not None:
            raise ValueError(
                "output_selector is only supported when the decoder does not "
                "require tokenized actions."
            )
        return _compute_tokenized_action_log_likelihood(
            policy=policy,
            features=features,
            actions=actions,
        )

    with EncoderCacheDisabled(decoder=policy.decoder):
        predictions = policy.algorithm.predict(
            features=features,
            network=policy.decoder,
        )
    selector = (
        output_selector if output_selector is not None else default_output_selector
    )
    return selector(predictions)

repeat_action_batch

repeat_action_batch(actions, repeat_count)

Repeat action tensors along the batch axis for perturbation methods.

Parameters:

Name	Type	Description	Default
actions	ActionBatch | None	Action tensors keyed by action component, or None when the attribution objective does not need action labels.	required
repeat_count	int	Number of copies to concatenate.	required

Returns:

Type	Description
ActionBatch | None	Repeated action batch, or None when actions is None.

Raises:

Type	Description
ValueError	If repeat_count is less than one.

Source code in src/versatil/explainability/attribution/objectives.py

def repeat_action_batch(
    actions: ActionBatch | None, repeat_count: int
) -> ActionBatch | None:
    """Repeat action tensors along the batch axis for perturbation methods.

    Args:
        actions: Action tensors keyed by action component, or ``None`` when the
            attribution objective does not need action labels.
        repeat_count: Number of copies to concatenate.

    Returns:
        Repeated action batch, or ``None`` when ``actions`` is ``None``.

    Raises:
        ValueError: If ``repeat_count`` is less than one.
    """
    if actions is None:
        return None
    if repeat_count < 1:
        raise ValueError(f"repeat_count must be positive. Got: {repeat_count}")
    return {
        key: torch.cat([value.clone() for _ in range(repeat_count)], dim=0)
        for key, value in actions.items()
    }