base

base

Base checkpoint loader with shared config, tokenizer, and metadata access.

BaseCheckpointLoader

BaseCheckpointLoader(device, checkpoint_path)

Base class for policy checkpoint loaders.

Handles configuration loading, tokenizer setup, and provides shared property accessors for observation/action spaces, horizons, denoising thresholds, and depth clamp ranges.

Subclasses implement concrete checkpoint restoration.

Initialize the base checkpoint loader.

Parameters:

Name	Type	Description	Default
device	device	Device to load the model onto.	required
checkpoint_path	str	Path to the checkpoint directory.	required

Source code in src/versatil/checkpoint_loading/base.py

def __init__(
    self,
    device: torch.device,
    checkpoint_path: str,
) -> None:
    """Initialize the base checkpoint loader.

    Args:
        device: Device to load the model onto.
        checkpoint_path: Path to the checkpoint directory.
    """
    self._device = device
    self._checkpoint_path = checkpoint_path
    self._tokenizer: Tokenizer | None = None
    self._config: MainConfig | None = None
    self._policy: Policy | None = None

device property

device

Get the checkpoint loading device.

checkpoint_path property

checkpoint_path

Get the checkpoint directory path.

config property

config

Get the loaded experiment configuration.

tokenizer property

tokenizer

Get the loaded tokenizer, if any.

policy property

policy

Get the restored policy.

observation_space property

observation_space

Get the policy's observation space.

action_space property

action_space

Get the policy's action space.

prediction_horizon property

prediction_horizon

Get the policy's prediction horizon.

observation_horizon property

observation_horizon

Get the decoder's observation horizon.

denoising_thresholds property

denoising_thresholds

Get denoising thresholds filtered to predicted action keys only.

Returns:

Type	Description
dict[str, float]	Dict mapping VersatIL action key to threshold. Empty if none set.

depth_clamp_ranges property

depth_clamp_ranges

Get per-camera depth clamping ranges from normalizer statistics.

Returns:

Type	Description
dict[str, tuple[float, float]]	Mapping from depth camera key to its (min, max) clamping range;
dict[str, tuple[float, float]]	cameras without normalizer statistics are omitted.

strip_compiled_prefixes

strip_compiled_prefixes(state_dict)

Remove torch.compile module prefixes from checkpoint keys.

Checkpoints saved from a compiled policy carry _orig_mod. segments in their state-dict keys; loading them into an uncompiled module would silently skip every affected weight under strict=False.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Tensor]	Checkpoint state dict, possibly from a compiled module.	required

Source code in src/versatil/checkpoint_loading/base.py

def strip_compiled_prefixes(
    state_dict: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Remove torch.compile module prefixes from checkpoint keys.

    Checkpoints saved from a compiled policy carry ``_orig_mod.`` segments in
    their state-dict keys; loading them into an uncompiled module would
    silently skip every affected weight under ``strict=False``.

    Args:
        state_dict: Checkpoint state dict, possibly from a compiled module.
    """
    compiled_marker = "_orig_mod."
    return {
        key.replace(compiled_marker, ""): value for key, value in state_dict.items()
    }

versatil_checkpoint_safe_globals

versatil_checkpoint_safe_globals()

Classes trusted when unpickling checkpoints with weights_only=True.

Lightning checkpoints embed the hyperparameters saved by save_hyperparameters, which for VersatIL are config dataclasses. Allowlisting exactly those classes keeps third-party checkpoints from executing arbitrary pickled code while our own checkpoints load.

Source code in src/versatil/checkpoint_loading/base.py

def versatil_checkpoint_safe_globals() -> list[type | object]:
    """Classes trusted when unpickling checkpoints with weights_only=True.

    Lightning checkpoints embed the hyperparameters saved by
    ``save_hyperparameters``, which for VersatIL are config dataclasses.
    Allowlisting exactly those classes keeps third-party checkpoints from
    executing arbitrary pickled code while our own checkpoints load.
    """
    configs_package = importlib.import_module("versatil.configs")
    # Lightning stores hyperparameters as an OmegaConf container, whose
    # pickle graph spans these classes.
    safe_classes: list[type | object] = [
        DictConfig,
        ListConfig,
        ContainerMetadata,
        Metadata,
        AnyNode,
        BooleanNode,
        BytesNode,
        EnumNode,
        FloatNode,
        IntegerNode,
        InterpolationResultNode,
        PathNode,
        StringNode,
        ValueNode,
        defaultdict,
        dict,
        list,
        tuple,
        set,
        int,
        float,
        str,
        bool,
        operator.getitem,
        typing.Any,
        # The pickle stream of structured configs references the deprecated
        # typing aliases literally, so the builtins cannot substitute here.
        typing.Dict,  # noqa: UP006
        typing.List,  # noqa: UP006
        typing.Optional,
        typing.Tuple,  # noqa: UP006
        typing.Union,
    ]
    for module_info in pkgutil.walk_packages(
        configs_package.__path__, prefix=f"{configs_package.__name__}."
    ):
        module = importlib.import_module(module_info.name)
        for _, member in inspect.getmembers(module, inspect.isclass):
            if member.__module__.startswith("versatil.configs"):
                safe_classes.append(member)
    return safe_classes