opentau.envs.subgoal

Eval-time subgoal image generators.

The pi07 low-level and pi07-paligemma low-level policies are trained with subgoal images keyed as subgoal{k} (one per camera in config.image_features) plus a single subgoal_is_pad bool. At eval time these keys are produced by a SubgoalImageGenerator that inspects the observation dict (post preprocess_observation + add_envs_task + add_eval_metadata) and returns the matching batch keys. opentau.envs.utils.add_subgoal_images() is the wiring helper that calls the generator and merges its output into the observation.

The only concrete generator implemented today is LiberoLastFrameSubgoalGenerator, which serves the TensorAuto/libero 20fps v2.1 relabel: it looks up the env’s task language (one of the 40 LIBERO task strings) in the dataset, samples a random matching episode at the start of each rollout, and returns that episode’s last frame from each camera as the subgoal.

Classes

LiberoLastFrameSubgoalGenerator([repo_id, ...])	Naive LIBERO subgoal generator backed by TensorAuto/libero.
RoboCasaGoalFrameSubgoalGenerator(...[, ...])	Serve harvested success goal-frames as eval-time subgoals, keyed by scene seed.
SubgoalImageGenerator(*args, **kwargs)	Eval-time contract for producing subgoal images.

class opentau.envs.subgoal.LiberoLastFrameSubgoalGenerator(repo_id: str = 'TensorAuto/libero', root: str | Path | None = None, resolution: tuple[int, int] = (256, 256), num_cams: int = 2, tolerance_s: float = 0.0001, revision: str | None = None, seed: int | None = None)

Bases: object

Naive LIBERO subgoal generator backed by TensorAuto/libero.

For each env in the rollout, looks up the task language string (one of the 40 LIBERO subtask strings) in the dataset’s episode index, samples a random matching episode at the start of the rollout, and returns that episode’s last frame from each video camera as the subgoal. Camera images are resized with aspect-ratio padding to match resolution (the same shape preprocess_observation produces for the live camera{k} keys).

The chosen episodes are pinned by start_episode() and reused across every __call__() in the rollout — fresh draws happen on the next rollout() call, which corresponds to the next env.reset(). If a prompt has no matching episode in the dataset, start_episode() raises ValueError.

Today only repo_id="TensorAuto/libero" is exercised; other LeRobot v2.1 repos with the same camera-name convention should work but are not tested.

Train↔eval distribution shift (intentional, naive baseline):

Training samples a within-episode future frame (BaseDataset._sample_subgoal_frame: uniform [t, t+4s] or segment end). This generator picks a different episode (matched only by task language) and always returns its last frame. That’s a different conditional distribution from training; the policy was not conditioned on a foreign trajectory’s terminus. This is deliberate as a first baseline — see the PR that introduced this module for the rationale. Follow-up improvements (uniform within-episode sampling, or sampling from the rollout’s own replay) would close most of the gap.

Memory:

The per-(episode, camera) frame cache is unbounded — it grows monotonically across rollouts and survives every mid-training eval block. Bounded in practice by the source dataset size (~660 MB worst case for full LIBERO at the default resolution), but a larger source would need an LRU. TODO if/when this generator gets pointed at a non-LIBERO repo.

__init__(repo_id: str = 'TensorAuto/libero', root: str | Path | None = None, resolution: tuple[int, int] = (256, 256), num_cams: int = 2, tolerance_s: float = 0.0001, revision: str | None = None, seed: int | None = None)

Load dataset metadata and build the task-language episode index.

Parameters:

• repo_id – HuggingFace dataset repo to source subgoal frames from. Must be a v2.1 LeRobot dataset registered in DATA_FEATURES_NAME_MAPPING.

• root – Optional local cache root for the dataset. Defaults to the OpenTau-managed location used by LeRobotDataset.

• resolution – Target (H, W) for the returned subgoal images, matching cfg.resolution.

• num_cams – Number of camera slots to populate. Capped by the cameras the source dataset actually exposes (typically 2 for LIBERO).

• tolerance_s – Timestamp tolerance for video frame decoding.

• revision – Optional explicit dataset revision (e.g. "v2.1", a branch, or a commit SHA). None falls back to CODEBASE_VERSION inside LeRobotDatasetMetadata — currently v2.1, matching TensorAuto/libero. Pin this when the codebase and dataset versions might drift apart and you need the eval-time subgoals to come from a specific relabel pass.

• seed – Seed for the per-instance random.Random used to pick episodes in start_episode(). None falls back to a default-constructed Random (entropy-seeded per process). Threading cfg.seed here is what makes episode picks reproducible across runs; the process- global random is not used so concurrent eval_policy_all workers don’t race on shared RNG state.

property fps: int

known_languages() → list[str]

start_episode(prompts: list[str], seeds: list[int] | None = None) → None

Sample one random matching episode per env, store for the rollout.

Parameters:

• prompts – Task language string per env in the vector.

• seeds – Unused — this generator matches by task language and samples a random matching episode, so the scene seed is irrelevant. Accepted to satisfy the SubgoalImageGenerator protocol.

Raises:

ValueError – If any prompt has no matching episode in lang_to_episodes. The error names the prompt and the count of known languages.

class opentau.envs.subgoal.RoboCasaGoalFrameSubgoalGenerator(manifest_dirs: list[str | Path], task: str, resolution: tuple[int, int] = (256, 256), num_cams: int = 3, decoder_preference: tuple[str, ...] = ('flow', 'discrete'))

Bases: object

Serve harvested success goal-frames as eval-time subgoals, keyed by scene seed.

Unlike LiberoLastFrameSubgoalGenerator (which matches by task language and samples a random matching episode), this generator is keyed by the exact per-episode scene seed. Under a pinned cfg.eval.seed each RoboCasa env runs a deterministic scene seed (seed + i), and we previously harvested the terminal camera frames of successful rollouts of those exact scenes — across several checkpoints — into goal_frames/ dirs, each with a manifest.csv (columns task,seed,decoder,success,camera0, camera1,...) plus the per-(task, seed, decoder, camera) .png files.

Given the (task, seed) an env is running, this returns the matching goal frames as subgoal{k}; scenes with no harvested success get subgoal_is_pad=True (no subgoal). That realises the “feed the successful last frame whenever available” semantics, and is the world-model upper-bound probe: the subgoal is the real terminus of a successful trajectory of the very scene being solved, so it measures how much a perfect goal-image predictor could help. The pi07/pi07-paligemma checkpoints were trained with subgoal_end_of_segment_prob > 0, so terminal frames are an in-distribution subgoal.

Scope and assumptions:

• Single task per instance. task is fixed at construction (it matches cfg.env.task, which RoboCasa eval sets to one task per invocation). The manifest’s task column is filtered to it.

• Camera alignment. camera{k} in the manifest maps directly to subgoal{k} (same camera{k} convention the env wrapper and preprocess_observation produce). The experiment must use the same env.camera_name ordering the frames were harvested under.

• Decoder preference. A scene may have a success frame from both the flow and discrete decoders (same scene, slightly different terminal pose). One is chosen per seed by decoder_preference; the chosen decoder’s frames are used for all cameras (never mixed).

• Union across dirs. Multiple goal_frames/ dirs (one per harvest checkpoint) are merged so coverage is the union; on a seed present in several dirs the best-ranked decoder wins, ties broken by dir order.

__init__(manifest_dirs: list[str | Path], task: str, resolution: tuple[int, int] = (256, 256), num_cams: int = 3, decoder_preference: tuple[str, ...] = ('flow', 'discrete'))

Build the seed -> {camera_idx: png_path} index for task.

Parameters:

• manifest_dirs – goal_frames/ directories, each with a manifest.csv and the harvested .png files. Merged into one index (union of scenes).

• task – RoboCasa task name to serve (filters the manifest task column); matches cfg.env.task.

• resolution – Target (H, W) for the returned subgoal images, matching cfg.resolution (the shape preprocess_observation produces for the live camera{k} keys).

• num_cams – Number of camera slots to populate; cameras with index >= num_cams in the manifest are ignored.

• decoder_preference – Decoder ranking when a seed has frames from more than one decoder; earlier = preferred.

property num_scenes: int

Number of distinct scene seeds with a harvested success frame.

start_episode(prompts: list[str], seeds: list[int] | None = None) → None

Pin the per-env scene seeds for this rollout.

Parameters:

• prompts – Task language per env (unused — task is fixed at construction; kept for protocol parity).

• seeds – Per-env scene seed for this rollout. None (rollout not manually seeded) means every env is served as padded (no subgoal), since without the seed we cannot identify the scene.

class opentau.envs.subgoal.SubgoalImageGenerator(*args, **kwargs)

Bases: Protocol

Eval-time contract for producing subgoal images.

Generators are constructed once at the top of opentau.scripts.eval.eval() and plumbed through to opentau.scripts.eval.rollout(), which calls start_episode() once at the start of each rollout (matching the “once per episode reset” cadence) and then __call__() on every step. Implementations are expected to memoise the per-rollout choice inside start_episode() so the per-step call is cheap.

__init__(*args, **kwargs)

start_episode(prompts: list[str], seeds: list[int] | None = None) → None

Re-sample the per-env subgoal source for a fresh rollout.

Called once at the top of opentau.scripts.eval.rollout() after env.reset() and add_envs_task have populated observation["prompt"].

Parameters:

• prompts – Task language string per env in the vector, length env.num_envs.

• seeds – Per-env scene seed used for env.reset(seed=...) this rollout, length env.num_envs (or None if the rollout was not manually seeded). Language-matched generators (e.g. LIBERO) ignore this; scene-seed-keyed generators (e.g. RoboCasaGoalFrameSubgoalGenerator) use it to serve the goal frame of the exact scene being run.