opentau.configs.policies

Policy configuration module.

This module provides the base PreTrainedConfig class for policy models, which defines the interface and common functionality for all policy configurations. It includes support for feature definitions, normalization modes, and loading configurations from pretrained models or local paths.

Functions

load_resolved_config_dict(source_path)	Resolve $ref includes in a config JSON and return the resulting dict.
load_stripped_config_to_tempfile(source_path)	Read a config JSON, resolve $ref includes, strip deprecated/removed fields in memory, write to a fresh temp file, and return its path.
strip_deprecated_fields_from_json(path)	Remove deprecated and removed fields from a config JSON file in-place.
warn_deprecated_latency_fields(config_path)	Emit a deprecation warning if a config JSON file contains latency fields.
warn_deprecated_latency_fields_from_dict(...)	Like warn_deprecated_latency_fields() but operating on an already-resolved config dict, so callers that already loaded the dict don't pay for a second $ref walk.
warn_removed_policy_fields(config_path)	Emit a deprecation warning if a config JSON file contains removed fields.
warn_removed_policy_fields_from_dict(data, ...)	Like warn_removed_policy_fields() but operating on an already-resolved config dict.
write_stripped_config_to_tempfile(data)	Strip deprecated/removed fields from data and write to a temp file.

Classes

PreTrainedConfig(n_obs_steps, ...)

Base configuration class for policy models.

class opentau.configs.policies.PreTrainedConfig(n_obs_steps: int = 1, normalization_mapping: dict[str, ~opentau.configs.types.NormalizationMode] = <factory>, input_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, output_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, device: str | None = None, use_amp: bool = False, use_torch_compile: bool = False, torch_compile_mode: str = 'default', pretrained_path: str | None = None, skip_normalization_weights: bool = False, skip_input_resolution_check: bool = False, save_normalization_stats: bool = True, dataset_names: list[str] | None = None, dataset_to_norm_index: dict[str, int] | None = None, cloud_vlm_latency_mean: float = 0.0, cloud_vlm_latency_std: float = 0.0, cloud_vlm_latency_lower: float = 0.0, cloud_vlm_latency_upper: float = 0.0, action_decoder_latency_mean: float = 0.0, action_decoder_latency_std: float = 0.0, action_decoder_latency_lower: float = 0.0, action_decoder_latency_upper: float = 0.0)

Bases: ChoiceRegistry, HubMixin, ABC

Base configuration class for policy models.

Parameters:

• n_obs_steps – Number of environment steps worth of observations to pass to the policy (takes the current step and additional steps going back).

• input_shapes – A dictionary defining the shapes of the input data for the policy.

• output_shapes – A dictionary defining the shapes of the output data for the policy.

• input_normalization_modes – A dictionary with key representing the modality and the value specifies the normalization mode to apply.

• output_normalization_modes – Similar dictionary as input_normalization_modes, but to unnormalize to the original scale.

• skip_normalization_weights – When loading via from_pretrained(), drop the saved normalize_* / unnormalize_* buffer tensors from the state dict before load_state_dict. The buffers freshly initialised from dataset_stats then survive — use this when finetuning a checkpoint whose saved normalization stats were aggregated over a different dataset mixture than the finetuning data. The buffers are registered as nn.Parameter(requires_grad=False), so training alone cannot recover from inheriting the wrong stats. Requires ``dataset_stats`` to be supplied to ``__init__`` (e.g. via opentau.policies.factory.make_policy()); otherwise the buffers stay at the inf sentinel from opentau.policies.normalize.create_stats_buffers() and the next forward crashes. One-shot (in-memory only): after the strip helper runs (whether keys were dropped or the “had no effect” warning fired), the flag is reset to False on the in-memory model’s config — a user who opted in once should not have to remember to flip it back before resume. Persistence requires save_pretrained: an in-process resume that goes train.py → save_checkpoint → cfg.save_pretrained(...) writes False to the new checkpoint’s config.json, so the next from_pretrained reads False and skips the strip. But re-running from_pretrained on the original source checkpoint (e.g. an interactive notebook session that hasn’t yet saved) will re-read True from the source config.json and re-strip. The model was trained against the saved stats, so switching the normalization mid-training only makes sense when followed by further training, not when loading purely for inference. Honored by every policy whose from_pretrained (or _load_as_safetensor) calls _strip_normalization_buffers_from_state_dict(). Defaults to False (no behaviour change).

• skip_input_resolution_check – Escape hatch for validate_input_resolution(). When True, a mismatch between the policy’s resize_imgs_with_padding and the (H, W) of the bound image features logs a loud warning instead of raising — even on the strict (training) path. Intended for deliberately resuming/finetuning a legacy checkpoint that was trained with the mismatch (i.e. with the policy silently letterboxing a second time inside prepare_images/prepare_videos); leave False for everything else. Eval-shaped policy construction never raises regardless of this flag (it warns), so plain evaluation of legacy checkpoints does not need it. Note that True persists into the config.json of every checkpoint the run saves — coherent for resuming that same mismatch-trained lineage, but set it back to false when fine-tuning such a checkpoint against a new dataset mixture, or a fresh mismatch there will only warn instead of raising. Defaults to False.

__init__(n_obs_steps: int = 1, normalization_mapping: dict[str, ~opentau.configs.types.NormalizationMode] = <factory>, input_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, output_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, device: str | None = None, use_amp: bool = False, use_torch_compile: bool = False, torch_compile_mode: str = 'default', pretrained_path: str | None = None, skip_normalization_weights: bool = False, skip_input_resolution_check: bool = False, save_normalization_stats: bool = True, dataset_names: list[str] | None = None, dataset_to_norm_index: dict[str, int] | None = None, cloud_vlm_latency_mean: float = 0.0, cloud_vlm_latency_std: float = 0.0, cloud_vlm_latency_lower: float = 0.0, cloud_vlm_latency_upper: float = 0.0, action_decoder_latency_mean: float = 0.0, action_decoder_latency_std: float = 0.0, action_decoder_latency_lower: float = 0.0, action_decoder_latency_upper: float = 0.0) → None

action_decoder_latency_lower: float = 0.0

action_decoder_latency_mean: float = 0.0

action_decoder_latency_std: float = 0.0

action_decoder_latency_upper: float = 0.0

abstract property action_delta_indices: list | None

Get indices for action delta features.

Returns:

List of indices indicating which action features should be treated as deltas, or None if no delta features are used.

property action_feature: PolicyFeature | None

Get the action feature from output features.

Returns:

The PolicyFeature with type ACTION if found, or None otherwise.

cloud_vlm_latency_lower: float = 0.0

cloud_vlm_latency_mean: float = 0.0

cloud_vlm_latency_std: float = 0.0

cloud_vlm_latency_upper: float = 0.0

dataset_names: list[str] | None = None

dataset_to_norm_index: dict[str, int] | None = None

device: str | None = None

property env_state_feature: PolicyFeature | None

Get the environment state feature from input features.

Returns:

The PolicyFeature with type ENV if found, or None otherwise.

classmethod from_pretrained(pretrained_name_or_path: str | Path, *, force_download: bool = False, resume_download: bool | None = None, proxies: dict | None = None, token: str | bool | None = None, cache_dir: str | Path | None = None, local_files_only: bool = False, revision: str | None = None, **policy_kwargs) → T

Load a policy configuration from a pretrained model or local path.

Parameters:

• cls – The class to instantiate.

• pretrained_name_or_path –

  Can be either:

  • A string, the model id of a pretrained config hosted inside a model repo on huggingface.co.

  • A path to a directory containing a configuration file saved using the _save_pretrained method.

• force_download – Whether to force (re-)downloading the config files and configuration from the HuggingFace Hub. Defaults to False.

• resume_download – Whether to resume downloading the config files. Defaults to None.

• proxies – Dictionary of proxies to use for requests. Defaults to None.

• token – The token to use as HTTP bearer authorization. If True, will use the token generated when running huggingface-cli login. Defaults to None.

• cache_dir – Path to a directory in which a downloaded pretrained model configuration should be cached. Defaults to None.

• local_files_only – Whether to only look at local files (i.e., do not try to download the config). Defaults to False.

• revision – The specific model version to use. It can be a branch name, a tag name, or a commit id. Defaults to None.

• **policy_kwargs – Additional keyword arguments. May include ‘cli_overrides’ for command-line argument overrides.

Returns:

An instance of the configuration class loaded from the specified path.

Raises:

FileNotFoundError – If the configuration file is not found on the HuggingFace Hub or in the local path.

abstract get_optimizer_preset() → OptimizerConfig

Get the default optimizer configuration for this policy.

Returns:

An OptimizerConfig instance with default settings for this policy type.

abstract get_scheduler_preset() → LRSchedulerConfig | None

Get the default learning rate scheduler configuration for this policy.

Returns:

An LRSchedulerConfig instance with default settings for this policy type, or None if no scheduler should be used.

property image_features: dict[str, PolicyFeature]

Get all visual/image features from input features.

Returns:

Dictionary mapping feature names to PolicyFeature instances with type VISUAL.

input_features: dict[str, PolicyFeature]

property input_image_size: tuple[int, int] | None

(H, W) the vision tower will actually receive.

This is resize_imgs_with_padding when the policy defines and sets it (the in-policy letterbox target), else the resolution of the bound image features (native pass-through), else None when neither is known (e.g. a bare config before feature binding).

n_obs_steps: int = 1

normalization_mapping: dict[str, NormalizationMode]

abstract property observation_delta_indices: list | None

Get indices for observation delta features.

Returns:

List of indices indicating which observation features should be treated as deltas, or None if no delta features are used.

output_features: dict[str, PolicyFeature]

pretrained_path: str | None = None

abstract property reward_delta_indices: list | None

Get indices for reward delta features.

Returns:

List of indices indicating which reward features should be treated as deltas, or None if no delta features are used.

property robot_state_feature: PolicyFeature | None

Get the robot state feature from input features.

Returns:

The PolicyFeature with type STATE if found, or None otherwise.

save_normalization_stats: bool = True

skip_input_resolution_check: bool = False

skip_normalization_weights: bool = False

torch_compile_mode: str = 'default'

property type: str

Get the type name of this configuration.

Returns:

The choice name of this configuration class.

use_amp: bool = False

use_torch_compile: bool = False

abstract validate_features() → None

Validate that the feature configuration is correct.

This method should check that all required features are present and have valid configurations.

Raises:

ValueError – If the feature configuration is invalid.

validate_input_resolution(*, strict: bool) → None

Check resize_imgs_with_padding against the bound image features.

The image features bound onto the config (from TrainPipelineConfig.resolution via the dataset mixture at train time, or carried inside a checkpoint’s config.json at load time) are the (H, W) the policy actually receives. A differing resize_imgs_with_padding means every frame gets silently letterboxed a second time inside the policy — downscaled and padded with black bars — which is almost never intended.

With resize_imgs_with_padding=None (native pass-through) there is no target to compare against, but the bound cameras must then agree with each other — with no in-policy resize step, a mixed-resolution camera set has no single geometry the vision tower could be built for, so that is flagged with the same strict/warn semantics.

No-op when the policy has no resize_imgs_with_padding field or no comparable image feature is bound.

Parameters:

strict – True for training-shaped construction — raise on mismatch (unless skip_input_resolution_check); False for eval/inference — warn loudly but keep loading, so legacy checkpoints trained with the mismatch remain evaluable.

Raises:

ValueError – On mismatch when strict is True and skip_input_resolution_check is False.

opentau.configs.policies.load_resolved_config_dict(source_path: str | Path) → dict

Resolve $ref includes in a config JSON and return the resulting dict.

Centralizes both the ref resolution and the top-level shape check so that every loader in this module fails consistently on a non-object root, and so that a single from_pretrained call can resolve once and pass the dict to all downstream helpers (warnings, stripping) instead of re-walking the ref tree from disk for each.

opentau.configs.policies.load_stripped_config_to_tempfile(source_path: str | Path) → Path

Read a config JSON, resolve $ref includes, strip deprecated/removed fields in memory, write to a fresh temp file, and return its path. Does not mutate source_path.

Use this on incoming config files (HF Hub downloads, user-supplied paths) before handing them to draccus.parse: HF cache paths are symlinks into a content-addressed blob store, so an in-place rewrite would corrupt the cache; user-supplied paths shouldn’t be mutated as a side effect of load.

See opentau.configs.refs for $ref semantics.

opentau.configs.policies.strip_deprecated_fields_from_json(path: Path) → None

Remove deprecated and removed fields from a config JSON file in-place.

Only safe to call on files we own (e.g. the JSON we just wrote inside a save directory). Do NOT use on user-supplied inputs or HF cache files — HF cache paths are symlinks into a content-addressed blob store, so an in-place rewrite mutates the blob and silently corrupts the cache. For incoming configs, use load_stripped_config_to_tempfile() instead.

opentau.configs.policies.warn_deprecated_latency_fields(config_path: str | Path) → None

Emit a deprecation warning if a config JSON file contains latency fields.

Checks both top-level fields and fields nested under a "policy" key. Should be called before loading a config so users are aware the fields will be ignored. Resolves $ref includes so fields hidden behind a reference are still detected.

opentau.configs.policies.warn_deprecated_latency_fields_from_dict(data: dict, config_path: str | Path) → None

Like warn_deprecated_latency_fields() but operating on an already-resolved config dict, so callers that already loaded the dict don’t pay for a second $ref walk. config_path is used only for the warning message.

opentau.configs.policies.warn_removed_policy_fields(config_path: str | Path) → None

Emit a deprecation warning if a config JSON file contains removed fields.

Removed fields (e.g. init_strategy) are stripped at load time so old saved configs continue to parse, but the user’s explicit choice is dropped on the floor — surface that with a warning so they know to re-save and update any downstream tooling that still emits the old key. Resolves $ref includes so fields hidden behind a reference are still detected.

opentau.configs.policies.warn_removed_policy_fields_from_dict(data: dict, config_path: str | Path) → None

Like warn_removed_policy_fields() but operating on an already-resolved config dict. config_path is used only for the warning message.

opentau.configs.policies.write_stripped_config_to_tempfile(data: dict) → Path

Strip deprecated/removed fields from data and write to a temp file.

data is treated as read-only (a defensive copy is made before stripping). Caller owns the returned path and must unlink it.