Concepts
This section explains the core concepts used in the OpenTau codebase.
Policies
Policies map observations (e.g., camera images, robot proprioceptive states) to actions (or action chunks).
Policies are implemented as PyTorch modules and inherit from opentau.policies.pretrained.PreTrainedPolicy.
Datasets
Datasets are used to handle data loading and processing. It supports downloading datasets from the Hugging Face Hub and loading them from local disk. The dataset format is versioned (currently v2.1) and utilizes parquet files for data and mp4 files for videos to ensure efficiency and portability. There are currently two types of datasets:
LeRobotDataset: For robotic data.VQADataset: For VLM training datasets.
These datasets are used to train policies.
DatasetMixture
To train policies on multiple datasets simultaneously, OpenTau uses opentau.datasets.dataset_mixture.WeightedDatasetMixture.
This class:
Combines multiple
LeRobotDatasetandVQADatasetinstances.Different weights can be assigned to each dataset to control the sampling frequency; if weights are omitted (or set to
nullin JSON), weights default to dataset lengths.Aggregates statistics from all constituent datasets to ensure consistent normalization across the mixture.
Resamples the action output frequency to match the
action_freqspecified in the configuration. Whenaction_freqisNone(the default), resampling is disabled — each dataset is sampled at its native fps and a single batch can mix samples spanning different real-time horizons (setaction_freqto a positive float to re-anchor every dataset to a common rate). When running mixed-frequency training, opt in toDatasetMixtureConfig.emit_fps=Trueso the per-sample effective rate is surfaced via thefpsstandard-format key (see Optional Standard-Format Keys) and the policy can condition on it.
Metadata
Metadata is crucial for defining the structure and statistics of a dataset. Handled by LeRobotDatasetMetadata and DatasetMetadata, it includes:
Info: Feature shapes, data types, FPS, and robot type.
Stats: Mean, standard deviation, min, and max values for each feature, used for normalization (e.g., standardizing images or normalizing action vectors).
Tasks: Natural language descriptions of the tasks contained in the dataset.
Metadata is stored in JSON files (info.json, stats.json) and JSONL files (tasks.jsonl) within the dataset directory.
Standard Data Format
To ensure compatibility across different datasets and policies, OpenTau introduces the Standard Data Format.
The Standard Data Format is the expected data format returned by torch.utils.data.Dataset’s __getitem__ and the expected input to torch.nn.Module’s forward method. Any new datasets, VLMs, or VLAs that get added to this repository need to adhere to this format. Data being passed to the model during inference should also adhere to this format.
Default format (n_obs_history=None):
{
"camera0": torch.Tensor, # shape (C, H, W) with values from [0, 1] and with the H, W resized to the config's specifications.
"camera1": torch.Tensor, # shape (C, H, W) with values from [0, 1] and with the H, W resized to the config's specifications.
# ...
"camera{num_cams-1}": torch.Tensor, # shape (C, H, W) with values from [0, 1] and with the H, W resized to the config's specifications.
"state": torch.Tensor, # shape (max_state_dim,)
"actions": torch.Tensor, # shape (action_chunk, max_action_dim)
"prompt": str, # the task prompt, e.g. "Pick up the object and place it on the table."
"response": str, # the response from the VLM for vision QA tasks. For LeRobotDataset, this will be an empty string.
"loss_type": str, # the loss type to be applied to this sample (either "CE" for cross entropy or "MSE" for mean squared error)
"img_is_pad": torch.BoolTensor, # shape (num_cams,) with values 0 or 1, where 1 indicates that the camera image is a padded image.
"action_is_pad": torch.BoolTensor, # shape (action_chunk,) with values 0 or 1, where 1 indicates that the action is a padded action.
"real_action_dim": torch.LongTensor, # scalar shape (); collates to (B,). Real (pre-pad) trailing dim of ``actions``. Used by per-policy flow-matching MSE to skip the zero-pad dim columns; see "Action-dim padding mask" below.
"obs_history_is_pad": torch.BoolTensor, # shape (1,) — always False when n_obs_history is None.
}
With observation history (n_obs_history=T):
{
"camera0": torch.Tensor, # shape (T, C, H, W) — T historical steps for camera 0.
# ...
"camera{num_cams-1}": torch.Tensor, # shape (T, C, H, W)
"state": torch.Tensor, # shape (T, max_state_dim)
"actions": torch.Tensor, # shape (action_chunk, max_action_dim)
# ... (prompt, response, loss_type unchanged)
"img_is_pad": torch.BoolTensor, # shape (num_cams,) — camera slot availability.
"action_is_pad": torch.BoolTensor, # shape (action_chunk,)
"real_action_dim": torch.LongTensor, # scalar shape (); collates to (B,). Real (pre-pad) trailing dim of ``actions``.
"obs_history_is_pad": torch.BoolTensor, # shape (T,) — True for timesteps outside the episode boundary.
}
When n_obs_history=T and history_interval=k, observations are sampled at timesteps
\(t - (T-1)k,\; t - (T-2)k,\; \ldots,\; t\) relative to the current timestep \(t\), where
the interval is measured in dataset steps (at the configured action_freq). For timesteps that
fall before the start of the current episode, the observation is clamped to the first step of the
episode and the corresponding entry in obs_history_is_pad is set to True.
The config file will have to provide the following information in TrainPipelineConfig:
H, W: The height and width of the camera images. These should be the same for all cameras.num_cams: The number of cameras for the cloud VLM in the dataset.max_state_dim: The maximum dimension of the state vector.max_action_dim: The maximum dimension of the action vector.action_chunk: The number of actions in the action vector. This is usually 1 for single action tasks, but can be more for multi-action tasks.
The following fields are set in DatasetMixtureConfig:
n_obs_history: Number of historical observation steps to include. WhenNone(default), the single-step format is used. When set to an integerT, cameras and state gain a leading temporal dimension of sizeT.history_interval: Step interval between historical observation steps. Defaults to1. Only relevant whenn_obs_historyis set.
Cameras should be labeled in order of importance (e.g. camera0 is the most important camera, camera1 is the second most important camera, etc.). The model dataset will select the most important cameras to use if num_cams is less than the number of cameras in the dataset.
Action-dim padding mask
For heterogeneous co-training across datasets with different native action
dimensionalities (e.g. a 7-DoF arm dataset and a 14-DoF bimanual dataset in
one mixture), actions is zero-padded along the last axis to
max_action_dim to keep batches rectangular. The per-sample
real_action_dim scalar records the real (pre-pad) trailing dim; each
policy’s flow-matching MSE on the velocity field uses it to skip the
zero-pad columns and only score the dims that the source dataset actually
uses. Without it, the action expert would be supervised against zero
targets on the padded tail dims — under per-dataset normalization (where
padded-dim stats are mean=0, std=0) that signal is a clean “predict 0
here” that contaminates samples in the same batch which do use those dims.
The name parallels (not collides with) the static original_action_dim
local variable in each policy’s inference path
(sample_actions / select_action), which is
self.config.action_feature.shape[0] — the deployed policy’s
homogeneous DoF, distinct from this per-sample dynamic key.
This is distinct from action_is_pad, which masks padded timesteps
along the action chunk. The two masks are AND-ed together inside each
policy’s MSE block: a slot in (B, chunk_size, max_action_dim) contributes
to the loss only when both its timestep is real (~action_is_pad) and
its dim is real (column index < real_action_dim). Policies with a
real-time-inference frozen prefix (pi05, pi05_mem, pi06, pi07,
pi07_paligemma) AND in a third condition, ~prefix_mask — see
flow_matching_masked_mse in opentau.policies.pi06.modeling_pi06
for the full reduction. pi0 has no such frozen prefix and uses only the
two conditions above.
When real_action_dim is absent (single-dataset configs, externally
constructed inference batches), all max_action_dim columns are treated
as real and loss behavior is bit-identical to pre-fix. VQA-style items emit
actions already shaped (action_chunk, max_action_dim) and rely on
the all-True action_is_pad to drive the loss to zero —
real_action_dim matches max_action_dim for those items, and the
all-True action_is_pad still zeros the contribution.
Optional Standard-Format Keys
On top of the core fields above, __getitem__ emits several optional
keys when the dataset has been enriched with segment metadata (see
Attaching segment metadata to a LeRobot v2.1 dataset) or for the subgoal images sampled from
future video frames. Each optional key is always present. Numeric
and image keys pair with an {key}_is_pad boolean flag — zero-filled
+ flag True means “unavailable or masked”. String keys
(response, memory, next_memory, robot_type,
control_mode) don’t get a separate flag: the empty string ""
is itself the pad signal, which also keeps the default PyTorch collate
happy (list of strings, same length as batch).
robot_type and control_mode are dataset-level identifiers
(constant for every sample within a given dataset, distinct across
datasets in a mixture batch) sourced directly from meta/info.json.
Like speed, mistake, and quality, they participate in the
metadata_drop_all_prob / metadata_drop_each_prob dropout rolls —
see Training-time dropout.
fps is the effective per-sample frame rate of the (possibly
resampled) action chunk: DatasetMixtureConfig.action_freq when set,
otherwise the dataset’s native meta.fps. The tokenized FPS: N, ``
segment adds ~3-4 BPE tokens to the metadata prefix, which still fits
comfortably inside the default ``metadata_max_length=52 — but long
robot_type strings combined with a fully-populated metadata batch
leave less headroom than before, and the underlying tokenizer call uses
truncation=True silently. Bump metadata_max_length (a field on
each pi07 / pi07_paligemma config) if you start seeing the trailing
Control: segment get clipped. Heterogeneous-frequency
mixtures (action_freq=None) need it so the policy can condition on
each sample’s rate — a 30 Hz chunk and a 50 Hz chunk carry different
real-time horizons even when both are chunk_size frames long. Unlike
the other metadata fields, fps does not participate in the
dropout rolls — it’s an intrinsic property of the chunk, not a noisy
label, so it’s always non-pad when emitted from a dataset that has a
real frame rate. Samples from VQA datasets (no temporal axis) emit
fps=0, fps_is_pad=True so heterogeneous VLA + VQA mixtures stay
schema-aligned across the batch; the policy’s prepare_metadata
then drops the FPS: segment for those rows. Emission is gated by
DatasetMixtureConfig.emit_fps (default False — pre-PR
checkpoints resume cleanly because the policy’s metadata prefix doesn’t
gain an unfamiliar FPS: segment). Flip to True for new training
runs that want per-sample fps conditioning, especially heterogeneous
mixtures where action_freq=None lets each dataset run at its native
rate. At inference, EnvMetadataConfig.emit_fps (same default
False) gates the eval-side broadcast of cfg.env.fps.
{
# ... core keys above ...
"robot_type": str, # e.g. "aloha", "panda", "human" — copied verbatim
# from `meta/info.json["robot_type"]` (a standard
# LeRobot v2 field). Empty string ("") when the key
# is absent or null (e.g. VQA datasets).
"control_mode": str, # One of {"joint", "ee", "mixed"} when the dataset
# opted in (see PR #183). Empty string ("") when
# `meta/info.json["control_mode"]` is absent.
"memory": str, # Cumulative subtask summary for the current frame's segment.
# Empty string ("") when memory_raw is absent
# (legacy / unannotated dataset).
"next_memory": str, # Memory string for frame t+1 (same as `memory` within a
# segment, differs at segment boundaries). Clipped at episode
# end. Empty string when unavailable.
"speed": torch.LongTensor, # Scalar in {0, 10, 20, ..., 100}; **per-task decile rank** of this
# episode's length-in-frames. Lower = faster, higher = slower.
# Computed by grouping episodes from ``meta/episodes.jsonl`` by
# task (``tasks[0]`` — the codebase assumes the list is length-1)
# and bucketing each episode's length against the
# ``[p5, p15, ..., p95]`` boundaries of *its own task's*
# length distribution. Episodes shorter than p5 bucket to 0;
# episodes at or above p95 bucket to 100; ties at a boundary
# land in the upper bucket (``p_X <= length < p_Y``). Per-task
# boundaries are persisted to ``meta/speed_percentiles.jsonl``;
# delete that file to recompute (e.g. after appending episodes).
# Tasks with fewer than 10 distinct episode lengths are treated
# as sparse and bucket every episode to 50 (the median-equivalent
# neutral default). The per-task framing makes the bucket carry
# task-relative information: "this is a fast example of *this*
# task" rather than "this episode lasted N seconds". See
# :mod:`opentau.datasets.speed_percentiles` for the helpers.
# Populated unconditionally — available on every LeRobotDataset
# regardless of whether the dataset went through
# ``attach_metadata``.
"speed_is_pad": torch.BoolTensor, # True only when the metadata drop rolls in
# _emit_optional_keys fire at training time. The percentile
# compute itself never produces a pad signal: sparse-task
# episodes still get a (neutral) bucket value rather than
# being padded.
"mistake": torch.BoolTensor, # Scalar; True iff the current segment's success flag is False.
"mistake_is_pad": torch.BoolTensor,
"quality": torch.LongTensor, # Scalar in {1,2,3,4,5}; episode-level quality score.
"quality_is_pad": torch.BoolTensor,
"fps": torch.LongTensor, # Scalar; effective per-sample frame rate of the action chunk.
# When `DatasetMixtureConfig.action_freq` is set, every dataset is
# resampled to that rate (via `resolve_delta_timestamps`) and `fps`
# reports `action_freq`. When `action_freq is None` (the default),
# the chunk runs at the dataset's native `meta.fps`. Gated by
# `DatasetMixtureConfig.emit_fps` (default `False` — opt-in); the
# key is omitted entirely when `emit_fps=False`. Does NOT
# participate in `metadata_drop_*_prob` — fps is an intrinsic
# property of the chunk, not a noisy label.
"fps_is_pad": torch.BoolTensor, # Always False when emitted from a real-rate sample. VQA samples
# (no temporal axis) emit `fps_is_pad=True` so heterogeneous
# VLA + VQA batches stay schema-aligned; `prepare_metadata`
# then drops the `FPS:` segment for those rows.
"subgoal0": torch.Tensor, # shape (3, H, W), values in [0,1]. A single future frame from
# camera0 sampled either at end-of-segment (with probability
# `subgoal_end_of_segment_prob`) or uniformly in [t, t+4 seconds].
# ...
"subgoal{num_cams-1}": torch.Tensor,
"subgoal_is_pad": torch.BoolTensor, # Single flag covering every `subgoalK`. Subgoals are either
# all present (annotated dataset, not dropped this step) or
# all padded (legacy dataset, or `subgoal_drop_prob` fired).
# `response` (already in the core fields) may be replaced with ""
# when `response_drop_prob` fires — consumers read "" as masked,
# same convention as `memory` / `next_memory`.
}
Subgoals are always rank-3 (3, H, W) regardless of
n_obs_history — they represent a single future target frame, not a
temporal window. All camera slots share a single subgoal_is_pad
flag because subgoals are all-or-none.
Subgoal image paths are read from meta/info.json under the
subgoals key. When the key is absent (the state of every LeRobot
dataset today), _load_subgoal_frames returns {} and every
subgoalK tensor comes out zero-filled with subgoal_is_pad=True.
Datasets opt in to subgoals by adding the key; the loader then uses the
frame-selection machinery (end-of-segment vs. uniform [t, t+4 s])
described below.
Training-time dropout
Six probability fields on DatasetMixtureConfig control how often
each optional key is masked during a single __getitem__ call.
Masks
are independent per sample (each call rolls fresh). DataLoader
workers seed their own torch RNG, so samples within a batch are
independent across workers; seed globally via torch.manual_seed(...)
for reproducibility.
Field |
Default |
Effect |
|---|---|---|
|
|
Zero-fills |
|
|
Zero-fills every |
|
|
Probability that a present subgoal is sourced from the end of
the current segment. Otherwise sampled uniformly in time from
the current timestamp through |
|
|
Replaces |
|
|
Masks |
|
|
Per-field independent mask roll for each of |
|
|
Whether the five drop rolls above also fire on the validation
split. Default is |
Note
fps is not in either drop pool — when emit_fps=True it
stays non-pad for every LeRobot sample regardless of the rolls. This
means under emit_fps=True, metadata_drop_all_prob=1.0 produces
a “fps-only metadata segment” rather than “no metadata segment at
all” (the policy’s has_metadata branch sees a non-empty metadata
mask and keeps the Metadata: FPS: N, `` block in the prefix). For a
true no-metadata ablation, keep the default ``emit_fps=False.
make_dataset enforces this by giving the validation subset its own
shallow-copied dataset instance with enable_optional_key_dropout
flipped accordingly; the underlying meta / hf_dataset objects
are still shared with the training subset, so the extra copy is cheap.
Legacy datasets that have not been passed through
opentau.scripts.attach_metadata still load: every optional key
appears with a zero/empty value and _is_pad=True, so policies that
consume these fields can train without gating on dataset provenance.
Configs
Configuration management is handled using Draccus.
The main configuration class is opentau.configs.train.TrainPipelineConfig, which orchestrates training settings,
policy configuration, and environment setup. Configs can be loaded from pretrained checkpoints to reproduce experiments.
Environments
Environments wrap simulation or real-robot interfaces compatible with OpenAI Gym/Gymnasium.
The factory src/opentau/envs/factory.py creates vectorized environments for efficient training and evaluation.
Currently, LIBERO and RoboCasa365 are supported, configured via opentau.envs.configs.LiberoEnv and opentau.envs.configs.RoboCasaEnv respectively.