opentau.policies.utils

Utility functions for policy implementations in OpenTau.

This module provides helper functions for managing data queues, inspecting model properties (device, dtype), determining output shapes, and logging model loading information.

Functions

assert_gemma3_input_resolution(...)	Fail fast when a Gemma3-family policy would run at a native resolution.
ce_per_sample(masked_ce, valid_mask)	Per-sample numerator/denominator for a masked token cross-entropy.
flow_matching_masked_mse(u_t, v_t, *, ...[, ...])	Masked MSE for flow-matching velocity-field training.
get_device_from_parameters(module)	Get a module's device by checking one of its parameters.
get_dtype_from_parameters(module)	Get a module's parameter dtype by checking one of its parameters.
get_output_shape(module, input_shape)	Calculates the output shape of a PyTorch module given an input shape.
log_model_loading_keys(missing_keys, ...)	Log missing and unexpected keys when loading a model.
make_action_dim_mask(real_action_dim, ...)	Per-sample bool mask over action dims; True for real dims, False for zero-pad.
populate_queues(queues, batch[, exclude_keys])	Populates queues with batch data.

Classes

PerSampleLoss(sum, count)

Per-sample decomposition of a masked loss, with the batch dim kept.

class opentau.policies.utils.PerSampleLoss(sum: Tensor, count: Tensor)

Bases: object

Per-sample decomposition of a masked loss, with the batch dim kept.

sum and count are both (B,) and hold, for each sample, the summed unmasked loss and the number of unmasked slots that fed it. The masked mean for any group of samples is Σsum / Σcount — carrying the (numerator, denominator) pair rather than a per-sample mean is what lets a caller (e.g. the validation loop) regroup samples by provenance and recover an exact masked mean per group. Averaging per-sample means instead would double-normalize and weight a 1-slot sample the same as a 200-slot one.

__init__(sum: Tensor, count: Tensor) → None

count: Tensor

sum: Tensor

opentau.policies.utils.assert_gemma3_input_resolution(input_image_size: tuple[int, int] | None, tower_image_size: int) → None

Fail fast when a Gemma3-family policy would run at a native resolution.

Gemma3MultiModalProjector hard-codes a square patches_per_image = image_size // patch_size reshape + avg-pool, so a non-default patch grid crashes deep inside the projector with an unrelated-looking reshape error. The pi06/pi07 constructors call this to surface the real diagnosis instead. (The PaliGemma-family policies support native resolutions.)

Parameters:

• input_image_size – (H, W) the vision tower will receive (PreTrainedConfig.input_image_size), or None when underivable (no check possible).

• tower_image_size – The Gemma 3 SigLIP config’s square image_size.

Raises:

ValueError – When input_image_size is known and differs from the tower’s square resolution.

opentau.policies.utils.ce_per_sample(masked_ce: Tensor, valid_mask: Tensor) → PerSampleLoss

Per-sample numerator/denominator for a masked token cross-entropy.

Policies compute their CE as F.cross_entropy(..., reduction="none") reshaped to (B, S) and zeroed at pad positions, then reduce it with .mean() to a scalar. This helper takes that same pad-zeroed (B, S) tensor plus the per-token validity mask and returns the per-sample (Σ over valid tokens, #valid tokens), so a caller can pool CE per provenance group as Σsum / Σcount — the mean cross-entropy per valid token. Multiple CE components (e.g. discrete-action + response) pool by adding their PerSampleLoss objects.

Note this normalizes by valid token count, unlike the legacy scalar .mean() which divides by the full B * S (pad slots included); the per-group breakdown is therefore over valid tokens only.

Parameters:

• masked_ce – (B, S) cross-entropy already zeroed at pad positions.

• valid_mask – (B, S) bool, True at non-pad (scored) tokens.

Returns:

PerSampleLoss whose sum and count are (B,).

opentau.policies.utils.flow_matching_masked_mse(u_t: Tensor, v_t: Tensor, *, max_action_dim: int, prefix_mask: Tensor | None = None, actions_is_pad: Tensor | None = None, real_action_dim: Tensor | None = None, return_per_sample: bool = False) → Tensor | tuple[Tensor, PerSampleLoss]

Masked MSE for flow-matching velocity-field training.

Shared across pi05, pi05_mem, pi06, pi07 (low_level), and pi07_paligemma (low_level). Builds a (B, chunk_size, max_action_dim) mask that AND-s together up to three conditions and reduces F.mse_loss(u_t, v_t) over the unmasked slots:

1. Frozen-prefix (RTI delay): ~prefix_mask — False where the model isn’t asked to predict (the action prefix is the actually executed action from a previous inference, frozen as ground truth). Pass None to disable (non-RTI policies); the helper builds an all-False prefix mask internally so every step is supervised.

2. Per-timestep chunk padding: ~actions_is_pad — False where the action chunk extends past episode end. Pass None to skip. Also covers VQA-style items (actions_is_pad all-True ⇒ loss = 0).

3. Per-sample real action dim: built from real_action_dim via make_action_dim_mask(). False on the zero-pad tail dims of each sample. Pass None to score all max_action_dim columns.

Parameters:

• u_t – Target velocity field, shape (B, chunk_size, D) (D ≥ max_action_dim).

• v_t – Predicted velocity field, same shape as u_t.

• max_action_dim – Number of leading action dims to score against; trailing dims are dropped before reduction. Keyword-only.

• prefix_mask – Optional bool (B, chunk_size) — True where the step is frozen (RTI delay). None ⇒ all-False (non-RTI behavior).

• actions_is_pad – Optional bool (B, chunk_size) — True where the action chunk is padded (no real action target). None ⇒ all-False.

• real_action_dim – Optional long (B,) — real (pre-pad) action dim per sample. None ⇒ all-True (every dim is real).

• return_per_sample – When True, additionally return a PerSampleLoss holding the per-sample (Σ over masked slots, #masked slots) so the caller can regroup the loss by provenance. The scalar is computed exactly as in the default path (bit-identical), so toggling this flag never perturbs the training reduction.

Returns:

Scalar tensor (masked mean of (u_t - v_t)**2 over the unmasked slots) when return_per_sample is False; otherwise (scalar, PerSampleLoss) where the per-sample sum/count are over the same masked slots (so each sample’s mean is sum / count).

opentau.policies.utils.get_device_from_parameters(module: Module) → device

Get a module’s device by checking one of its parameters.

Note

Assumes that all parameters have the same device.

Parameters:

module – The PyTorch module to inspect.

Returns:

The device of the module’s parameters.

Return type:

torch.device

opentau.policies.utils.get_dtype_from_parameters(module: Module) → dtype

Get a module’s parameter dtype by checking one of its parameters.

Note

Assumes that all parameters have the same dtype.

Parameters:

module – The PyTorch module to inspect.

Returns:

The data type of the module’s parameters.

Return type:

torch.dtype

opentau.policies.utils.get_output_shape(module: Module, input_shape: tuple) → tuple

Calculates the output shape of a PyTorch module given an input shape.

Parameters:

• module – A PyTorch module.

• input_shape – A tuple representing the input shape, e.g., (batch_size, channels, height, width).

Returns:

The output shape of the module.

Return type:

tuple

opentau.policies.utils.log_model_loading_keys(missing_keys: list[str], unexpected_keys: list[str]) → None

Log missing and unexpected keys when loading a model.

Parameters:

• missing_keys – Keys that were expected but not found.

• unexpected_keys – Keys that were found but not expected.

opentau.policies.utils.make_action_dim_mask(real_action_dim: Tensor | None, max_action_dim: int, batch_size: int, device: device) → Tensor

Per-sample bool mask over action dims; True for real dims, False for zero-pad.

Heterogeneous datasets are zero-padded to max_action_dim along the last action axis to keep batches rectangular, but the flow-matching MSE on the velocity field should only score real dims for each sample. This helper builds the per-dim mask that callers AND into their existing per-timestep mask before reducing.

Parameters:

• real_action_dim – Optional (B,) long tensor of the real (pre-pad) action dimensionality for each sample (the batch key emitted by LeRobotDataset._to_standard_data_format). When None, the returned mask is all-True so the dim-mask AND in the caller’s reduction is a no-op (pi0 additionally harmonized its .mean() to sum / mask.sum() in this PR — see the PR body’s “pi0 loss-magnitude shift” note; the dim-mask itself is still a no-op when real_action_dim is None).

• max_action_dim – The padded action dim (last-axis length of actions).

• batch_size – Used to construct the all-True fallback shape; when real_action_dim is provided, must match real_action_dim.shape[0].

• device – Output device.

Returns:

(batch_size, max_action_dim) bool tensor.

opentau.policies.utils.populate_queues(queues: dict[str, deque], batch: dict[str, Tensor], exclude_keys: list[str] | None = None) → dict[str, deque]

Populates queues with batch data.

If a queue is not full (e.g. at the start of an episode), it is filled by repeating the first observation. Otherwise, the latest observation is appended.

Parameters:

• queues – A dictionary of deques to be populated.

• batch – A dictionary containing the data to add to the queues.

• exclude_keys – A list of keys to exclude from population. Defaults to None.

Returns:

The updated dictionary of queues.

Return type:

dict[str, deque]