opentau.policies.normalize

Per-dataset Normalize / Unnormalize for VLA policies.

Each feature buffer has shape (D, *feat_shape) where D is the number of datasets the policy was trained on. forward accepts a per-sample dataset_index: torch.LongTensor of shape (B,) and gathers the right row per sample before broadcasting.

Functions

create_stats_buffers(features, norm_map[, ...])

Create per-feature stat buffers shaped (D, *feat_shape).

resolve_num_datasets(per_dataset_stats, ...)

Decide the leading dim for the stacked stats buffers.

warn_missing_keys(features, batch, mode)

Warns if expected features are missing from the batch.

Classes

Normalize(features, norm_map[, ...])

Normalizes data (e.g. observation.image) per-sample using the row of the stacked stat buffer selected by dataset_index.

Unnormalize(features, norm_map[, ...])

Inverse of Normalize: scales outputs back to the per-sample dataset range.
class opentau.policies.normalize.Normalize(features: dict[str, PolicyFeature], norm_map: dict[str, NormalizationMode], per_dataset_stats: list[dict[str, dict[str, Tensor | ndarray]]] | None = None, dataset_names: list[str] | None = None, num_datasets: int | None = None)[source]

Bases: Module

Normalizes data (e.g. observation.image) per-sample using the row of the stacked stat buffer selected by dataset_index.

__init__(features: dict[str, PolicyFeature], norm_map: dict[str, NormalizationMode], per_dataset_stats: list[dict[str, dict[str, Tensor | ndarray]]] | None = None, dataset_names: list[str] | None = None, num_datasets: int | None = None)[source]

Initializes the Normalize module.

Parameters:

  • features – Mapping feature name -> PolicyFeature.

  • norm_map – Mapping FeatureType -> NormalizationMode.

  • per_dataset_stats – Ordered list of per-dataset stat dicts. None during checkpoint-load construction (the safetensors load overwrites the buffers afterwards); in that case num_datasets must be supplied so the inf-init buffers have the right leading dim.

  • dataset_names – Ordered list of dataset name strings, parallel to per_dataset_stats. Stored as a Python attribute, not a buffer (strings aren’t tensors). Persisted into config.json via PreTrainedConfig.dataset_names.

  • num_datasets – Required when per_dataset_stats is None.

forward(batch: dict[str, Tensor], dataset_index: Tensor) → dict[str, Tensor][source]

Normalizes the batch data per-sample.

Parameters:

  • batch – Dictionary containing the data to normalize.

  • dataset_index – LongTensor of shape (B,) giving the source dataset row for each sample.

Returns:

The normalized batch (shallow-copied; inputs are not mutated).

class opentau.policies.normalize.Unnormalize(features: dict[str, PolicyFeature], norm_map: dict[str, NormalizationMode], per_dataset_stats: list[dict[str, dict[str, Tensor | ndarray]]] | None = None, dataset_names: list[str] | None = None, num_datasets: int | None = None)[source]

Bases: Module

Inverse of Normalize: scales outputs back to the per-sample dataset range.

__init__(features: dict[str, PolicyFeature], norm_map: dict[str, NormalizationMode], per_dataset_stats: list[dict[str, dict[str, Tensor | ndarray]]] | None = None, dataset_names: list[str] | None = None, num_datasets: int | None = None)[source]

See Normalize.__init__ for argument semantics.

forward(batch: dict[str, Tensor], dataset_index: Tensor) → dict[str, Tensor][source]

Unnormalizes the batch data per-sample.

Parameters:

  • batch – Dictionary containing the data to unnormalize.

  • dataset_index – LongTensor of shape (B,).

opentau.policies.normalize.create_stats_buffers(features: dict[str, PolicyFeature], norm_map: dict[str, NormalizationMode], per_dataset_stats: list[dict[str, dict[str, Tensor | ndarray]]] | None = None, num_datasets: int | None = None) → dict[str, dict[str, ParameterDict]][source]

Create per-feature stat buffers shaped (D, *feat_shape).

Parameters:

  • features – Mapping feature name -> PolicyFeature.

  • norm_map – Mapping FeatureType -> NormalizationMode.

  • per_dataset_stats – Ordered list (one entry per dataset) of dicts shaped {feature_name: {"mean": ..., "std": ...}} or {... "min": ..., "max": ...}. When None the buffers are initialised to +inf so the assertion in Normalize.forward fires until a checkpoint load or _inject_stats overwrites them.

  • num_datasets – Required when per_dataset_stats is None; ignored otherwise. Sets the leading dim of the inf-init buffers.

Returns:

param, “std”: param}) etc.

Return type:

dict feature_name -> nn.ParameterDict({“mean”

opentau.policies.normalize.resolve_num_datasets(per_dataset_stats: list | None, dataset_names: list | None, config) → int[source]

Decide the leading dim for the stacked stats buffers.

Preference order:

  1. len(per_dataset_stats) if provided.

  2. len(dataset_names) if provided.

  3. len(config.dataset_names) if the policy config remembers a trained-on list (checkpoint-load construction path).

  4. 1 — legacy single-dataset default.

opentau.policies.normalize.warn_missing_keys(features: dict[str, PolicyFeature], batch: dict[str, Tensor], mode: str) → None[source]

Warns if expected features are missing from the batch.

Parameters:

  • features – Dictionary of expected policy features.

  • batch – Dictionary containing the data batch.

  • mode – The operation mode (e.g., “normalization” or “unnormalization”) for the warning message.