opentau.policies.pretrained

Base class for pre-trained policies in OpenTau.

This module defines the abstract base class PreTrainedPolicy which handles loading, saving, and basic interface requirements for all policy implementations in the OpenTau library. It integrates with Hugging Face Hub for model sharing and safetensors for efficient serialization.

Functions

is_norm_buffer_key(key)

Return True iff key names a Normalize/Unnormalize stat parameter.

Classes

PreTrainedPolicy(config, *inputs, **kwargs)

Base class for all policy models in OpenTau.

Exceptions

ProjectionRemapError

Raised when a per-(robot_type, control_mode) projection weight in a checkpoint cannot be reconciled with the policy's current group set (e.g. a per-group checkpoint loaded into a non-per-group policy, or a multi-group source whose group ordering is unknown).

class opentau.policies.pretrained.PreTrainedPolicy(config: PreTrainedConfig, *inputs, **kwargs)

Bases: Module, HubMixin, ABC

Base class for all policy models in OpenTau.

This class extends nn.Module and HubMixin to provide common functionality for policy models, including configuration management, model loading/saving, and abstract methods that all policies must implement.

config

The configuration instance for this policy.

__init__(config: PreTrainedConfig, *inputs, **kwargs)

Initializes the PreTrainedPolicy.

Parameters:

• config – The configuration object for the policy.

• *inputs – Variable length argument list.

• **kwargs – Arbitrary keyword arguments.

Raises:

ValueError – If config is not an instance of PreTrainedConfig.

config_class: None

The configuration class associated with this policy. Must be defined in subclasses.

abstract forward(batch: dict[str, Tensor]) → tuple[Tensor, dict | None]

Performs a forward pass of the policy.

Parameters:

batch – A dictionary of input tensors.

Returns:

A tuple containing:

• The loss tensor.

• An optional dictionary of metrics or auxiliary outputs. Apart from the loss, items should be logging-friendly native Python types.

Return type:

tuple[Tensor, dict | None]

classmethod from_pretrained(pretrained_name_or_path: str | Path, *, config: PreTrainedConfig | None = None, 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, strict: bool = False, **kwargs) → T

Loads a pretrained policy from a local path or the Hugging Face Hub.

The policy is set in evaluation mode by default using policy.eval() (dropout modules are deactivated). To train it, you should first set it back in training mode with policy.train().

Parameters:

• pretrained_name_or_path – The name or path of the pretrained model.

• config – Optional configuration object. If None, it will be loaded from the pretrained model.

• force_download – Whether to force download the model weights.

• resume_download – Whether to resume an interrupted download.

• proxies – Proxy configuration for downloading.

• token – Hugging Face token for authentication.

• cache_dir – Directory to cache downloaded files.

• local_files_only – Whether to only look for local files.

• revision – The specific model version to use (branch, tag, or commit hash).

• strict – Whether to strictly enforce matching keys in state_dict.

• **kwargs – Additional keyword arguments passed to the constructor.

Returns:

An instance of the loaded policy.

Return type:

T

Raises:

FileNotFoundError – If the model file is not found.

abstract get_optim_params() → dict

Returns the policy-specific parameters dict to be passed on to the optimizer.

Returns:

A dictionary of parameters to optimize.

Return type:

dict

maybe_compile_for_training() → None

Optionally torch.compile the heavy compute submodule for training.

Controlled by config.use_torch_compile. When enabled, compiles the inner flow-matching module (self.model on pi05 / pi07) in place via torch.nn.Module.compile(). The in-place form only swaps the submodule’s __call__ dispatch for a compiled one — it does not wrap the module in a torch._dynamo.OptimizedModule and does not prefix state_dict keys with _orig_mod.. So self.parameters() and the on-disk checkpoint layout are unchanged: the optimizer built afterwards sees the same params, and resume / from_pretrained stay compatible with both compiled and uncompiled checkpoints.

The policy wrapper’s own forward (tokenization, normalization, the .cpu().tolist() / numpy / string preprocessing) is intentionally left uncompiled — it is cheap relative to the transformer and is full of host-side ops that would only force graph breaks.

Note: this relies on the training forward calling the submodule via self.model(...) (__call__), not self.model.forward(...) — the latter bypasses the compiled dispatch installed by nn.Module.compile and would silently no-op the compile.

No-op unless config.use_torch_compile is True. The caller (train.py) is responsible for rejecting parameter-sharding backends (DeepSpeed ZeRO-3 / FSDP) before this runs.

name: None

The name of the policy. Must be defined in subclasses.

abstract reset()

Resets the policy state.

This method should be called whenever the environment is reset. It handles tasks like clearing caches or resetting internal states for stateful policies.

abstract select_action(batch: dict[str, Tensor]) → Tensor

Selects an action based on the input batch.

This method handles action selection during inference, including caching for stateful policies (e.g. RNNs, Transformers).

Parameters:

batch – A dictionary of observation tensors.

Returns:

The selected action(s).

Return type:

Tensor

supports_torch_compile: bool = False

Whether maybe_compile_for_training() may compile this policy’s self.model. Default False: opt-in per policy, because the in-place nn.Module.compile only takes effect if the policy’s training forward invokes the submodule via self.model(...) (__call__) rather than self.model.forward(...). Subclasses whose forward has been switched to self.model(...) (currently PI05Policy and PI07LowLevelPolicy) set this True. Leaving it False makes use_torch_compile=True a loud no-op on unwired policies instead of a silent compile-but-never-dispatch.

exception opentau.policies.pretrained.ProjectionRemapError

Bases: ValueError

Raised when a per-(robot_type, control_mode) projection weight in a checkpoint cannot be reconciled with the policy’s current group set (e.g. a per-group checkpoint loaded into a non-per-group policy, or a multi-group source whose group ordering is unknown). Subclasses ValueError so existing except ValueError handlers still catch it, but is distinct so the load path can re-raise it past the broad except Exception that otherwise swallows load errors into a warning.

opentau.policies.pretrained.is_norm_buffer_key(key: str) → bool

Return True iff key names a Normalize/Unnormalize stat parameter.