opentau.policies.outlier_utils

Pure, collective-free detection of normalized state/action outliers.

Shared by the pi07 and pi07_paligemma low-level policies. Detection runs inside forward so it MUST stay rule-5 safe (see CLAUDE.md rule 5): it only reads the batch and returns records — no logging, no cross-step dedup, and no collectives — keeping the forward graph (and the collective counts under FSDP / ZeRO) identical across ranks regardless of what local data a rank holds.

The cross-rank merge, cross-step dedup, and the actual logging.warning live in the training loop (log_outlier_records_distributed in opentau/scripts/train.py), which runs on rank 0 after gathering offenders from every rank — so the warning always reaches wandb regardless of which rank held the offending sample.

Functions

detect_state_action_outliers(batch, threshold)

Return the worst (source, key, dim) outlier records in batch.

Classes

OutlierRecord

One worst (source, key, dim) offender found in a local batch.

class opentau.policies.outlier_utils.OutlierRecord

Bases: TypedDict

One worst (source, key, dim) offender found in a local batch.

value is the largest |normalized value| seen for this (source, key, dim) in the batch; source / episode / frame are the provenance of that worst sample (None when the batch lacks the field). source is the mixture-level per-entry dataset name (dataset_repo_id, unique across colliding mixture entries) when the batch carries one, else the sample-level source feature-mapping key. Only plain Python scalars are stored so the list survives the gather_object round-trip used to ship records to rank 0.

dim: int

episode: object

frame: object

key: str

source: object

value: float

opentau.policies.outlier_utils.detect_state_action_outliers(batch: dict[str, Tensor], threshold: float | None) → list[OutlierRecord]

Return the worst (source, key, dim) outlier records in batch.

A normalized value far from unit scale almost always means bad normalization stats (e.g. near-zero std on a constant dim) or corrupt data. This finds the offending dims so the training loop can warn about them, recording the dataset identity / episode_index / frame_index (when present in the batch) to trace a poorly-normalized dim back to the dataset/frame to inspect. The dataset identity is dataset_repo_id (the mixture-level per-entry name, unique even for colliding entries that share a repo_id and control_mode) when present, else the sample-level source key.

Pure and collective-free: it reads batch without mutating it, fires no collective, and does NO cross-step dedup or logging — so the forward graph, and therefore the collective counts under FSDP / ZeRO, stay identical across ranks regardless of what any rank’s local data contains (CLAUDE.md rule 5). The returned records are merged across ranks, deduped, and logged on rank 0 by log_outlier_records_distributed in the training loop.

Parameters:

• batch – Training batch after input/target normalization. Reads state (B, [T,] D) and actions (B, chunk, D); the zero-padded tail dims never trigger. Timesteps the model does not attend to (padded history via obs_history_is_pad / padded action steps via action_is_pad) are excluded so a masked-out frame can’t trip the check — the current state frame is always kept.

• threshold – Absolute-value ceiling. None or <= 0 disables the check entirely (returns [] with no device sync).

Returns:

One OutlierRecord per offending (source, key, dim) in this batch (the worst |value| per tuple). Empty when disabled or clean.