policy

policy

Policy module that handles the sequence of input encoding, output decoding, and loss computation.

Policy

Policy(encoding_pipeline, algorithm, decoder, observation_space, action_space, prediction_horizon, observation_horizon, loss, device, metadata_passthrough=None)

Bases: Module

General policy class that orchestrates observation encoding, action decoding, and loss computation.

Initialize policy.

Parameters:

Name	Type	Description	Default
encoding_pipeline	EncodingPipeline	Observation encoding pipeline.	required
algorithm	DecodingAlgorithm	Decoding algorithm (diffusion, flow matching, etc.).	required
decoder	ActionDecoder	Action decoder architecture.	required
observation_space	ObservationSpace	Observation space configuration.	required
action_space	ActionSpace	Action space configuration.	required
prediction_horizon	int	Number of future actions to predict.	required
observation_horizon	int	Number of past observations to condition on.	required
loss	BaseLoss	Loss module for training.	required
device	str	Device to run on.	required
metadata_passthrough	dict[str, dict[str, str]] | None	Mapping from source dictionaries to metadata keys for logging/visualization.	None

Source code in src/versatil/models/policy.py

def __init__(
    self,
    encoding_pipeline: EncodingPipeline,
    algorithm: DecodingAlgorithm,
    decoder: ActionDecoder,
    observation_space: ObservationSpace,
    action_space: ActionSpace,
    prediction_horizon: int,
    observation_horizon: int,
    loss: BaseLoss,
    device: str,
    metadata_passthrough: dict[str, dict[str, str]] | None = None,
) -> None:
    """Initialize policy.

    Args:
        encoding_pipeline: Observation encoding pipeline.
        algorithm: Decoding algorithm (diffusion, flow matching, etc.).
        decoder: Action decoder architecture.
        observation_space: Observation space configuration.
        action_space: Action space configuration.
        prediction_horizon: Number of future actions to predict.
        observation_horizon: Number of past observations to condition on.
        loss: Loss module for training.
        device: Device to run on.
        metadata_passthrough: Mapping from source dictionaries to metadata
            keys for logging/visualization.
    """
    super().__init__()
    self.encoding_pipeline = encoding_pipeline
    self.algorithm = algorithm
    self.decoder = decoder
    self.observation_space = observation_space
    self.action_space = action_space
    self.prediction_horizon = prediction_horizon
    self.observation_horizon = observation_horizon
    self.loss_module = loss
    self.device = torch.device(device)
    self.metadata_passthrough = self._resolve_metadata_passthrough(
        metadata_passthrough=metadata_passthrough
    )
    self.normalizer: LinearNormalizer = LinearNormalizer()
    self.tokenizer = None  # Set later via set_tokenizer()
    self.denoising_thresholds = DictOfTensorMixin()

input_keys property

input_keys

Sorted observation keys the policy expects as input.

output_keys property

output_keys

Action keys the policy produces, in action-space metadata order.

set_normalizer

set_normalizer(normalizer)

Set normalizer for observations and actions.

Source code in src/versatil/models/policy.py

def set_normalizer(self, normalizer: LinearNormalizer) -> None:
    """Set normalizer for observations and actions."""
    self.normalizer.load_state_dict(normalizer.state_dict())
    self.normalizer.to(self.device)
    self.decoder.set_normalizer(self.normalizer)

set_tokenizer

set_tokenizer(tokenizer)

Set tokenizer and pass it to the decoder.

Source code in src/versatil/models/policy.py

def set_tokenizer(self, tokenizer: Tokenizer | None) -> None:
    """Set tokenizer and pass it to the decoder."""
    self.tokenizer = tokenizer
    self.encoding_pipeline.set_tokenizer(tokenizer)
    self.decoder.set_tokenizer(tokenizer)

set_denoising_thresholds

set_denoising_thresholds(thresholds)

Set the denoising thresholds from training data.

Parameters:

Name	Type	Description	Default
thresholds	dict[str, float]	Dictionary mapping observation keys to their denoising thresholds. May be empty for precomputed actions.	required

Note

These thresholds are computed from the dataset's action processor and stored via DictOfTensorMixin to persist through checkpointing.

Source code in src/versatil/models/policy.py

def set_denoising_thresholds(self, thresholds: dict[str, float]) -> None:
    """Set the denoising thresholds from training data.

    Args:
        thresholds: Dictionary mapping observation keys to their denoising thresholds.
            May be empty for precomputed actions.

    Note:
        These thresholds are computed from the dataset's action processor and stored
        via DictOfTensorMixin to persist through checkpointing.
    """
    for key, value in thresholds.items():
        self.denoising_thresholds.params_dict[key] = nn.Parameter(
            torch.tensor(value), requires_grad=False
        )

get_denoising_thresholds

get_denoising_thresholds()

Return the stored denoising thresholds as plain floats.

Source code in src/versatil/models/policy.py

def get_denoising_thresholds(self) -> dict[str, float]:
    """Return the stored denoising thresholds as plain floats."""
    return {
        key: float(parameter.item())
        for key, parameter in self.denoising_thresholds.params_dict.items()
    }

set_gripper_class_weights

set_gripper_class_weights(pos_weight)

Set positive class weight for GripperLoss components in the loss module.

Parameters:

Name	Type	Description	Default
pos_weight	Tensor | None	Tensor with positive class weight for BCE loss, or None to disable.	required

Source code in src/versatil/models/policy.py

def set_gripper_class_weights(self, pos_weight: torch.Tensor | None) -> None:
    """Set positive class weight for GripperLoss components in the loss module.

    Args:
        pos_weight: Tensor with positive class weight for BCE loss, or None to disable.
    """
    for module in self.loss_module.modules():
        if isinstance(module, GripperLoss):
            module.pos_weight = pos_weight

forward

forward(batch)

Forward pass through observation encoding → action decoding.

Parameters:

Name	Type	Description	Default
batch	dict[str, dict[str, Tensor]]	A batch dictionary containing normalized observations and actions dictionaries. Each is a dict of tensors.	required

Returns:

Type	Description
dict[str, Tensor]	Decoder output dictionary containing action predictions and any architecture-specific outputs.

Source code in src/versatil/models/policy.py

def forward(
    self, batch: dict[str, dict[str, torch.Tensor]]
) -> dict[str, torch.Tensor]:
    """Forward pass through observation encoding → action decoding.

    Args:
        batch: A batch dictionary containing normalized observations and actions dictionaries. Each is a dict of tensors.

    Returns:
        Decoder output dictionary containing action predictions and any architecture-specific outputs.
    """
    observation = self._strip_metadata_passthrough_observations(
        observation=batch[SampleKey.OBSERVATION.value]
    )
    actions = self._filter_predicted_actions(
        actions=batch.get(SampleKey.ACTION.value)
    )
    features = self._build_algorithm_features(observation=observation)
    return self.algorithm.forward(
        features=features, actions=actions, network=self.decoder
    )

compute_loss

compute_loss(batch)

Compute loss using the configured loss module.

The algorithm determines what the regression targets are via get_targets. For BC this is the ground-truth actions; for flow matching it is the target velocity field; for diffusion it depends on the prediction type (noise, sample, or velocity).

Parameters:

Name	Type	Description	Default
batch	dict[str, dict[str, Tensor]]	Batch dictionary containing observations and actions	required

Returns:

Type	Description
LossOutput	LossOutput with total loss and component losses

Source code in src/versatil/models/policy.py

def compute_loss(
    self,
    batch: dict[str, dict[str, torch.Tensor]],
) -> LossOutput:
    """Compute loss using the configured loss module.

    The algorithm determines what the regression targets are via
    ``get_targets``. For BC this is the ground-truth actions; for
    flow matching it is the target velocity field; for diffusion it
    depends on the prediction type (noise, sample, or velocity).

    Args:
        batch: Batch dictionary containing observations and actions

    Returns:
        LossOutput with total loss and component losses
    """
    output = self.forward(batch)
    ground_truth_actions = batch[SampleKey.ACTION.value]
    targets = self.algorithm.get_targets(
        algorithm_output=output,
        ground_truth_actions=ground_truth_actions,
    )
    loss_output = self.loss_module(
        predictions=output,
        targets=targets,
        is_pad=ground_truth_actions.get(SampleKey.IS_PAD_ACTION.value),
    )
    metadata = self._collect_metadata_passthrough(
        batch=batch,
        predictions=output,
    )
    if not metadata:
        return loss_output
    return LossOutput(
        total_loss=loss_output.total_loss,
        component_losses=loss_output.component_losses,
        metadata={**loss_output.metadata, **metadata},
    )

predict_action

predict_action(obs_dict)

Predict actions from observations.

Parameters:

Name	Type	Description	Default
obs_dict	dict[str, Tensor]	Dictionary of observation tensors	required

Returns:

Type	Description
dict[str, Tensor]	Predicted actions (on same device as policy)

Source code in src/versatil/models/policy.py

def predict_action(
    self,
    obs_dict: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Predict actions from observations.

    Args:
        obs_dict: Dictionary of observation tensors

    Returns:
        Predicted actions (on same device as policy)
    """
    obs_dict = to_device(obs_dict, device=self.device)
    normalized_observation = normalize_observation(
        observation=obs_dict,
        normalizer=self.normalizer,
        observation_space=self.observation_space,
    )
    normalized_observation = self._strip_metadata_passthrough_observations(
        observation=normalized_observation
    )
    if (
        self.tokenizer is not None
        and self.tokenizer.observation_tokenizer is not None
    ):
        normalized_observation = tokenize_observation(
            observation=normalized_observation,
            obs_tokenizer=self.tokenizer.observation_tokenizer,
            batched=True,
        )
    features = self._build_algorithm_features(observation=normalized_observation)
    predictions = self.algorithm.predict(features=features, network=self.decoder)
    if DecoderOutputKey.PREDICTED_ACTION_TOKENS.value in predictions:
        action_tokens = predictions[DecoderOutputKey.PREDICTED_ACTION_TOKENS.value]
        if self.tokenizer is None or self.tokenizer.action_tokenizer is None:
            raise RuntimeError(
                "Action tokenizer not set. Cannot detokenize actions."
            )
        normalized_actions = detokenize_actions(
            action_tokens=action_tokens,
            action_tokenizer=self.tokenizer.action_tokenizer,
            action_space=self.action_space,
        )
        normalized_actions = to_device(normalized_actions, device=self.device)
    else:
        normalized_actions = predictions
    actions = unnormalize_actions(
        normalized_actions=normalized_actions,
        normalizer=self.normalizer,
        action_space=self.action_space,
    )
    return actions

build_algorithm_features

build_algorithm_features(observation, encoding_pipeline, decoder, algorithm_injected_keys=None)

Encode observations and select the features the decoder declared.

Merges raw observations with encoding-pipeline outputs, then filters to the decoder's decoder_input.keys allowlist plus their padding masks. Both Policy and ExportablePolicy must build features through this function so that exported models see exactly the training-time inputs.

Parameters:

Name	Type	Description	Default
observation	dict[str, Tensor]	Raw observation dictionary.	required
encoding_pipeline	EncodingPipeline	Pipeline producing encoded features.	required
decoder	ActionDecoder	Decoder whose input specification selects the features.	required
algorithm_injected_keys	set[str] | None	Decoder-declared keys the algorithm adds later (latents, timesteps); they are not required here.	None

Raises:

Type	Description
ValueError	If the decoder requests keys that are neither raw observations, encoding-pipeline outputs, nor algorithm-injected.

Source code in src/versatil/models/policy.py

def build_algorithm_features(
    observation: dict[str, torch.Tensor],
    encoding_pipeline: EncodingPipeline,
    decoder: ActionDecoder,
    algorithm_injected_keys: set[str] | None = None,
) -> dict[str, torch.Tensor]:
    """Encode observations and select the features the decoder declared.

    Merges raw observations with encoding-pipeline outputs, then filters to
    the decoder's ``decoder_input.keys`` allowlist plus their padding masks.
    Both ``Policy`` and ``ExportablePolicy`` must build features through this
    function so that exported models see exactly the training-time inputs.

    Args:
        observation: Raw observation dictionary.
        encoding_pipeline: Pipeline producing encoded features.
        decoder: Decoder whose input specification selects the features.
        algorithm_injected_keys: Decoder-declared keys the algorithm adds
            later (latents, timesteps); they are not required here.

    Raises:
        ValueError: If the decoder requests keys that are neither raw
            observations, encoding-pipeline outputs, nor algorithm-injected.
    """
    injected_keys = algorithm_injected_keys or set()
    encoded_features = encoding_pipeline(observation)
    available_features = {**observation, **encoded_features}
    selected_features: dict[str, torch.Tensor] = {}
    missing_keys: list[str] = []
    for key in decoder.decoder_input.keys:
        if key not in available_features:
            if key not in injected_keys:
                missing_keys.append(key)
            continue
        selected_features[key] = available_features[key]
        padding_key = f"{key}_{EncoderOutputKeys.PADDING_MASK.value}"
        if padding_key in available_features:
            selected_features[padding_key] = available_features[padding_key]

    if missing_keys:
        raise ValueError(
            f"Decoder requested input keys {missing_keys}, but they were not "
            f"available from raw observations or the encoding pipeline. "
            f"Available keys: {sorted(available_features.keys())}."
        )
    return selected_features