compressed_runtime

compressed_runtime ¶

Compressed policy inference runtime.

CompressedPolicyRuntime ¶

CompressedPolicyRuntime(device, checkpoint_path, compile_model=True)

Bases: PolicyRuntime

Inference runtime for compressed policy checkpoints.

Initialize the compressed policy runtime.

Parameters:

Name	Type	Description	Default
`device`	`device`	Device to load the model onto.	required
`checkpoint_path`	`str`	Path to the compressed checkpoint directory containing compression_metadata.json, the model artifact (Torch Export `.pt2` or ExecuTorch `.pte`), and normalizer.pt.	required
`compile_model`	`bool`	Whether to compile the model with torch.compile before inference. For PT2E models on CPU this enables inductor int8 kernel fusion and is essential for performance. For other workflows it applies standard inductor compilation.	`True`

Source code in src/versatil/inference/policy_runtime/compressed_runtime.py

def __init__(
    self,
    device: torch.device,
    checkpoint_path: str,
    compile_model: bool = True,
) -> None:
    """Initialize the compressed policy runtime.

    Args:
        device: Device to load the model onto.
        checkpoint_path: Path to the compressed checkpoint directory
            containing compression_metadata.json, the model artifact
            (Torch Export ``.pt2`` or ExecuTorch ``.pte``), and
            normalizer.pt.
        compile_model: Whether to compile the model with
            torch.compile before inference. For PT2E models on
            CPU this enables inductor int8 kernel fusion and is
            essential for performance. For other workflows it
            applies standard inductor compilation.
    """
    self._compile_model = compile_model
    self._compressed_model: nn.Module | None = None
    checkpoint_loader = CompressedCheckpointLoader(
        device=device,
        checkpoint_path=checkpoint_path,
    )
    super().__init__(
        checkpoint_loader=checkpoint_loader,
        client_identifier=checkpoint_path,
    )
    self._compressed_model = self._load_artifact_model(
        model_path=checkpoint_loader.model_path,
        workflow=checkpoint_loader.workflow,
    )

input_keys `property` ¶

input_keys

Get compressed model input key ordering.

output_keys `property` ¶

output_keys

Get compressed model output key ordering.

run_inference ¶

run_inference(obs_dict)

Run compressed policy inference with normalization.

Normalizes observations, optionally tokenizes, converts to positional tensors, runs the compressed model, and converts output back to an unnormalized action dict.

Parameters:

Name	Type	Description	Default
`obs_dict`	`dict[str, Tensor]`	Observation dictionary for the policy.	required

Returns:

Type	Description
`dict[str, Tensor]`	Unnormalized action dictionary.

Source code in src/versatil/inference/policy_runtime/compressed_runtime.py

def run_inference(
    self, obs_dict: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]:
    """Run compressed policy inference with normalization.

    Normalizes observations, optionally tokenizes, converts to
    positional tensors, runs the compressed model, and converts
    output back to an unnormalized action dict.

    Args:
        obs_dict: Observation dictionary for the policy.

    Returns:
        Unnormalized action dictionary.
    """
    obs_dict = to_device(obs_dict, device=self.device)
    normalized_obs = normalize_observation(
        observation=obs_dict,
        normalizer=self.checkpoint_loader.normalizer,
        observation_space=self.observation_space,
    )
    tokenizer = self.tokenizer
    if tokenizer is not None and tokenizer.observation_tokenizer is not None:
        normalized_obs = tokenize_observation(
            observation=normalized_obs,
            obs_tokenizer=tokenizer.observation_tokenizer,
            batched=True,
        )
    observation_tensors = tuple(normalized_obs[key] for key in self.input_keys)
    output_tensors = self._run_compressed_model(
        observation_tensors=observation_tensors,
    )
    if len(output_tensors) != len(self.output_keys):
        raise ValueError(
            f"Compressed model returned {len(output_tensors)} tensors, "
            f"but metadata declares {len(self.output_keys)} output keys."
        )
    normalized_actions = {
        key: output_tensors[index] for index, key in enumerate(self.output_keys)
    }
    return unnormalize_actions(
        normalized_actions=normalized_actions,
        normalizer=self.checkpoint_loader.normalizer,
        action_space=self.action_space,
    )

compressed_runtime