#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Default configuration classes for datasets, evaluation, and logging.
This module provides default configuration classes for:
- Dataset configuration and dataset mixtures
- Weights & Biases (wandb) logging configuration
- Evaluation settings and parameters
"""
import warnings
from dataclasses import dataclass, field
from typing import Literal
import draccus
import numpy as np
from draccus.parsers.encoding import encode_dataclass
from opentau import (
policies, # noqa: F401
)
from opentau.datasets.transforms import ImageTransformsConfig
from opentau.datasets.video_utils import get_safe_default_codec
# --- Custom NumPy encoder registration ---
# For decoding from cmd/yaml
draccus.decode.register(np.ndarray, np.asarray)
# For encoding to yaml
draccus.encode.register(np.ndarray, lambda x: x.tolist())
# Registry keys some config entry RESOLVES through (its effective key) this
# process. Distinguishes a genuine mixture-entry collision (warn: an earlier
# entry's registry consumers would silently see the new columns) from the
# silent-by-design overwrites: a single entry overriding a built-in default,
# and the plain repo_id back-compat slot in the joint/ee pattern.
_CONFIG_EFFECTIVE_MAPPING_KEYS: set[str] = set()
[docs]
@dataclass
class DatasetConfig:
"""Configuration for a dataset.
You may provide a list of datasets here. `train.py` creates them all and
concatenates them. Note: only data keys common between the datasets are kept.
Each dataset gets an additional transform that inserts the "dataset_index"
into the returned item. The index mapping is made according to the order in
which the datasets are provided.
Args:
repo_id: HuggingFace repository ID for the dataset. Exactly one of
`repo_id` or `vqa` must be set.
vqa: VQA dataset identifier. Exactly one of `repo_id` or
`vqa` must be set.
root: Root directory where the dataset will be stored (e.g. 'dataset/path').
Defaults to None.
episodes: List of episode indices to use from the dataset. If None, all
episodes are used. Defaults to None.
excluded_episodes: List of episode indices to drop from this dataset.
Takes precedence over `episodes`: an index present in both is
excluded. If None (default), no episodes are excluded. Note: on
legacy v2.0 datasets (no per-episode stats) the listed episodes are
dropped from training, but normalization stats stay the global
(all-episode) aggregate — only v2.1+ can recompute them.
image_transforms: Configuration for image transformations. Defaults to
ImageTransformsConfig().
revision: Git revision of the dataset repository to use. Defaults to None.
use_imagenet_stats: Whether to use ImageNet statistics for normalization.
Defaults to True.
video_backend: Video codec backend to use. Defaults to a safe default codec.
stats: Dictionary of statistics for normalization, keyed by feature name.
Each value is a dictionary with 'mean' and 'std' arrays. Defaults to None.
data_features_name_mapping: Optional mapping from standard feature
names (``camera0``/``camera1``/..., ``state``, ``actions``,
``prompt``, ``response``, ``mistake``, ``success``) to this
dataset's own column names. The ``mistake`` and ``success`` roles
feed the optional ``mistake`` metadata key: map a
mistake-polarity column (True/1 = something went wrong) to
``mistake``, or a success-polarity column (e.g. DROID's
``is_episode_successful``) to ``success`` — polarity is expressed
by which role you map, so no inversion flag is needed. When both
resolve, ``mistake`` wins (it is segment-grained). ``mistake``
must name a per-frame column; for an episode-level outcome map
``success`` instead, which resolves from a per-frame column or a
per-episode key in the episodes metadata and also drives the
value-function return bins.
Two mixture entries may share a ``repo_id`` and ``control_mode``
while declaring different mappings (e.g. two camera views of one
repo): each dataset instance resolves its own entry's mapping.
The global registry keeps only the last registration for
annotation-script consumers — a warning fires when entries
disagree. Defaults to None.
robot_type: Optional override for the dataset's ``robot_type`` metadata
field. When provided (including the empty string), takes precedence
over the value loaded from ``meta/info.json``. ``None`` (default)
leaves the loaded value untouched.
control_mode: Optional override for the dataset's ``control_mode``
metadata field. When provided (including the empty string), takes
precedence over the value loaded from ``meta/info.json``. ``None``
(default) leaves the loaded value untouched.
tolerance_s: Optional per-dataset override for the timestamp-sync
tolerance (in seconds) passed to ``LeRobotDataset``'s load-time
``check_timestamps_sync`` call. ``None`` (default) inherits the
mixture-wide ``DatasetMixtureConfig.tolerance_s`` value. Set this
to a larger value (e.g. ``1e-3``) when a single dataset in the
mixture has slightly off-fps timestamps but you don't want to
loosen the check for the others. Must be ``>= 0`` when set.
This also applies to frames decoded from video files: the same
value is used as the per-frame match tolerance in
``query_video_frames_*``, so loosening it widens the
video-frame match window as well.
skip_timestamp_check: Optional per-dataset override that bypasses the
load-time ``check_timestamps_sync`` call entirely. ``None``
(default) inherits the mixture-wide
``DatasetMixtureConfig.skip_timestamp_check``. ``True`` skips the
check (a warning is logged); ``False`` forces the check on for this
dataset even if the mixture default is ``True``. Does not affect
the record-time check inside ``add_episode``.
prompt_substitutions: Optional mapping from an on-disk task string
(exact match against ``meta/tasks.*``) to a non-empty list of
non-empty substitute prompts. At fetch time a matching sample's task is
ALWAYS replaced by a uniform random draw from its list (include
the original in the list if it should still appear); unmapped
tasks pass through unchanged. Applies to the train split only
unless ``DatasetMixtureConfig.val_enable_prompt_substitution`` is
set. Keys that match no on-disk task string raise at dataset
init. ``response``/memory CE targets are NOT rewritten, so
substitutes must be semantic paraphrases of the original task.
Only settable for LeRobot datasets (``repo_id``), not VQA
entries. Like every per-dataset field it has no CLI override
path — set it in the JSON config; an external fragment can be
inlined with ``{"$ref": "path/to/fragment.json"}`` (see
``opentau/configs/refs.py``). Draws use the default torch RNG
(see the ``DatasetMixtureConfig`` note). Defaults to None.
Raises:
ValueError: If both or neither of ``repo_id`` and ``vqa`` are set, if
``tolerance_s`` or ``val_split_ratio`` is out of range, or if
``prompt_substitutions`` is set on a VQA entry, maps a task
literally named ``"$ref"``, has a non-string key, or contains an
empty substitute list or empty/non-string entries.
"""
repo_id: str | None = None
vqa: str | None = None
# Root directory where the dataset will be stored (e.g. 'dataset/path').
root: str | None = None
episodes: list[int] | None = None
excluded_episodes: list[int] | None = None
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
revision: str | None = None
use_imagenet_stats: bool = True
video_backend: str = field(default_factory=get_safe_default_codec)
stats: dict[str, dict[str, np.ndarray]] | None = None
# optional standard data format mapping for the dataset if mapping is not already in standard_data_format_mapping.py
data_features_name_mapping: dict[str, str] | None = None
# Optional per-dataset prompt substitution. Maps an on-disk task string
# (exact match against meta/tasks.*) to a non-empty list of substitute
# prompts; at fetch time a matching sample's task is ALWAYS replaced by a
# uniform random draw from the list (include the original in the list if
# it should still appear). Unmapped tasks pass through unchanged. Train
# split only unless `DatasetMixtureConfig.val_enable_prompt_substitution`
# is set. Config-file-only (per-dataset fields have no CLI override path);
# an external fragment can be inlined with `{"$ref": "fragment.json"}`.
# Note: `response`/memory CE targets are NOT rewritten — substitutes must
# be semantic paraphrases of the original task.
prompt_substitutions: dict[str, list[str]] | None = None
# Optional overrides for the metadata fields read from `meta/info.json`.
# `None` means "do not override". Any string value (including "") is
# written through to `dataset.meta.info[...]` after the dataset is built.
robot_type: str | None = None
control_mode: str | None = None
# Per-dataset overrides for the load-time timestamp-sync check. `None`
# inherits from `DatasetMixtureConfig.{tolerance_s, skip_timestamp_check}`.
# A non-None value here wins over the mixture default for this dataset
# only — useful when one dataset in a mixture has off-fps timestamps.
tolerance_s: float | None = None
skip_timestamp_check: bool | None = None
# Per-dataset override for the train/validation split ratio. `None`
# inherits from `DatasetMixtureConfig.val_split_ratio`. A non-None value
# (including `0.0`, which opts this dataset out of validation) wins over the
# mixture default for this dataset only — useful when one dataset in a
# mixture wants a different validation fraction than the rest. Must be in
# `[0, 1]` when set. Only consulted when `TrainPipelineConfig.val_freq > 0`.
val_split_ratio: float | None = None
def __post_init__(self):
"""Validate dataset configuration and register custom mappings if provided."""
if (self.repo_id is None) == (self.vqa is None):
raise ValueError("Exactly one of `repo_id` or `vqa` for Dataset config should be set.")
if self.tolerance_s is not None and self.tolerance_s < 0:
raise ValueError(
f"`DatasetConfig.tolerance_s` must be >= 0 (or None to inherit), "
f"got {self.tolerance_s} for {self.repo_id or self.vqa}."
)
if self.val_split_ratio is not None and not (0.0 <= self.val_split_ratio <= 1.0):
raise ValueError(
f"`DatasetConfig.val_split_ratio` must be in [0, 1] (or None to inherit), "
f"got {self.val_split_ratio} for {self.repo_id or self.vqa}."
)
if self.prompt_substitutions is not None:
if self.vqa is not None:
raise ValueError(
f"`DatasetConfig.prompt_substitutions` only applies to LeRobot datasets "
f"(`repo_id`); VQA dataset {self.vqa!r} builds its own prompt."
)
for task, subs in self.prompt_substitutions.items():
if not isinstance(task, str):
raise ValueError(
f"`prompt_substitutions` keys must be on-disk task strings, got "
f"{task!r} for {self.repo_id}."
)
if task == "$ref":
raise ValueError(
"`prompt_substitutions` cannot map a task literally named '$ref': the "
"config loader treats any JSON object containing a '$ref' key as a file "
"include directive (see opentau/configs/refs.py)."
)
if not isinstance(subs, list) or len(subs) == 0:
raise ValueError(
f"`prompt_substitutions[{task!r}]` must be a non-empty list of strings, "
f"got {subs!r} for {self.repo_id}."
)
if not all(isinstance(s, str) and s for s in subs):
raise ValueError(
f"`prompt_substitutions[{task!r}]` contains empty or non-string entries: "
f"{subs!r} for {self.repo_id}. With always-replace semantics an empty "
f"substitute would silently train matching samples on an empty prompt."
)
# If data_features_name_mapping is provided, upsert it into the global
# DATA_FEATURES_NAME_MAPPING. Register under the plain repo_id (back-compat
# fallback, last-wins) AND, when this entry carries a real control mode,
# under the composite `repo_id::control_mode` key. The composite key lets
# two entries that share a repo_id but use different action columns
# (e.g. control_mode="joint" -> action_joint vs control_mode="ee" ->
# action_ee) coexist instead of the second silently clobbering the first.
# Dataset instances resolve their own entry's mapping per-instance
# (BaseDataset._get_name_map), so a same-key clobber here no longer
# affects training reads — but the registry stays the source for
# annotation scripts and direct constructions, so warn when two
# entries disagree on the same key.
if self.data_features_name_mapping is not None:
from opentau.datasets.standard_data_format_mapping import (
DATA_FEATURES_NAME_MAPPING,
feature_mapping_key,
)
keys = []
effective = None
if self.repo_id is not None:
effective = feature_mapping_key(self.repo_id, self.control_mode)
keys = list(dict.fromkeys([self.repo_id, effective]))
for key in keys:
# Warn only when overwriting a key some earlier config entry
# RESOLVES through (its effective key) with a different
# mapping — that entry's registry consumers would silently
# see this entry's columns. Both the effective write and the
# plain-repo_id back-compat write are checked (a plain-key
# entry's effective registration can be clobbered by a later
# composite-key entry's back-compat write). Overwriting an
# un-tracked key — a built-in default, or the back-compat
# slot in the joint/ee pattern — stays silent by design.
previous = DATA_FEATURES_NAME_MAPPING.get(key)
if (
key in _CONFIG_EFFECTIVE_MAPPING_KEYS
and previous is not None
and previous != self.data_features_name_mapping
):
warnings.warn(
f"data_features_name_mapping: registry key {key!r} is being overwritten "
"with a different mapping (another mixture entry resolves its mapping "
"via this key). Each dataset instance still uses its own entry's "
"mapping, but registry consumers (annotation scripts, direct "
"constructions) will see only the last one registered.",
stacklevel=2,
)
DATA_FEATURES_NAME_MAPPING[key] = self.data_features_name_mapping
if effective is not None:
_CONFIG_EFFECTIVE_MAPPING_KEYS.add(effective)
[docs]
@dataclass
class DatasetMixtureConfig:
r"""Configuration for a mixture of multiple datasets.
This configuration allows combining multiple datasets with specified weights
for training. The datasets are sampled according to their weights during
training, and features are resampled to a common action frequency.
Args:
datasets: List of dataset configs to be used in the mixture.
weights: Optional list of weights for each dataset in the mixture. Must be the
same length as `datasets` when provided. If None, weights are inferred
from dataset lengths. Defaults to None.
action_freq: Frequency at which actions from the dataset mixture are
resampled, in Hz. ``None`` (default) disables resampling — each
dataset is sampled at its native fps, so a single batch can mix
samples from sources running at different rates (predicting
``chunk_size`` consecutive native frames per sample). Set a
positive float to resample every dataset in the mixture to that
common rate via nearest-neighbor frame selection. When using
``None``, prefer also setting ``emit_fps=True`` so the policy
can condition on the per-sample rate.
image_resample_strategy: Resample strategy for image features. Must be
one of 'linear' or 'nearest'. Defaults to 'nearest'.
vector_resample_strategy: Resample strategy for non-image features, such
as action or state. Must be one of 'linear' or 'nearest'.
Defaults to 'nearest'.
val_split_ratio: Mixture-wide default fraction of each dataset reserved
for the validation split (only used when
``TrainPipelineConfig.val_freq > 0``). A per-dataset
``DatasetConfig.val_split_ratio`` overrides this value for that
dataset; ``None`` there inherits this mixture default. Must be in
``[0, 1]``. Defaults to 0.05.
n_obs_history: Number of historical observation steps to include. When
set to ``T``, each camera returns shape ``(T, C, H, W)`` and state
returns shape ``(T, max_state_dim)``. When ``None``, the default
single-step behavior is preserved with rank-3 camera tensors
``(C, H, W)`` and rank-1 state tensors ``(max_state_dim,)``.
Note that ``n_obs_history=1`` produces rank-4 camera tensors
``(1, C, H, W)`` with a leading singleton dimension, while state
collapses to rank-1 ``(max_state_dim,)`` (the length-1
delta-timestamps query is squeezed like the ``None`` case), so
downstream consumers must rank-normalize state themselves.
Defaults to ``None``. The temporal stride
between sampled observations is read from the policy config's
``history_interval`` attribute (defaults to 1 when the policy
doesn't define one), so observations are sampled at timesteps
:math:`t - (T-1)k,\; t - (T-2)k,\; \ldots,\; t`.
history_state_drop_prob: Probability of dropping the observation
*history* during a single ``__getitem__`` call. When it fires, the
historical steps are masked via ``obs_history_is_pad`` (set all True)
and the historical camera frames are zeroed; the current step —
current ``observation.state`` and current camera frame — is kept.
``state`` is deliberately NOT zeroed here: it is MEAN_STD-normalized
downstream, so the dropped history is zeroed *after* normalization
inside the policy (zeroing a raw state pre-normalization would map
0 to ``-mean/std``, an out-of-distribution extreme). Must be in
``[0, 1]``. Defaults to 0.3.
subgoal_drop_prob: Probability of dropping all subgoal images during a
single ``__getitem__`` call. Must be in ``[0, 1]``. Defaults to 0.75.
subgoal_end_of_segment_prob: Probability of sampling the subgoal frame
at the end of the current segment (vs. uniformly in the next 4s of
wall-clock time). Must be in ``[0, 1]``. Defaults to 0.25.
response_drop_prob: Probability of dropping the ``response`` (subtask
text) during a single ``__getitem__`` call. Only rolled when
subgoals are not dropped. Must be in ``[0, 1]``. Defaults to 0.3.
metadata_drop_all_prob: Probability of dropping ``speed``, ``mistake``,
``quality``, ``robot_type``, and ``control_mode`` together during a
single ``__getitem__`` call. Must be in ``[0, 1]``. Defaults to 0.15.
metadata_drop_each_prob: Per-field independent drop probability for
``speed``, ``mistake``, ``quality``, ``robot_type``, and
``control_mode``. Only rolled when ``metadata_drop_all_prob`` did not
fire. Must be in ``[0, 1]``.
Defaults to 0.05.
val_enable_optional_key_dropout: Whether to apply the five
``*_drop_prob`` rolls above to the validation split. Defaults to
``False`` — validation evaluates on un-masked samples so metrics
aren't polluted by training-time augmentation. Subgoal *frame*
sampling (end-of-segment vs. uniform in the next 4s) stays active
either way; only the masking logic is gated.
val_enable_prompt_substitution: Whether the per-dataset
``DatasetConfig.prompt_substitutions`` swaps also fire on the
validation split. Defaults to ``False`` — validation evaluates on
the on-disk prompts. Flip to ``True`` when val loss should match
the training prompt distribution (with always-replace semantics
the original prompt never appears in training unless listed among
its own substitutes).
require_non_empty_robot_type: If True, every dataset in the mixture
must have a non-empty ``robot_type`` after the optional
``DatasetConfig.robot_type`` override has been applied. Defaults to
``False`` (empty / missing values are allowed).
require_non_empty_control_mode: If True, every dataset in the mixture
must have a non-empty ``control_mode`` after the optional
``DatasetConfig.control_mode`` override has been applied. Defaults
to ``False`` (empty / missing values are allowed).
emit_fps: Whether ``__getitem__`` returns the *effective*
per-sample frame rate (``action_freq`` if set, else the
dataset's native ``meta.fps``) as the ``fps`` metadata key
(``torch.long`` scalar, paired with ``fps_is_pad=False``).
Default ``False`` — fps conditioning is an opt-in feature so
pre-PR checkpoints resume without the policy's metadata
prefix gaining an unfamiliar ``FPS:`` segment. Flip to
``True`` for new training runs that want per-sample
frame-rate conditioning (especially heterogeneous-frequency
mixtures where ``action_freq=None`` lets each dataset run at
its native rate). Unlike the other metadata fields, ``fps``
is **not** rolled by ``metadata_drop_*_prob`` — it's an
intrinsic property of the chunk, not a noisy label, so it
is always present (non-pad) for LeRobot samples when
``emit_fps=True``. VQA samples (no temporal axis) emit
``fps=0, fps_is_pad=True`` regardless so heterogeneous
VLA + VQA batches stay schema-aligned.
tolerance_s: Mixture-wide default tolerance (in seconds) for the
load-time ``check_timestamps_sync`` call inside
``LeRobotDataset.__init__``. Each dataset's frame-to-frame
timestamp spacing must lie within ``1/fps +/- tolerance_s`` or the
check raises. Defaults to ``1e-4``. A per-dataset
``DatasetConfig.tolerance_s`` overrides this value when set. Must
be ``>= 0``. This also applies to frames decoded from video
files: the same value is used as the per-frame match tolerance
in ``query_video_frames_*``, so loosening it widens the
video-frame match window as well.
skip_timestamp_check: If True, bypass the load-time
``check_timestamps_sync`` call for every dataset in the mixture
(a warning is logged per dataset). Useful as a debug knob when
the timing data is known-bad but you still want the mixture to
load. Defaults to ``False``. A per-dataset
``DatasetConfig.skip_timestamp_check`` overrides this value when
set. Does not affect the record-time check inside ``add_episode``.
Note:
Dropout rolls use the default torch RNG. PyTorch DataLoader workers
auto-seed each process's torch RNG from the base seed + worker id, so
workers sample independently. For reproducibility the caller should
seed via ``torch.manual_seed(...)`` in the main process before
constructing the DataLoader.
Raises:
ValueError: If `weights` is provided and its length doesn't match
`datasets`, if `action_freq` is not None and not positive, if
resample strategies are invalid, or if any drop probability is
outside ``[0, 1]``.
"""
# List of dataset configs to be used in the mixture.
datasets: list[DatasetConfig] = field(default_factory=list)
# Optional list of weights for each dataset in the mixture.
# Must be the same length as `datasets` when provided.
weights: list[float] | None = None
# Frequency at which the actions from dataset mixture are resampled, in Hz.
# ``None`` disables resampling — each dataset is sampled at its native fps.
action_freq: float | None = None
# Resample strategy for image features
image_resample_strategy: str = "nearest"
# Resample strategy for non-image features, such as action or state
vector_resample_strategy: str = "nearest"
# Mixture-wide default ratio of each dataset to be used for validation.
# If `val_freq` is set to 0, a validation dataset will not be created and this value will be ignored.
# A per-dataset `DatasetConfig.val_split_ratio` overrides this for that
# dataset (`None` there inherits this value).
# Defaults to 0.05.
val_split_ratio: float = 0.05
# Number of historical observation steps. None preserves default single-step behavior.
n_obs_history: int | None = None
# Training-time dropout probabilities for optional sample keys.
history_state_drop_prob: float = 0.3
subgoal_drop_prob: float = 0.75
subgoal_end_of_segment_prob: float = 0.25
response_drop_prob: float = 0.3
metadata_drop_all_prob: float = 0.15
metadata_drop_each_prob: float = 0.05
# Whether the above dropout rolls also fire on the validation split.
# Default keeps validation deterministic-ish (no masking); subgoal frame
# selection stays random either way.
val_enable_optional_key_dropout: bool = False
# Whether per-dataset `DatasetConfig.prompt_substitutions` swaps also fire
# on the validation split. Default keeps validation on the on-disk prompts.
val_enable_prompt_substitution: bool = False
# When True, require every dataset in the mixture to have a non-empty
# robot_type / control_mode after `DatasetConfig.{robot_type,control_mode}`
# overrides have been applied. Defaults are False — empty values are
# tolerated unless the caller opts in to the stricter check.
require_non_empty_robot_type: bool = False
require_non_empty_control_mode: bool = False
# Whether `__getitem__` emits the effective per-sample fps as the `fps`
# metadata key. Default `False` so pre-PR checkpoints resume cleanly
# (no new `FPS:` segment in the policy's metadata prefix). Flip to
# `True` for new training runs that want per-sample fps conditioning;
# especially relevant for heterogeneous-frequency mixtures
# (`action_freq=None`). Independent of `metadata_drop_*_prob` — fps
# is intrinsic to the chunk, not a noisy label, so it is always
# present (never padded) for LeRobot samples when this is True.
emit_fps: bool = False
# Mixture-wide defaults for the load-time timestamp-sync check. Each
# dataset can override these via `DatasetConfig.{tolerance_s,
# skip_timestamp_check}`. The default tolerance matches
# `LeRobotDataset.__init__`'s historical default.
tolerance_s: float = 1e-4
skip_timestamp_check: bool = False
def __post_init__(self):
"""Validate dataset mixture configuration."""
if self.weights is not None and len(self.datasets) != len(self.weights):
raise ValueError("The length of `weights` must match the length of `datasets`.")
if self.action_freq is not None and self.action_freq <= 0:
raise ValueError(f"`action_freq` must be a positive number or None, got {self.action_freq}.")
if self.image_resample_strategy not in ["linear", "nearest"]:
raise ValueError(
f"`image_resample_strategy` must be one of ['linear', 'nearest'], got {self.image_resample_strategy}."
)
if self.vector_resample_strategy not in ["linear", "nearest"]:
raise ValueError(
f"`vector_resample_strategy` must be one of ['linear', 'nearest'], got {self.vector_resample_strategy}."
)
if self.val_split_ratio < 0 or self.val_split_ratio > 1:
raise ValueError(f"`val_split_ratio` must be between 0 and 1, got {self.val_split_ratio}.")
if self.tolerance_s < 0:
raise ValueError(f"`tolerance_s` must be >= 0, got {self.tolerance_s}.")
if self.n_obs_history is not None and (
not isinstance(self.n_obs_history, int) or self.n_obs_history < 1
):
raise ValueError(f"`n_obs_history` must be None or a positive integer, got {self.n_obs_history}.")
for name in (
"history_state_drop_prob",
"subgoal_drop_prob",
"subgoal_end_of_segment_prob",
"response_drop_prob",
"metadata_drop_all_prob",
"metadata_drop_each_prob",
):
value = getattr(self, name)
if not 0.0 <= value <= 1.0:
raise ValueError(f"`{name}` must be in [0, 1], got {value}.")
[docs]
@dataclass
class WandBConfig:
"""Configuration for Weights & Biases (wandb) logging.
Args:
enable: Enable Weights & Biases logging. Defaults to False.
entity: The entity name in Weights & Biases, e.g. your username or your
team name. Defaults to None.
project: The project name in Weights & Biases, e.g. "pi0". Defaults to "opentau".
run_id: If provided, the run will be forked from this run ID. Defaults to None.
name: Name of the run, shown in the UI. Defaults to None.
notes: Description of the run, shown in the UI. If None and `enable` is True,
will prompt the user for input. Defaults to None.
tags: Tags to be added to the run in the UI, e.g. ["robot", "v1.0"].
Defaults to empty list.
group: Used to group runs in the UI, e.g. "experiment_1", "experiment_2".
Defaults to None.
job_type: Used to group runs in the UI, e.g. "train", "eval", "test".
Defaults to None.
mode: Allowed values: 'online', 'offline', 'disabled'. Defaults to None
(which uses 'online').
on_resume: What to do when resuming a run that has a `run_id` and a
training `step` (i.e. resuming from a checkpoint). Allowed values:
'fork' (default) creates a NEW run that branches from the parent at
the resume step via wandb's `fork_from`; 'continue' resumes the same
run in place (`resume='allow'`). Forking keeps the parent run
immutable and records server-side lineage; the cost is that a job
preempted/requeued N times produces a chain of N linked runs.
allow_resume: DEPRECATED, use `on_resume` instead. If set, it is mapped
to `on_resume` for backward compatibility (True -> 'continue',
False -> 'fork') and a `FutureWarning` is emitted. Defaults to None.
disable_artifact: Set to True to disable saving an artifact despite
`training.save_checkpoint=True`. Defaults to False.
disable_video: Set to True to skip logging eval grid-summary videos to
wandb. Defaults to False (videos are logged).
"""
enable: bool = False # Enable Weights & Biases logging.
entity: str | None = None # The entity name in Weights & Biases, e.g. your username or your team name
project: str = "opentau" # The project name in Weights & Biases, e.g. "pi0"
run_id: str | None = None # If provided, the run will be forked from this run ID.
name: str | None = None # Name of the run, shown in the UI
notes: str | None = None # Description of the run, shown in the UI
tags: list[str] = field(
default_factory=list
) # Tags to be added to the run in the UI, e.g. ["robot", "v1.0"]
group: str | None = None # Used to group runs in the UI, e.g. "experiment_1", "experiment_2"
job_type: str | None = None # Used to group runs in the UI, e.g. "train", "eval", "test"
mode: str | None = None # Allowed values: 'online', 'offline' 'disabled'. Defaults to 'online'
# On checkpoint resume: 'fork' a new branched run (default) or 'continue' the same run in place.
on_resume: Literal["fork", "continue"] = "fork"
# DEPRECATED, use `on_resume`. Mapped to `on_resume` for back-compat (True->continue, False->fork).
allow_resume: bool | None = None
# Set to true to disable saving an artifact despite training.save_checkpoint=True
disable_artifact: bool = False
# Set to true to skip logging eval grid-summary videos to wandb.
disable_video: bool = False
def __post_init__(self):
"""Map the deprecated `allow_resume` alias, validate, and prompt for notes."""
# Back-compat: `allow_resume` is deprecated in favor of `on_resume`.
if self.allow_resume is not None:
warnings.warn(
"`wandb.allow_resume` is deprecated; use `wandb.on_resume='fork'|'continue'`.",
FutureWarning,
stacklevel=2,
)
# True historically meant "resume the run in place" -> 'continue'.
mapped = "continue" if self.allow_resume else "fork"
# Known limitation: because "fork" is also the default, we cannot tell an explicit
# `on_resume="fork"` apart from the default, so a deprecated `allow_resume=True` paired
# with an explicit `on_resume="fork"` silently lets the alias win (-> "continue"). The
# conflict warning below only fires for the unambiguous case (alias maps to "fork" while
# `on_resume` is "continue"). Both inputs together are a deprecated-config edge case.
if self.on_resume == "fork":
# `on_resume` left at its default -> let the deprecated alias drive behavior.
self.on_resume = mapped
elif self.on_resume != mapped:
warnings.warn(
f"Both `allow_resume`={self.allow_resume} and `on_resume`='{self.on_resume}' "
"are set; using `on_resume`.",
FutureWarning,
stacklevel=2,
)
# Normalize so a re-saved checkpoint config does not re-trigger this warning forever.
self.allow_resume = None
if self.on_resume not in ("fork", "continue"):
raise ValueError(f"`on_resume` must be 'fork' or 'continue', got {self.on_resume!r}.")
# Prompt user for wandb notes if enabled and notes are not provided.
if not self.enable or self.notes is not None:
return
confirm = False
while not confirm:
self.notes = input("Please enter a description for wandb logging:\n")
confirm = input("Confirm (y/N): ").strip().lower() == "y"
[docs]
def to_wandb_kwargs(self, step=None):
"""Convert configuration to keyword arguments for wandb.init().
Args:
step: Optional training step number. If provided along with `run_id`,
used for resuming or forking runs. Defaults to None.
Returns:
Dictionary of keyword arguments suitable for passing to wandb.init().
"""
kwargs = encode_dataclass(self)
# OpenTau-side switches that are not valid `wandb.init()` arguments.
excluded_keys = [
"enable",
"disable_artifact",
"disable_video",
"project",
"on_resume",
"allow_resume",
]
for ek in excluded_keys:
kwargs.pop(ek, None)
run_id = kwargs.pop("run_id", None)
# If both `run_id` and `step` are provided, we are resuming from a checkpoint and
# handle the resuming or forking logic. Use `step is not None` (not truthiness) so a
# resume from a step-0 checkpoint still forks at `?_step=0`.
if run_id is not None and step is not None:
offline = self.mode in ("offline", "disabled")
if self.on_resume == "continue":
# Resume the same run in place.
kwargs["id"] = run_id
kwargs["resume"] = "allow"
elif not offline:
# Fork a new run that branches from the parent at `step`. Server-side lineage
# requires an online run, hence the offline guard below.
kwargs["fork_from"] = f"{run_id}?_step={step}"
note = f"Forked from run {run_id} at step {step}."
kwargs["notes"] = note if kwargs.get("notes") is None else f"{kwargs['notes']}\n{note}"
else:
# Offline/disabled: lineage cannot be resolved, so start a fresh run and only
# record the intended fork in the notes.
note = f"(offline) would fork from run {run_id} at step {step}."
kwargs["notes"] = note if kwargs.get("notes") is None else f"{kwargs['notes']}\n{note}"
return kwargs
[docs]
@dataclass
class EvalConfig:
"""Configuration for evaluation settings.
Args:
n_episodes: Number of episodes to run during evaluation. Defaults to 16.
batch_size: Number of environments to use in a gym.vector.VectorEnv.
Only used for environments that are not already vectorized.
Defaults to 16.
use_async_envs: Whether to use asynchronous environments (multiprocessing).
Defaults to True. RoboCasa eval *requires* the async backend for any
multi-env (batch_size > 1) build: its per-env EGL/GL offscreen render
contexts cross-contaminate in a single process, so SyncVectorEnv feeds
the policy the wrong env's pixels (issue #449). A multi-env RoboCasa
Sync build is therefore auto-promoted to async; LIBERO is unaffected.
max_episodes_rendered: Maximum number of episodes to render as videos.
Defaults to 16.
grid_size: Grid dimensions for video summary (rows, cols). If None, will
be auto-calculated as a square grid. Defaults to None.
video_crf: H.264 constant-rate-factor for the uploaded grid-summary video
(higher = smaller file / lower quality, 0-51). Defaults to 30.
video_preset: x264 encode preset (ultrafast..veryslow); an encode-speed
vs compression-ratio knob that does not change the quality target.
Defaults to "veryfast".
video_frame_stride: Keep only every k-th frame of the grid-summary video
(k>1 shrinks the upload ~linearly and speeds playback up k x).
Defaults to 2.
keep_per_episode_videos: If False, delete the per-episode
eval_episode_*.mp4 clips after the grid summary is built (they are
never uploaded to wandb). Defaults to False.
recording_root: Root directory for saving evaluation recordings.
Defaults to None.
seed: Master seed for the eval simulations (env scene generation). When
set, takes precedence over the top-level `cfg.seed` for seeding the
eval environments; when None (default), falls back to `cfg.seed`.
Does not affect the global `set_seed`. Defaults to None.
decorrelate_rank_seeds: If True, each accelerator rank evaluates a
distinct, orthogonal slice of scenes (for tasks deliberately
replicated across ranks to gain coverage). If False (default), all
ranks seed identically, so the eval is reproducible across world
sizes. Defaults to False.
Raises:
ValueError: If `batch_size` is greater than `n_episodes`, or if any of
`video_crf`, `video_preset`, `video_frame_stride` is out of range.
"""
n_episodes: int = 16
# `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv. (Only used for environments that are not already vectorized.)
batch_size: int = 16
# `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
use_async_envs: bool = True
max_episodes_rendered: int = 16
# Grid dimensions for video summary (rows, cols). If None, will be auto-calculated as square grid.
grid_size: tuple[int, int] | None = None
# ---- Eval grid-summary video encoding (wandb upload storage footprint) ----
# Only the grid summary is uploaded to wandb, so these knobs control the size
# of the wandb media. H.264 constant-rate-factor: higher = smaller / lower
# quality (0=lossless, 23=x264 default, 51=worst). Defaults to 30.
video_crf: int = 30
# x264 encode preset (ultrafast..veryslow): encode-speed vs compression-ratio.
# Does NOT change the CRF quality target. "veryfast" keeps encoding off the
# eval critical path. Defaults to "veryfast".
video_preset: str = "veryfast"
# Write only every k-th frame of the grid summary. k>1 shrinks the upload
# ~linearly and plays back k x faster (at the cost of temporal smoothness).
# Defaults to 2.
video_frame_stride: int = 2
# If False, per-episode eval_episode_*.mp4 clips are deleted once the grid
# summary is built (they are never uploaded to wandb and are the bulk of
# local disk usage). Set True to keep them for inspection. Defaults to False.
keep_per_episode_videos: bool = False
recording_root: str | None = None
# Master seed for the evaluation *simulations* (env scene generation). When
# set, it takes precedence over the top-level `cfg.seed` for seeding the eval
# environments, so the eval scene set can be pinned independently of the
# training/global seed (e.g. hold the eval scenes fixed while sweeping the
# training seed, or vary the eval scenes without disturbing model init /
# dataset shuffling). When None (default), the eval falls back to the
# top-level `cfg.seed`. Resolved by `scripts/eval.py::_resolve_eval_seed`;
# only affects the env/scene seeding, NOT the global `set_seed`.
seed: int | None = None
# Whether each accelerator rank should evaluate a *different* set of scenes.
# Default False: every rank seeds its environments identically, so the scene a
# given (task, episode) maps to does not depend on the world size / node count
# — the eval is reproducible whether it runs on 1 GPU or 16, and two ranks that
# happen to share a task evaluate the same scenes. Set True only when you have
# deliberately replicated a task across ranks (listed it N x to fill N ranks via
# the round-robin sharding) to get N x *distinct* scenes instead of redundant
# copies; each rank then gets an orthogonal, non-overlapping slice of the seed
# line. Trades world-size reproducibility for extra per-task coverage. Plumbed
# by `scripts/eval.py::eval_policy` via `_rank_seed_offset`.
decorrelate_rank_seeds: bool = False
# When set, harvest "goal frames": for every SUCCESSFUL eval episode, save the final raw camera
# frames (all cameras) to this dir as `.png` files named `{task}__seed{seed}__{decoder}__{camera}.png`,
# plus an appended `manifest.csv`. Keyed by the per-episode scene seed so the same scene can be
# matched across checkpoints (pin `eval.seed` so the seed→scene map is identical everywhere) — a
# scene one checkpoint fails can be backfilled from another that solves it. None disables capture.
goal_frames_dir: str | None = None
# Optional comma-separated explicit list of scene seeds to evaluate (e.g. "705,708,712"), overriding
# the contiguous start_seed..start_seed+n_episodes-1 range. Used for goal-frame backfill: re-run only
# the exact scenes a prior checkpoint failed. n_episodes becomes len(list); the final batch is
# right-padded by repeating the last seed to fill the vector env. None = use the contiguous range.
seed_list: str | None = None
# Which training-time norm head to use when calling `policy.select_action`
# on eval observations. Either:
# - set both `robot_type` and `control_mode` to address the head by
# `(robot_type, control_mode)` (preferred for multi-head checkpoints
# trained against the new per-`(robot_type, control_mode)`
# aggregation),
# - or set `dataset_repo_id` to a training-time dataset name (the
# policy maps it to the norm head via its persisted
# `dataset_to_norm_index`; back-compat path that also works on
# legacy per-dataset checkpoints).
# When all three are ``None`` (default), single-head policies fall back
# to the `_resolve_dataset_index` zero-default; multi-head ones raise.
# The robot_type / control_mode pair takes precedence over
# `dataset_repo_id` when both are set. Plumbed by
# `scripts/eval.py::rollout`.
dataset_repo_id: str | None = None
robot_type: str | None = None
control_mode: str | None = None
def __post_init__(self):
"""Validate evaluation configuration."""
if self.batch_size > self.n_episodes:
raise ValueError(
"The eval batch size is greater than the number of eval episodes "
f"({self.batch_size} > {self.n_episodes}). As a result, {self.batch_size} "
f"eval environments will be instantiated, but only {self.n_episodes} will be used. "
"This might significantly slow down evaluation. To fix this, you should update your command "
f"to increase the number of episodes to match the batch size (e.g. `eval.n_episodes={self.batch_size}`), "
f"or lower the batch size (e.g. `eval.batch_size={self.n_episodes}`)."
)
if not 0 <= self.video_crf <= 51:
raise ValueError(f"eval.video_crf must be in [0, 51], got {self.video_crf}.")
if self.video_frame_stride < 1:
raise ValueError(f"eval.video_frame_stride must be >= 1, got {self.video_frame_stride}.")
valid_presets = {
"ultrafast",
"superfast",
"veryfast",
"faster",
"fast",
"medium",
"slow",
"slower",
"veryslow",
}
if self.video_preset not in valid_presets:
raise ValueError(
f"eval.video_preset must be one of {sorted(valid_presets)}, got {self.video_preset!r}."
)