diff --git a/configs/sim_libero.yaml b/configs/sim_libero.yaml index e5ef5bc3..d09c0c4c 100644 --- a/configs/sim_libero.yaml +++ b/configs/sim_libero.yaml @@ -18,12 +18,20 @@ EVALUATION: task_id: 0 num_trials: 50 output_dir: ./evaluate_results/libero/${hydra:runtime.choices.task}/${now:%Y%m%d_%H%M%S} + # Keep the historical single-seed behavior unless these are overridden + # independently. env_seed controls simulator construction; policy_seed + # controls FastWAM inference randomness. + env_seed: ${seed} + policy_seed: ${seed} # Runtime behavior env_num: 1 num_steps_wait: 30 replan_steps: 10 binarize_gripper: true + # Raw convention after action de-normalization. LIBERO-Plus uses + # signed_open_negative; released FastWAM LIBERO data uses zero_one_open_positive. + gripper_action_format: zero_one_open_positive use_action_ensembler: false visualize_future_video: false diff --git a/experiments/libero/README_LIBERO_PLUS_TABLE13.md b/experiments/libero/README_LIBERO_PLUS_TABLE13.md new file mode 100644 index 00000000..9c0106d3 --- /dev/null +++ b/experiments/libero/README_LIBERO_PLUS_TABLE13.md @@ -0,0 +1,215 @@ +# FastWAM LIBERO-Plus Table 13 evaluation + +The persistent evaluator follows the LIBERO-Plus Table 13 inventory and +aggregation convention: + +- all four base suites and all 10,030 classified variants; +- exactly one rollout per variant in full mode; +- Camera, Robot, Language, Light, Background, Noise, and Layout categories; +- `Average = total successes / 10,030`, not the unweighted mean of categories. + +The inventory comes from +`LIBERO-plus/libero/libero/benchmark/task_classification.json`. Each GPU owns a +configurable number of persistent FastWAM worker replicas. All replicas consume +one shared dynamic queue, and result JSON files are independently resumable. + +The full launcher defaults to eight workers per GPU (`W8`), or 64 workers on +eight GPUs. The smoke launcher defaults to one worker per GPU because its +default inventory contains only 14 tasks. + +## Prerequisites + +Install FastWAM and LIBERO-Plus, extract the LIBERO-Plus assets, and prepare: + +- a FastWAM `fastwam.pt` checkpoint; +- its matching LIBERO-Plus `dataset_stats.json`; +- the Wan model files referenced by `configs/sim_libero.yaml`. + +Download `assets.zip` from the official +[`Sylvest/LIBERO-plus`](https://huggingface.co/datasets/Sylvest/LIBERO-plus/tree/main) +repository and extract it so the LIBERO-Plus checkout contains +`libero/libero/assets`. + +This implementation was validated against +[`sylvestf/LIBERO-plus`](https://github.com/sylvestf/LIBERO-plus) commit +`4976dc30028e805ff8094b55501d532c48fec182`. + +The launchers have no host-specific defaults. Set: + +```bash +export LIBERO_PLUS_ROOT="" +export DATASET_STATS_PATH="/dataset_stats.json" +export MODEL_BASE_PATH="" +``` + +Optional variables include `PYTHON_BIN`, `DEPS_ROOT`, `GPU_IDS`, +`WORKERS_PER_GPU`, `TEXT_EMBEDDING_CACHE_DIR`, `CHECKPOINT_LOAD_PATH`, +`TOKENIZER_MODEL_ID`, `NVIDIA_EGL_ROOT`, and `EGL_FALLBACK_GPU`. Both launchers +default `TOKENIZER_MODEL_ID` to `Wan-AI/Wan2.2-TI2V-5B`; use that same value +when generating the prompt cache. + +## Table 13 prompt cache + +W8 requires a complete text-embedding cache. Table 13 contains 10,030 task +entries and 10,002 unique exact `task.language` strings; a cache made only from +the canonical training prompts is incomplete. Without a cache, every worker +loads a separate UMT5 encoder and eight workers cannot fit on one GPU. + +First export a small synthetic dataset containing the exact benchmark +instructions: + +```bash +export TABLE13_PROMPT_DATASET="" +export TABLE13_PROMPT_CACHE="" + +python3 scripts/export_libero_plus_eval_prompts.py \ + --libero-plus-root "${LIBERO_PLUS_ROOT}" \ + --output-dir "${TABLE13_PROMPT_DATASET}" +``` + +Then generate the cache. Batch size 1 matches online single-prompt inference; +`overwrite=false` makes the command resumable: + +```bash +export DIFFSYNTH_MODEL_BASE_PATH="${MODEL_BASE_PATH}" + +torchrun --standalone --nproc_per_node=8 scripts/precompute_text_embeds.py \ + task=libero_uncond_2cam224_1e-4 \ + 'data.train.dataset_dirs=[${oc.env:TABLE13_PROMPT_DATASET}]' \ + 'data.train.text_embedding_cache_dir=${oc.env:TABLE13_PROMPT_CACHE}' \ + model.redirect_common_files=false \ + model.tokenizer_model_id=Wan-AI/Wan2.2-TI2V-5B \ + +text_embedding_batch_size=1 \ + +overwrite=false +``` + +The evaluator validates all selected prompt filenames before starting any GPU +worker. Prompt tensors are loaded lazily per task, so a worker does not preload +all 10,002 embeddings onto its GPU. + +## Smoke evaluation + +The default smoke selects 14 tasks, two from each category, using seed 42: + +```bash +cd FastWAM + +./scripts/eval_fastwam_libero_plus_smoke_8gpu.sh \ + "/fastwam.pt" \ + ./evaluate_results/libero_plus/smoke +``` + +Override `SMOKE_TASKS`, `SMOKE_SEED`, `SMOKE_TRIALS`, `GPU_IDS`, or +`SAVE_VIDEOS=0` as needed. + +`SMOKE_SEED` only controls which task variants are selected. Simulator and +policy randomness can be controlled independently: + +```bash +ENV_SEED=42 \ +POLICY_SEED=50 \ +TABLE13_CATEGORIES=Robot \ +SMOKE_TASKS=64 \ +SAVE_VIDEOS=0 \ +./scripts/eval_fastwam_libero_plus_smoke_8gpu.sh \ + "/fastwam.pt" \ + ./evaluate_results/libero_plus/robot_env42_policy50 +``` + +`TABLE13_CATEGORIES` accepts a comma-separated list of the short Table 13 +labels `Camera`, `Robot`, `Language`, `Light`, `Background`, `Noise`, and +`Layout`, or their canonical classification names. Smoke sampling is balanced +over the selected categories. The direct Python CLI also accepts repeated +`--category` arguments. Setting a category filter with the full launcher runs +every variant in the selected categories; without it, full mode remains the +official 10,030-task evaluation. + +If `ENV_SEED` or `POLICY_SEED` is unset, that value falls back to the Hydra +`cfg.seed`, exactly matching the historical one-seed behavior. With the +official one-trial protocol, every task still uses +`task_suite.get_task_init_states(task_id)[0]`: `ENV_SEED` controls simulator +construction, while `POLICY_SEED` controls FastWAM inference randomness. +`manifest.json` records both controls and the normalized category filter, and +each new task result records the effective environment and policy seeds. + +Before a full run, exercise the exact W8 topology with one task per worker: + +```bash +WORKERS_PER_GPU=8 \ +EGL_LOCK_SCOPE=gpu \ +SMOKE_TASKS=64 \ +SMOKE_MAX_TASKS_PER_WORKER=1 \ +SAVE_VIDEOS=0 \ +TEXT_EMBEDDING_CACHE_DIR="${TABLE13_PROMPT_CACHE}" \ +./scripts/eval_fastwam_libero_plus_smoke_8gpu.sh \ + "/fastwam.pt" \ + ./evaluate_results/libero_plus/w8_smoke +``` + +## Full Table 13 evaluation + +Full mode evaluates all 10,030 variants, one rollout each, and disables videos: + +```bash +export TEXT_EMBEDDING_CACHE_DIR="${TABLE13_PROMPT_CACHE}" + +./scripts/eval_fastwam_libero_plus_full_8gpu.sh \ + "/fastwam.pt" \ + ./evaluate_results/libero_plus/table13_full +``` + +The full launcher defaults to `WORKERS_PER_GPU=8` and `EGL_LOCK_SCOPE=gpu`. +Override `WORKERS_PER_GPU=1 EGL_LOCK_SCOPE=global` for the legacy W1 topology. + +Both launchers enable `--resume`. Re-running with the same output directory +skips valid results and retries missing tasks. `WORKERS_PER_GPU`, +`EGL_LOCK_SCOPE`, and `CHECKPOINT_LOAD_PATH` are runtime topology choices and +may change on resume. Manifest validation still rejects a different checkpoint, +prompt-cache setting, seed control, category filter, or other semantic +configuration. + +Legacy manifests created before independent seed controls remain resumable when +both new seed options are unset and no category filter is requested; this is +the same `cfg.seed` behavior they originally used. An explicit seed or category +filter is a semantic change and therefore requires a new output directory. + +Consequently, an older W1 output created with online prompt encoding cannot be +continued as W8 in the same output directory, because W8 requires a cache. +Start a new W8 output directory for that case. W1 and W8 can share an output +only when both invocations use the same prompt-cache setting. + +For faster startup from non-persistent local storage, copy `fastwam.pt` locally +and set `CHECKPOINT_LOAD_PATH`. The persistent `--checkpoint` path remains the +manifest identity; the evaluator checks that both files have equal size and +SHA256 before spawning workers. + +## Progress and outputs + +```bash +tail -f ./evaluate_results/libero_plus/table13_full/eval.log +tail -f ./evaluate_results/libero_plus/table13_full/worker_logs/gpu0-w0.log +``` + +The parent waits until every worker reports ready before releasing the shared +queue. Ctrl-C or SIGTERM lets workers finish their current task, preserves +completed JSON files, and exits with a message to rerun the same command with +`--resume`. A per-output lock rejects concurrent evaluator parents. + +Successful completion produces: + +- `table13.csv`; +- `table13_summary.json`; +- `task_results.csv`; +- `results//task_NNNN.json`; +- `manifest.json`. + +Regenerate and require a complete report with: + +```bash +python3 experiments/libero/summarize_libero_plus.py \ + --output-dir ./evaluate_results/libero_plus/table13_full \ + --require-complete +``` + +Training-time preview videos and loss/image metrics are not simulator rollout +success metrics. Use this evaluator for benchmark success rates. diff --git a/experiments/libero/README_SELF_COLLECT_ROLLOUTS.md b/experiments/libero/README_SELF_COLLECT_ROLLOUTS.md new file mode 100644 index 00000000..647a20f2 --- /dev/null +++ b/experiments/libero/README_SELF_COLLECT_ROLLOUTS.md @@ -0,0 +1,375 @@ +# FastWAM LIBERO-Plus success-rollout collection + +This pipeline runs a FastWAM checkpoint in LIBERO-Plus, keeps successful +episodes, and publishes a standalone LeRobot v2.1 dataset. Collection is +multi-GPU, resumable, and crash-safe: workers write independent staging +bundles, while the parent process publishes the final dataset only after every +planned attempt has completed and every artifact has passed validation. + +## What a collected training sample contains + +A LIBERO-Plus benchmark task is an environment specification, not a row that a +trainer can consume. The collector turns a successful simulator rollout into: + +- `observation.images.front`: 256 x 256 RGB H.264 video at 20 FPS; +- `observation.images.wrist`: 256 x 256 RGB H.264 video at 20 FPS; +- `observation.state`: 8-D end-effector pose plus gripper state; +- `action`: the executed 7-D de-normalized LIBERO action; +- timestamps, frame/episode/global indices, task index, and provenance. + +The released Robot Initial States category has 1,550 variants of 40 canonical +instructions. A benchmark name may append text such as +`view 0 0 100 0 0 initstate 101`. The evaluation-aligned state-0 preset +defaults to `--prompt-mode benchmark`: the policy and saved dataset task label +both use the exact official `task.language`, including the technical Table 13 +suffix. Across Table 13, non-Language perturbation categories retain their +technical suffixes; Language uses its official paraphrased instruction. +Set `PROMPT_MODE=canonical` only when deliberately falling back to the +suffix-free 40-instruction policy input. + +Changing a rollout seed does not select a new robot initialization class. +These controls are deliberately independent: + +- `robot_init_id`: the `_initstate_N` class encoded by the BDDL variant; +- `env_seed`: environment reset and randomized scene state; +- `policy_seed`: FastWAM diffusion/policy randomness; +- `init_state_index`: saved simulator state used by `fixed_init` or the + evaluation-aligned preset. + +## Prerequisites + +Prepare the following before launching: + +1. this FastWAM checkout and its Python dependencies (`pip install -e .`); +2. a LIBERO-Plus checkout containing `task_classification.json`, BDDL files, + init-state files, and the extracted `libero/libero/assets` directory; +3. the Wan2.2 TI2V 5B model files; +4. a FastWAM `fastwam.pt` checkpoint and its matching LIBERO-Plus + `dataset_stats.json`; +5. the original LIBERO-Plus LeRobot dataset, used as a schema/task reference. + +The official +[LIBERO-Plus assets repository](https://huggingface.co/datasets/Sylvest/LIBERO-plus/tree/main) +provides `assets.zip`. Extract it so the checkout contains +`libero/libero/assets`. The reference LeRobot dataset is published as +[`Sylvest/libero_plus_lerobot`](https://huggingface.co/datasets/Sylvest/libero_plus_lerobot). +Do not commit an assets symlink that points into a machine-local cache. + +The pipeline was validated against +[`sylvestf/LIBERO-plus`](https://github.com/sylvestf/LIBERO-plus) commit +`4976dc30028e805ff8094b55501d532c48fec182`. Install the simulator dependencies +described by that repository in addition to `pip install -e .`; the FastWAM +package metadata does not install LIBERO-Plus itself. + +The wrapper intentionally contains no machine-specific paths. Set these +environment variables: + +| Variable | Required | Meaning | +|---|---:|---| +| `LIBERO_PLUS_ROOT` | yes | LIBERO-Plus repository root | +| `REFERENCE_DATASET` | yes | original LIBERO-Plus LeRobot dataset | +| `MODEL_BASE_PATH` | yes | model root containing the configured Wan files | +| `CHECKPOINT` | yes | checkpoint identity recorded in manifest/provenance | +| `DATASET_STATS_PATH` | yes | matching normalization statistics | +| `CHECKPOINT_LOAD_PATH` | no | byte-identical local checkpoint copy; defaults to `CHECKPOINT` | +| `TEXT_EMBEDDING_CACHE_DIR` | no | exact prompt-embedding cache; blank loads the text encoder | +| `DEPS_ROOT` | no | extra vendored site-packages directory prepended to `PYTHONPATH` | +| `PYTHON_BIN` | no | Python command; defaults to `python3` | +| `GPU_IDS` | no | comma-separated physical GPU IDs; defaults to `0,1,2,3,4,5,6,7` | +| `WORKERS_PER_GPU` | no | persistent policy workers per GPU; safe default is `1` | +| `EGL_LOCK_SCOPE` | no | `global` or `gpu`; safe default is `global` | +| `EGL_FALLBACK_GPU` | no | extra physical GPU exposed for unusual EGL setups; normally unset | +| `MODEL_REDIRECT_COMMON_FILES` | no | Hydra model override; defaults to `false` for the tested recipe | +| `TOKENIZER_MODEL_ID` | no | tokenizer source; defaults to the tested Wan2.2 TI2V 5B ID | + +Example configuration, using placeholders rather than host-specific paths: + +```bash +export LIBERO_PLUS_ROOT="" +export REFERENCE_DATASET="" +export MODEL_BASE_PATH="" +export CHECKPOINT="/fastwam.pt" +export DATASET_STATS_PATH="/dataset_stats.json" +``` + +`CHECKPOINT_LOAD_PATH`, when supplied, is checked against `CHECKPOINT` by both +byte size and SHA256 before workers start. + +Start with one worker per GPU and a global EGL initialization lock. Increase +`WORKERS_PER_GPU` only after measuring model memory and simulator stability. +`EGL_LOCK_SCOPE=gpu` can improve startup throughput, but it is appropriate only +after the target host has demonstrated that concurrent EGL initialization is +safe. + +## Collect Robot Initial States + +For augmentation intended to reproduce the official one-trial Table 13 +protocol, use the dedicated preset: + +```bash +cd FastWAM + +POLICY_SEEDS=42,43 \ +./scripts/collect_fastwam_libero_plus_robot_eval_state0_8gpu.sh \ + ./outputs/table13_robot_eval_state0 \ + --status-every 10 +``` + +For every selected exact Robot task, `table13_eval_state0` does all three of +the following: + +1. restores `task_suite.get_task_init_states(task_id)[0]`; +2. fixes `env_seed=42`; +3. treats `POLICY_SEEDS` / `--policy-seeds` as the only changing randomness. + +It also defaults to the exact official benchmark prompt used by full +evaluation. The collector resolves that prompt from the same +`task_suite.get_task(task_id).language` source as the full evaluator, checks at +runtime that the policy prompt equals the environment prompt, and saves that +exact string as the LeRobot task label. Provenance additionally keeps +`canonical_task` and `benchmark_prompt` as separate fields. + +To reproduce an older canonical-prompt collection instead: + +```bash +PROMPT_MODE=canonical \ +./scripts/collect_fastwam_libero_plus_robot_eval_state0_8gpu.sh \ + ./outputs/table13_robot_eval_state0_canonical +``` + +Every attempt and successful episode records `env_seed`, `policy_seed`, +`initialization`, and `init_state_index` separately. The collection manifest +also freezes the initial-state expression and seed strategy. This prevents a +resumed collection from silently mixing random resets or environment seeds. + +After the base attempts finish, fill underrepresented exact variants until +they have two cumulative successes: + +```bash +./scripts/retry_fastwam_libero_plus_robot_eval_state0_8gpu.sh \ + ./outputs/table13_robot_eval_state0 +``` + +The retry preset prioritizes variants with fewer saved successes, keeps +`env_seed=42`, and advances only `policy_seed` deterministically from +`RETRY_INDEX_START` (default 44). A variant retires after base + retry +successes reaches 2 or `RETRY_MAX_ATTEMPTS`/the runtime budget is exhausted. +Override `--target-successes-per-variant` when a different repeatability +threshold is intentional. + +To stop at an absolute wall-clock time while still merging successes that +finish in flight, pass a timezone-aware ISO-8601 deadline directly: + +```bash +./scripts/retry_fastwam_libero_plus_robot_eval_state0_8gpu.sh \ + ./outputs/table13_robot_eval_state0 \ + --wall-deadline '2030-01-01T18:00:00+08:00' +``` + +The deadline includes startup and worker-readiness time. At the deadline the +Python coordinator stops dispatching new attempts, lets in-flight attempts +finish within the graceful-shutdown allowance, and finalizes the merged +dataset. This is independent of `RETRY_MAX_RUNTIME_HOURS`, which is a +cumulative active-rollout budget; whichever limit is reached first wins. The +normalized deadline is frozen in the campaign manifest, so resume the same +campaign with the same value. To extend it, start a new campaign name and use +a non-overlapping `RETRY_INDEX_START`. + +The older, distribution-expanding random-reset mode remains available: + +Table 13 source, two paired environment/policy seeds: + +```bash +cd FastWAM + +./scripts/collect_fastwam_libero_plus_robot_8gpu.sh \ + ./outputs/table13_robot_seed42_43 \ + --task-source table13 \ + --mode full \ + --rollout-seeds 42,43 \ + --initialization random_reset \ + --prompt-mode canonical \ + --status-every 10 +``` + +Held-out robot initialization classes: + +```bash +./scripts/collect_fastwam_libero_plus_robot_8gpu.sh \ + ./outputs/heldout_robot_seed42_43 \ + --task-source heldout \ + --mode full \ + --heldout-robot-init-ids 5,111,227 \ + --rollout-seeds 42,43 \ + --initialization random_reset \ + --prompt-mode canonical +``` + +`random_reset` keeps the selected robot variant and samples a valid scene from +`env_seed`. For Table 13 tasks, generic `fixed_init` restores selected pruned +MuJoCo states but does not constrain the seed relationship: + +```bash +./scripts/collect_fastwam_libero_plus_robot_8gpu.sh ./outputs/fixed_init \ + --task-source table13 \ + --mode full \ + --rollout-seeds 42,43 \ + --initialization fixed_init \ + --init-state-indices 0 +``` + +Held-out virtual variants require `random_reset`; using a representative base +task's saved state would overwrite the held-out robot initialization, so the +collector rejects held-out tasks with either fixed-state mode. + +Use repeated `--seed-pair ENV:POLICY` arguments to vary environment and policy +seeds independently. In `table13_eval_state0`, every explicit ENV must be 42. +Use repeated `--task-key SUITE/TASK_ID` arguments for an exact smoke subset. + +The wrapper always enables `--resume`. Re-running the same command validates +the immutable manifest, skips completed successes and ordinary rollout +failures, and retries worker exceptions. A file lock prevents two collectors +from writing the same output directory. + +To inspect task selection and manifest construction without loading models: + +```bash +./scripts/collect_fastwam_libero_plus_robot_8gpu.sh ./outputs/prepare_check \ + --task-source heldout \ + --mode smoke \ + --smoke-tasks 8 \ + --prepare-only +``` + +For a graceful stop, send `SIGINT` to the collector parent (or press Ctrl-C in +its terminal). Workers finish their current attempt, and the same command can +resume the remaining jobs. Do not delete `.collector.lock` while the parent is +alive. + +## Prompt-embedding cache + +Without `TEXT_EMBEDDING_CACHE_DIR`, every worker loads the T5 text encoder. A +cache substantially reduces per-worker GPU memory and enables more workers per +GPU. + +Generate the 40-instruction canonical cache from the reference dataset with: + +```bash +export TEXT_EMBEDDING_CACHE_DIR="" +export DIFFSYNTH_MODEL_BASE_PATH="${MODEL_BASE_PATH}" + +python3 scripts/precompute_text_embeds.py \ + task=libero_uncond_2cam224_1e-4 \ + 'data.train.dataset_dirs=[${oc.env:REFERENCE_DATASET}]' \ + 'data.train.text_embedding_cache_dir=${oc.env:TEXT_EMBEDDING_CACHE_DIR}' \ + model.redirect_common_files=false \ + model.tokenizer_model_id=Wan-AI/Wan2.2-TI2V-5B \ + +text_embedding_batch_size=1 +``` + +The model/tokenizer overrides used for cache generation must match those used +by the policy workers. + +Choose the cache that exactly matches `PROMPT_MODE`: + +- `canonical` uses the 40 suffix-free instruction cache generated above; +- `benchmark` uses the 10,002-prompt Table 13 full-evaluation cache. A full + Robot collection selects 1,550 of those exact official prompts. + +Prompt caches are keyed by the exact full prompt. The benchmark cache is not a +superset of the canonical cache because a suffixed and suffix-free instruction +have different hashes. Before writing a new manifest or spawning any GPU +worker, the collector checks every selected prompt using the evaluator's exact +prompt template, SHA256 filename, context length, and encoder suffix. A +mismatch fails once in the parent with the missing count and an example path. + +Workers eagerly move at most 256 distinct prompt contexts to their policy GPU. +The 40-prompt canonical set is therefore preloaded. Benchmark-sized sets are +loaded lazily as tasks reach each worker, avoiding a 1,550-prompt device-cache +copy in every worker; parent preflight still validates the complete selected +set before workers start. + +Cache generation is batch-shape-sensitive in BF16. For a run that began with +online single-prompt encoding, generate and use a cache with +`text_embedding_batch_size=1`; do not switch mid-run to a cache generated in a +larger batch. The cache loader checks exact prompt hashes, shapes, dtypes, and +finite values, then keeps contexts resident on the model device. + +New collection manifests record whether prompt encoding is online or cached. +For a cache, they also record a path-independent fingerprint of every `.pt` +file, so moving an identical cache is allowed but changing cache contents on +resume is rejected. Legacy schema-1 manifests cannot make this check and emit a +warning; their original prompt mode must be preserved manually. + +Schema-2 manifests also record SHA256 identities for the checkpoint, +normalization statistics, task classification, and the reference dataset's +schema/task metadata. This makes semantic input changes fail early on resume +while still allowing identical files, the output tree, and the GPU topology to +move to a different host. + +## Output + +```text +OUTPUT/ +├── attempts/ # one JSON result per completed rollout +├── errors/ # retryable worker exceptions +├── staging/success/JOB_ID/ # crash-safe successful episode bundles +├── worker_logs/ +├── collection_manifest.json +├── collection_summary.json # written after all attempts complete +├── final_report.json # written after deep dataset validation +└── lerobot_dataset/ + ├── data/chunk-*/episode_*.parquet + ├── videos/chunk-*/observation.images.{front,wrist}/episode_*.mp4 + └── meta/{info.json,tasks.jsonl,episodes.jsonl,episodes_stats.jsonl,provenance.jsonl} +``` + +Only successful rollouts become episodes. Failed rollouts remain as attempt +metadata and are not included in `lerobot_dataset`. + +The success-staging transaction is written before its attempt JSON. If a +process stops in that small window, resume validates file sizes and SHA256 +digests and reconstructs the missing attempt record. The finalizer validates +every Parquet file and decodes every MP4 before atomically publishing +`lerobot_dataset`. + +## Continue training + +Add only `OUTPUT/lerobot_dataset` to the training data roots. Reuse the +normalization statistics from the checkpoint being continued: + +```yaml +data: + train: + dataset_dirs: + - + - /lerobot_dataset + dataset_repeats: [1, 1] + pretrained_norm_stats: /dataset_stats.json +``` + +If the base and augmentation roots contain `B` and `A` frames and the +augmentation repeat is `R`, its approximate sampling fraction is: + +```text +R * A / (B + R * A) +``` + +Choose `R` after collection from the final frame count. Do not recompute +normalization statistics when continuing the checkpoint. Dataset task labels +match the prompt seen by the collecting policy: canonical runs can reuse the +40-instruction cache, while benchmark runs retain exact Table 13 suffixes and +need the corresponding benchmark prompt cache during training. + +## Evaluation integrity + +`--task-source table13` directly collects from released Table 13 test +variants. Training on those rollouts contaminates any later score on the same +rows and must be reported as test-distribution self-training, not clean +generalization. + +For a cleaner augmentation study, use `--task-source heldout` and reserve +disjoint `(base_task, robot_init_id, init_state_index, env_seed)` keys for +validation. Success-only collection also introduces policy-selection bias: +hard variants that the current policy never solves cannot enter the resulting +dataset. diff --git a/experiments/libero/collect_libero_plus_self_rollouts.py b/experiments/libero/collect_libero_plus_self_rollouts.py new file mode 100644 index 00000000..57b54f76 --- /dev/null +++ b/experiments/libero/collect_libero_plus_self_rollouts.py @@ -0,0 +1,2168 @@ +#!/usr/bin/env python3 +"""Collect successful FastWAM LIBERO-Plus rollouts as a LeRobot v2.1 dataset. + +The collector keeps worker output transaction-safe: GPU workers only write +per-attempt result files and per-success staging bundles. The parent process +then deterministically finalizes every successful bundle into one standalone +LeRobot dataset that can be appended to LightX2V's ``dataset_dirs``. +""" + +from __future__ import annotations + +import argparse +import ctypes +import fcntl +import hashlib +import json +import multiprocessing as mp +import os +import queue +import random +import re +import signal +import sys +import time +import traceback +from collections import deque +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +from libero_plus_eval_utils import ( + SUITE_ORDER, + LiberoPlusTask, + default_classification_path, + load_task_classification, +) + + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_CONFIG_DIR = PROJECT_ROOT / "configs" +DEFAULT_LIBERO_PLUS_ROOT = PROJECT_ROOT.parent / "LIBERO-plus" +ROBOT_CATEGORY = "Robot Initial States" +RANDOM_RESET_INITIALIZATION = "random_reset" +FIXED_INIT_INITIALIZATION = "fixed_init" +TABLE13_EVAL_STATE0_INITIALIZATION = "table13_eval_state0" +TABLE13_EVAL_ENV_SEED = 42 +MAX_EAGER_PROMPT_CONTEXTS = 256 + +# These robot initialization classes exist in LIBERO-Plus but are not used by +# any Robot Initial States row in the released Table 13 classification. +DEFAULT_HELDOUT_ROBOT_INIT_IDS = ( + 5, + 9, + 12, + 18, + 21, + 22, + 25, + 26, + 28, + 29, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 41, + 46, + 48, + 49, + 52, + 53, + 56, + 59, + 61, + 64, + 66, + 68, + 72, + 75, + 76, + 79, + 82, + 85, + 86, + 91, + 92, + 93, + 96, + 98, + 111, + 122, + 134, + 166, + 179, + 182, + 227, +) + + +@dataclass(frozen=True) +class CollectionTask: + suite: str + task_id: int + classification_id: int + classification_name: str + category: str + difficulty_level: int | None + base_name: str + canonical_task: str + robot_init_id: int + task_source: str + + @property + def key(self) -> str: + if self.task_source == "table13": + return f"{self.suite}/{self.task_id}" + return f"{self.suite}/{self.base_name}/initstate_{self.robot_init_id}" + + def to_dict(self) -> dict: + return asdict(self) + + @classmethod + def from_dict(cls, value: dict) -> "CollectionTask": + return cls( + suite=str(value["suite"]), + task_id=int(value["task_id"]), + classification_id=int(value["classification_id"]), + classification_name=str(value["classification_name"]), + category=str(value["category"]), + difficulty_level=( + None + if value.get("difficulty_level") is None + else int(value["difficulty_level"]) + ), + base_name=str(value["base_name"]), + canonical_task=str(value["canonical_task"]), + robot_init_id=int(value["robot_init_id"]), + task_source=str(value["task_source"]), + ) + + +@dataclass(frozen=True) +class CollectionJob: + task: CollectionTask + env_seed: int + policy_seed: int + initialization: str + init_state_index: int + prompt_mode: str + + @property + def job_id(self) -> str: + init_label = ( + "random" + if self.initialization == RANDOM_RESET_INITIALIZATION + else f"fixed{self.init_state_index:02d}" + ) + source_label = "t13" if self.task.task_source == "table13" else "heldout" + raw = ( + f"{source_label}__{self.task.suite}__task{self.task.task_id:04d}" + f"__robot{self.task.robot_init_id:03d}__{init_label}" + f"__env{self.env_seed}__policy{self.policy_seed}" + ) + digest = hashlib.sha1( + json.dumps(asdict(self), sort_keys=True).encode("utf-8") + ).hexdigest()[:10] + return f"{raw}__{digest}" + + def to_dict(self) -> dict: + payload = asdict(self) + payload["job_id"] = self.job_id + return payload + + @classmethod + def from_dict(cls, value: dict) -> "CollectionJob": + return cls( + task=CollectionTask.from_dict(value["task"]), + env_seed=int(value["env_seed"]), + policy_seed=int(value["policy_seed"]), + initialization=str(value["initialization"]), + init_state_index=int(value["init_state_index"]), + prompt_mode=str(value["prompt_mode"]), + ) + + +def _atomic_json_dump(payload: dict, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + with temporary_path.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, ensure_ascii=False, indent=2) + handle.write("\n") + os.replace(temporary_path, path) + + +def _read_json(path: Path) -> dict: + with path.open(encoding="utf-8") as handle: + return json.load(handle) + + +def _result_path(output_dir: Path, job: CollectionJob) -> Path: + return output_dir / "attempts" / f"{job.job_id}.json" + + +def _error_path(output_dir: Path, job: CollectionJob) -> Path: + return output_dir / "errors" / f"{job.job_id}.json" + + +def _is_complete_attempt(path: Path, job: CollectionJob) -> bool: + if not path.is_file(): + return False + try: + result = _read_json(path) + complete = ( + result.get("job_id") == job.job_id + and result.get("job") == job.to_dict() + and isinstance(result.get("success"), bool) + ) + if complete and bool(result["success"]): + artifact_path = result.get("artifact_path") + expected_artifact = ( + path.parent.parent / "staging" / "success" / job.job_id + ).resolve() + complete = ( + bool(artifact_path) + and expected_artifact.is_dir() + and (expected_artifact / "manifest.json").is_file() + and (expected_artifact / "provenance.json").is_file() + ) + return complete + except (OSError, TypeError, ValueError, json.JSONDecodeError): + return False + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(4 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _prompt_embedding_identity(cache_dir: str | None) -> dict: + """Build a path-independent identity for the policy's prompt input mode.""" + if not cache_dir: + return {"mode": "online"} + + root = Path(cache_dir).expanduser().resolve() + cache_files = sorted(path for path in root.rglob("*.pt") if path.is_file()) + if not cache_files: + raise ValueError(f"Text embedding cache contains no .pt files: {root}") + + digest = hashlib.sha256() + total_bytes = 0 + for path in cache_files: + relative_path = path.relative_to(root).as_posix() + size = path.stat().st_size + file_sha256 = _sha256_file(path) + total_bytes += size + digest.update(relative_path.encode("utf-8")) + digest.update(b"\0") + digest.update(str(size).encode("ascii")) + digest.update(b"\0") + digest.update(file_sha256.encode("ascii")) + digest.update(b"\n") + return { + "mode": "cache", + "file_count": len(cache_files), + "total_bytes": total_bytes, + "sha256": digest.hexdigest(), + } + + +def _model_task_descriptions( + tasks: list[CollectionTask], + *, + prompt_mode: str, + unique: bool = True, +) -> list[str]: + if prompt_mode == "canonical": + descriptions = [task.canonical_task for task in tasks] + elif prompt_mode == "benchmark": + table13_tasks = [ + task for task in tasks if task.task_source == "table13" + ] + heldout_tasks = [ + task for task in tasks if task.task_source == "heldout" + ] + if table13_tasks and heldout_tasks: + raise ValueError( + "Cannot resolve benchmark prompts for mixed table13/heldout tasks." + ) + if table13_tasks: + # Reuse the full evaluator's official benchmark lookup. It verifies + # classification name/task id alignment and returns exact + # task_suite.get_task(task_id).language strings. + from eval_libero_plus_persistent import ( + _load_exact_task_descriptions, + ) + + descriptions = _load_exact_task_descriptions( + [ + LiberoPlusTask( + suite=task.suite, + task_id=task.task_id, + classification_id=task.classification_id, + name=task.classification_name, + category=task.category, + difficulty_level=task.difficulty_level, + ) + for task in table13_tasks + ] + ) + else: + descriptions = [ + ( + f"{task.canonical_task} view 0 0 100 0 0 " + f"initstate {task.robot_init_id}" + ) + for task in heldout_tasks + ] + else: + raise ValueError(f"Unsupported prompt mode: {prompt_mode!r}") + return sorted(set(descriptions)) if unique else descriptions + + +def _preflight_text_embedding_cache( + args: argparse.Namespace, + tasks: list[CollectionTask], + *, + output_dir: Path, + physical_gpu_id: int, +) -> dict[str, int] | None: + """Validate every selected model prompt before any worker is spawned.""" + if not args.text_embedding_cache_dir: + return None + + # Reuse the evaluator's exact full-prompt template, hash, encoder suffix, + # and filename validation instead of maintaining a collector-only variant. + from eval_libero_plus_persistent import ( + WAN_PROMPT_CONTEXT_LEN, + _compose_worker_config, + _validate_text_embedding_cache, + ) + + cache_cfg = _compose_worker_config( + config_dir=str(Path(args.config_dir).expanduser().resolve()), + task_config=args.task_config, + checkpoint=str( + Path(args.checkpoint_load_path or args.checkpoint) + .expanduser() + .resolve() + ), + dataset_stats=str(Path(args.dataset_stats).expanduser().resolve()), + output_dir=str(output_dir), + physical_gpu_id=int(physical_gpu_id), + num_trials=1, + save_videos=False, + gripper_action_format=args.gripper_action_format, + hydra_overrides=list(args.override), + text_embedding_cache_dir=args.text_embedding_cache_dir, + ) + context_len = int( + cache_cfg.model.get("tokenizer_max_len", WAN_PROMPT_CONTEXT_LEN) + ) + descriptions = _model_task_descriptions( + tasks, + prompt_mode=args.prompt_mode, + unique=False, + ) + report = _validate_text_embedding_cache( + args.text_embedding_cache_dir, + descriptions, + context_len=context_len, + ) + print( + "Validated collection text embedding cache before worker spawn: " + f"tasks={report['tasks']} " + f"unique_prompts={report['unique_prompts']} " + f"context_len={context_len} " + f"path={Path(args.text_embedding_cache_dir).expanduser().resolve()}", + flush=True, + ) + return report + + +def _prepare_worker_prompt_contexts( + evaluator, + task_descriptions: list[str], + *, + worker_slot: str, +) -> int: + """Eagerly load small prompt sets and lazily load benchmark-sized sets.""" + unique_descriptions = list(dict.fromkeys(task_descriptions)) + if len(unique_descriptions) > MAX_EAGER_PROMPT_CONTEXTS: + print( + "[prompt-cache] " + f"worker={worker_slot} unique_prompts={len(unique_descriptions)} " + f"exceeds eager_limit={MAX_EAGER_PROMPT_CONTEXTS}; " + "using lazy per-prompt device loading.", + flush=True, + ) + return 0 + return evaluator.preload_prompt_contexts(unique_descriptions) + + +def _recover_staged_attempt( + output_dir: Path, + staging_root: Path, + job: CollectionJob, + *, + checkpoint: str, +) -> bool: + """Recover an attempt JSON published just after its success staging bundle.""" + result_path = _result_path(output_dir, job) + artifact_path = staging_root / "success" / job.job_id + if _is_complete_attempt(result_path, job): + result = _read_json(result_path) + if artifact_path.exists() and result.get("success") is not True: + raise ValueError( + f"Conflicting failure attempt and success staging: {job.job_id}" + ) + return False + + if not artifact_path.exists(): + return False + if not artifact_path.is_dir(): + raise ValueError(f"Staged success is not a directory: {artifact_path}") + + manifest_path = artifact_path / "manifest.json" + provenance_path = artifact_path / "provenance.json" + manifest = _read_json(manifest_path) + if manifest.get("job_id") != job.job_id: + raise ValueError(f"Staged manifest job_id mismatch: {manifest_path}") + + records = manifest.get("files") + if not isinstance(records, list): + raise ValueError(f"Staged manifest has no file records: {manifest_path}") + expected_files = { + "trajectory.npz", + "front.mp4", + "wrist.mp4", + "provenance.json", + } + recorded_files: set[str] = set() + for record in records: + if not isinstance(record, dict): + raise ValueError(f"Invalid staged file record: {manifest_path}") + relative = Path(str(record.get("path", ""))) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"Unsafe staged file path {relative}: {manifest_path}") + staged_file = artifact_path / relative + if not staged_file.is_file(): + raise FileNotFoundError(staged_file) + if staged_file.stat().st_size != int(record.get("size", -1)): + raise ValueError(f"Staged file size mismatch: {staged_file}") + if _sha256_file(staged_file) != record.get("sha256"): + raise ValueError(f"Staged file checksum mismatch: {staged_file}") + recorded_files.add(relative.as_posix()) + if recorded_files != expected_files: + raise ValueError( + f"Staged file set mismatch for {job.job_id}: {sorted(recorded_files)}" + ) + + provenance = _read_json(provenance_path) + if ( + provenance.get("job_id") != job.job_id + or provenance.get("job") != job.to_dict() + or provenance.get("success") is not True + or provenance.get("checkpoint") != checkpoint + ): + raise ValueError(f"Staged provenance mismatch: {provenance_path}") + result = { + key: value + for key, value in provenance.items() + if key not in {"task", "dataset_task", "bddl_file"} + } + result["artifact_path"] = str(artifact_path) + _atomic_json_dump(result, result_path) + _error_path(output_dir, job).unlink(missing_ok=True) + return True + + +def _parse_int_csv(value: str) -> list[int]: + values = [int(part.strip()) for part in value.split(",") if part.strip()] + if not values: + raise ValueError(f"Expected at least one integer, got {value!r}.") + if len(set(values)) != len(values): + raise ValueError(f"Values must be unique, got {values}.") + return values + + +def _parse_gpu_ids(value: str) -> list[int]: + gpu_ids = _parse_int_csv(value) + if any(gpu_id < 0 for gpu_id in gpu_ids): + raise ValueError(f"GPU ids must be non-negative, got {gpu_ids}.") + return gpu_ids + + +def _robot_init_id(name: str) -> int: + match = re.search(r"_initstate_(\d+)(?:_|$)", name) + if match is None: + raise ValueError(f"Robot task has no _initstate_N suffix: {name!r}") + return int(match.group(1)) + + +def _base_name(name: str) -> str: + if "_view_" not in name: + raise ValueError(f"LIBERO-Plus variant has no _view_ suffix: {name!r}") + return name.split("_view_", 1)[0] + + +def _canonical_task(name: str) -> str: + """Return the base LIBERO instruction encoded in a variant filename.""" + base_name = _base_name(name) + # Match LIBERO's filename-to-language handling for LIBERO-10 scene names. + if base_name and base_name[0].isupper() and "SCENE" in base_name: + scene_offset = 8 if "SCENE10" in base_name else 7 + base_name = base_name[base_name.find("SCENE") + scene_offset :] + return base_name.replace("_", " ") + + +def _load_reference_tasks(dataset_root: Path) -> set[str]: + path = dataset_root / "meta" / "tasks.jsonl" + if not path.is_file(): + raise FileNotFoundError(f"Missing reference tasks metadata: {path}") + tasks: set[str] = set() + with path.open(encoding="utf-8") as handle: + for line in handle: + if line.strip(): + tasks.add(str(json.loads(line)["task"])) + return tasks + + +def _table13_robot_tasks( + classification_path: Path, + reference_dataset: Path, +) -> list[CollectionTask]: + reference_tasks = _load_reference_tasks(reference_dataset) + tasks = [] + for task in load_task_classification(classification_path, require_full=True): + if task.category != ROBOT_CATEGORY: + continue + canonical = _canonical_task(task.name) + if canonical not in reference_tasks: + raise ValueError( + f"Canonical Robot task is absent from reference tasks.jsonl: {canonical!r} " + f"(source={task.key}, name={task.name!r})" + ) + tasks.append( + CollectionTask( + suite=task.suite, + task_id=task.task_id, + classification_id=task.classification_id, + classification_name=task.name, + category=task.category, + difficulty_level=task.difficulty_level, + base_name=_base_name(task.name), + canonical_task=canonical, + robot_init_id=_robot_init_id(task.name), + task_source="table13", + ) + ) + if len(tasks) != 1550: + raise ValueError(f"Expected 1550 Table 13 Robot tasks, got {len(tasks)}.") + return tasks + + +def _make_heldout_tasks( + table13_tasks: list[CollectionTask], + robot_init_ids: list[int], +) -> list[CollectionTask]: + """Create virtual Robot variants from the 40 base tasks. + + ``task_id`` points to a representative official variant and is used only + to load that base task's pruned init-state file when fixed initialization + is requested. The environment itself receives the virtual BDDL filename. + """ + by_base: dict[tuple[str, str], CollectionTask] = {} + for task in table13_tasks: + by_base.setdefault((task.suite, task.base_name), task) + + tasks: list[CollectionTask] = [] + for (suite, base_name), representative in sorted(by_base.items()): + for robot_init_id in robot_init_ids: + virtual_name = f"{base_name}_view_0_0_100_0_0_initstate_{robot_init_id}" + tasks.append( + CollectionTask( + suite=suite, + task_id=representative.task_id, + classification_id=-1, + classification_name=virtual_name, + category=ROBOT_CATEGORY, + difficulty_level=None, + base_name=base_name, + canonical_task=representative.canonical_task, + robot_init_id=robot_init_id, + task_source="heldout", + ) + ) + if len(by_base) != 40: + raise ValueError(f"Expected 40 unique Robot base tasks, got {len(by_base)}.") + return tasks + + +def _balanced_smoke_sample( + tasks: list[CollectionTask], + *, + limit: int, + seed: int, +) -> list[CollectionTask]: + if limit < 1: + raise ValueError(f"--smoke-tasks must be positive, got {limit}.") + if limit >= len(tasks): + return list(tasks) + rng = random.Random(seed) + buckets = { + suite: [task for task in tasks if task.suite == suite] for suite in SUITE_ORDER + } + base_quota, remainder = divmod(limit, len(SUITE_ORDER)) + quotas = {suite: base_quota for suite in SUITE_ORDER} + for suite in rng.sample(list(SUITE_ORDER), k=remainder): + quotas[suite] += 1 + selected: list[CollectionTask] = [] + for suite in SUITE_ORDER: + selected.extend(rng.sample(buckets[suite], k=quotas[suite])) + rng.shuffle(selected) + return selected + + +def _select_tasks( + tasks: list[CollectionTask], + *, + mode: str, + smoke_tasks: int, + selection_seed: int, + task_keys: list[str], +) -> list[CollectionTask]: + if task_keys: + lookup: dict[str, CollectionTask] = {} + ambiguous: set[str] = set() + for task in tasks: + aliases = { + task.key, + f"{task.suite}/{task.task_id}", + f"{task.suite}/{task.classification_name}", + } + for alias in aliases: + if alias in lookup and lookup[alias] != task: + ambiguous.add(alias) + else: + lookup[alias] = task + selected = [] + for key in task_keys: + if key in ambiguous: + raise ValueError( + f"Task key {key!r} is ambiguous for this source; use the full " + "virtual key containing /initstate_N." + ) + if key not in lookup: + raise KeyError(f"Unknown collection task key: {key!r}") + selected.append(lookup[key]) + return list(dict.fromkeys(selected)) + if mode == "full": + return list(tasks) + return _balanced_smoke_sample( + tasks, + limit=smoke_tasks, + seed=selection_seed, + ) + + +def _seed_pairs(args: argparse.Namespace) -> list[tuple[int, int]]: + if ( + args.initialization != TABLE13_EVAL_STATE0_INITIALIZATION + and args.policy_seeds is not None + ): + raise ValueError( + "--policy-seeds is only valid with " + f"--initialization {TABLE13_EVAL_STATE0_INITIALIZATION}." + ) + if args.seed_pair: + if args.policy_seeds is not None: + raise ValueError("--seed-pair and --policy-seeds cannot be combined.") + pairs = [] + for raw in args.seed_pair: + try: + env_seed_raw, policy_seed_raw = raw.split(":", 1) + pairs.append((int(env_seed_raw), int(policy_seed_raw))) + except ValueError as error: + raise ValueError( + f"Invalid --seed-pair {raw!r}; expected ENV_SEED:POLICY_SEED." + ) from error + if args.initialization == TABLE13_EVAL_STATE0_INITIALIZATION: + invalid_env_seeds = sorted( + {env_seed for env_seed, _ in pairs} + - {TABLE13_EVAL_ENV_SEED} + ) + if invalid_env_seeds: + raise ValueError( + f"{TABLE13_EVAL_STATE0_INITIALIZATION} fixes env_seed=" + f"{TABLE13_EVAL_ENV_SEED}; invalid env seeds: " + f"{invalid_env_seeds}." + ) + elif args.initialization == TABLE13_EVAL_STATE0_INITIALIZATION: + policy_seeds = _parse_int_csv( + args.policy_seeds + if args.policy_seeds is not None + else args.rollout_seeds + ) + pairs = [ + (TABLE13_EVAL_ENV_SEED, policy_seed) + for policy_seed in policy_seeds + ] + else: + seeds = _parse_int_csv(args.rollout_seeds) + pairs = [(seed, seed) for seed in seeds] + if len(set(pairs)) != len(pairs): + raise ValueError(f"Seed pairs must be unique, got {pairs}.") + return pairs + + +def _build_jobs( + tasks: list[CollectionTask], + *, + seed_pairs: list[tuple[int, int]], + initialization: str, + init_state_indices: list[int], + prompt_mode: str, +) -> list[CollectionJob]: + if initialization == RANDOM_RESET_INITIALIZATION: + indices = [-1] + elif initialization == TABLE13_EVAL_STATE0_INITIALIZATION: + indices = [0] + else: + indices = init_state_indices + jobs = [ + CollectionJob( + task=task, + env_seed=env_seed, + policy_seed=policy_seed, + initialization=initialization, + init_state_index=init_state_index, + prompt_mode=prompt_mode, + ) + for task in tasks + for init_state_index in indices + for env_seed, policy_seed in seed_pairs + ] + job_ids = [job.job_id for job in jobs] + if len(set(job_ids)) != len(job_ids): + raise RuntimeError("Internal error: duplicate collection job ids.") + return jobs + + +def _ensure_libero_config(libero_plus_root: Path, output_dir: Path) -> Path: + config_dir = output_dir / "runtime" / "libero_config" + benchmark_root = libero_plus_root / "libero" / "libero" + payload = { + "benchmark_root": str(benchmark_root), + "bddl_files": str(benchmark_root / "bddl_files"), + "init_states": str(benchmark_root / "init_files"), + "datasets": str(output_dir / "runtime" / "libero_datasets"), + "assets": str(benchmark_root / "assets"), + } + _atomic_json_dump(payload, config_dir / "config.yaml") + Path(payload["datasets"]).mkdir(parents=True, exist_ok=True) + return config_dir + + +def _prepare_libero_runtime( + libero_plus_root: Path, + output_dir: Path, +) -> Path: + """Install a non-interactive LIBERO config before importing ``libero``.""" + config_dir = _ensure_libero_config(libero_plus_root, output_dir) + os.environ["LIBERO_CONFIG_PATH"] = str(config_dir) + return config_dir + + +def _file_identity(path: Path) -> dict: + path = path.expanduser().resolve() + return { + "size": path.stat().st_size, + "sha256": _sha256_file(path), + } + + +def _validate_inputs(args: argparse.Namespace, classification_path: Path) -> dict: + checkpoint_path = Path(args.checkpoint).expanduser().resolve() + checkpoint_load_path = ( + checkpoint_path + if args.checkpoint_load_path is None + else Path(args.checkpoint_load_path).expanduser().resolve() + ) + reference_dataset = Path(args.reference_dataset).expanduser().resolve() + paths = { + "checkpoint": checkpoint_path, + "checkpoint load path": checkpoint_load_path, + "dataset stats": Path(args.dataset_stats).expanduser().resolve(), + "reference dataset": reference_dataset, + "reference info": reference_dataset / "meta" / "info.json", + "reference tasks": reference_dataset / "meta" / "tasks.jsonl", + "classification": classification_path, + "Hydra config dir": Path(args.config_dir).expanduser().resolve(), + "model base path": Path(args.model_base_path).expanduser().resolve(), + "LIBERO-Plus bddl_files": ( + Path(args.libero_plus_root).expanduser().resolve() + / "libero" + / "libero" + / "bddl_files" + ), + "LIBERO-Plus init_files": ( + Path(args.libero_plus_root).expanduser().resolve() + / "libero" + / "libero" + / "init_files" + ), + "LIBERO-Plus assets": ( + Path(args.libero_plus_root).expanduser().resolve() + / "libero" + / "libero" + / "assets" + ), + } + if args.text_embedding_cache_dir: + paths["text embedding cache"] = ( + Path(args.text_embedding_cache_dir).expanduser().resolve() + ) + missing = [f"{label}: {path}" for label, path in paths.items() if not path.exists()] + if missing: + raise FileNotFoundError( + "Missing rollout collection inputs:\n- " + "\n- ".join(missing) + ) + if int(args.workers_per_gpu) < 1: + raise ValueError("--workers-per-gpu must be at least 1.") + if args.egl_fallback_gpu is not None and int(args.egl_fallback_gpu) < 0: + raise ValueError("--egl-fallback-gpu must be non-negative.") + if float(args.worker_ready_timeout) <= 0: + raise ValueError("--worker-ready-timeout must be positive.") + if float(args.status_every) <= 0: + raise ValueError("--status-every must be positive.") + if ( + args.task_source == "heldout" + and args.initialization != RANDOM_RESET_INITIALIZATION + ): + raise ValueError( + "--task-source heldout requires --initialization random_reset. " + "A representative task's saved state would overwrite the held-out " + "robot initialization selected by its virtual BDDL." + ) + if args.initialization == TABLE13_EVAL_STATE0_INITIALIZATION: + if args.task_source != "table13": + raise ValueError( + f"{TABLE13_EVAL_STATE0_INITIALIZATION} requires " + "--task-source table13." + ) + if _parse_int_csv(args.init_state_indices) != [0]: + raise ValueError( + f"{TABLE13_EVAL_STATE0_INITIALIZATION} always uses " + "task_suite.get_task_init_states(task_id)[0]; " + "--init-state-indices must be exactly 0." + ) + input_files = { + "checkpoint": _file_identity(checkpoint_path), + "dataset_stats": _file_identity( + Path(args.dataset_stats).expanduser().resolve() + ), + "classification": _file_identity(classification_path), + "reference_info": _file_identity( + reference_dataset / "meta" / "info.json" + ), + "reference_tasks": _file_identity( + reference_dataset / "meta" / "tasks.jsonl" + ), + } + checkpoint_identity = input_files["checkpoint"] + if checkpoint_load_path.stat().st_size != checkpoint_identity["size"]: + raise ValueError( + "Checkpoint load path size differs from the manifest checkpoint: " + f"{checkpoint_load_path} ({checkpoint_load_path.stat().st_size}) != " + f"{checkpoint_path} ({checkpoint_identity['size']})." + ) + if checkpoint_load_path != checkpoint_path: + load_path_sha256 = _sha256_file(checkpoint_load_path) + if load_path_sha256 != checkpoint_identity["sha256"]: + raise ValueError( + "Checkpoint load path SHA256 differs from the manifest checkpoint: " + f"{checkpoint_load_path} ({load_path_sha256}) != " + f"{checkpoint_path} ({checkpoint_identity['sha256']})." + ) + return input_files + + +def _make_runtime_task(collection_task: CollectionTask, official_task): + if collection_task.task_source == "table13": + if str(official_task.name) != collection_task.classification_name: + raise ValueError( + f"Classification/benchmark mismatch for {collection_task.key}: " + f"{official_task.name!r} != {collection_task.classification_name!r}" + ) + return official_task + + from libero.libero.benchmark import Task + + virtual_name = collection_task.classification_name + return Task( + name=virtual_name, + language=( + f"{collection_task.canonical_task} view 0 0 100 0 0 " + f"initstate {collection_task.robot_init_id}" + ), + problem=official_task.problem, + problem_folder=official_task.problem_folder, + bddl_file=f"{virtual_name}.bddl", + init_states_file=f"{virtual_name}.pruned_init", + ) + + +def _validate_benchmark_prompt_alignment( + job: CollectionJob, + *, + model_prompt: str, + benchmark_prompt: str, +) -> None: + if ( + job.prompt_mode == "benchmark" + and model_prompt != benchmark_prompt + ): + raise ValueError( + "Benchmark prompt mismatch between policy input and " + f"environment task.language for {job.task.key}: " + f"model={model_prompt!r} env={benchmark_prompt!r}" + ) + + +def _collect_episode( + evaluator, job: CollectionJob +) -> tuple[bool, dict[str, np.ndarray], dict]: + import torch + from omegaconf import open_dict + + from action_ensembler import ActionEnsembler + from eval_libero_single import ( + _extract_sim_state, + _get_max_steps, + _predict_action_chunk, + ) + from experiments.libero.libero_utils import ( + LIBERO_ENV_RESOLUTION, + get_libero_dummy_action, + get_libero_env, + get_libero_image, + ) + + task_spec = job.task + task_suite = evaluator._get_suite(task_spec.suite) + if not 0 <= task_spec.task_id < int(task_suite.n_tasks): + raise IndexError( + f"Representative task id {task_spec.task_id} is out of range for " + f"{task_spec.suite} (n_tasks={task_suite.n_tasks})." + ) + official_task = task_suite.get_task(task_spec.task_id) + runtime_task = _make_runtime_task(task_spec, official_task) + model_task = ( + task_spec.canonical_task + if job.prompt_mode == "canonical" + else str(runtime_task.language) + ) + prompt_context = evaluator.get_prompt_context(model_task) + + with open_dict(evaluator.cfg): + evaluator.cfg.seed = int(job.policy_seed) + evaluator.cfg.EVALUATION.task_suite_name = task_spec.suite + evaluator.cfg.EVALUATION.task_id = int(task_spec.task_id) + evaluator.cfg.EVALUATION.category_value = task_spec.category + + evaluator._restore_post_load_rng_state() + env_setup_started_at = time.time() + env, benchmark_prompt = get_libero_env( + runtime_task, + LIBERO_ENV_RESOLUTION, + int(job.env_seed), + ) + env_setup_duration = time.time() - env_setup_started_at + try: + _validate_benchmark_prompt_alignment( + job, + model_prompt=model_task, + benchmark_prompt=benchmark_prompt, + ) + obs = env.reset() + if job.initialization != RANDOM_RESET_INITIALIZATION: + initial_states = list(task_suite.get_task_init_states(task_spec.task_id)) + if not 0 <= job.init_state_index < len(initial_states): + raise IndexError( + f"init_state_index={job.init_state_index} is out of range " + f"for {task_spec.key} ({len(initial_states)} states)." + ) + obs = env.set_init_state(initial_states[job.init_state_index]) + + num_steps_wait = int(evaluator.cfg.EVALUATION.get("num_steps_wait", 30)) + settled_success = False + for _ in range(num_steps_wait): + obs, _, settled_success, _ = env.step(get_libero_dummy_action()) + if settled_success: + break + if settled_success: + return ( + False, + {}, + { + "reason": "task_succeeded_during_unrecorded_settling", + "env_setup_duration": env_setup_duration, + "benchmark_prompt": benchmark_prompt, + "model_prompt": model_task, + }, + ) + + max_steps = _get_max_steps(task_spec.suite) + replan_steps = int(evaluator.cfg.EVALUATION.get("replan_steps", 5)) + use_action_ensembler = bool( + evaluator.cfg.EVALUATION.get("use_action_ensembler", False) + ) + ensembler = ActionEnsembler() if use_action_ensembler else None + if ensembler is not None: + ensembler.reset() + + fronts: list[np.ndarray] = [] + wrists: list[np.ndarray] = [] + states: list[np.ndarray] = [] + actions: list[np.ndarray] = [] + pending_actions: list[list[float]] = [] + success = False + inference_calls = 0 + + episode_started_at = time.time() + for policy_step in range(max_steps): + simulator_t = num_steps_wait + policy_step + if not pending_actions: + action_chunk, imgs, _ = _predict_action_chunk( + obs=obs, + task_description=model_task, + model=evaluator.model, + processor=evaluator.processor, + cfg=evaluator.cfg, + action_horizon=evaluator.action_horizon, + input_w=evaluator.input_w, + input_h=evaluator.input_h, + model_device=evaluator.model_device, + prompt_context=prompt_context, + ) + inference_calls += 1 + if ensembler is not None: + ensembler.add_actions(action_chunk, simulator_t) + pending_actions = [ + ensembler.get_action(timestamp).tolist() + for timestamp in range( + simulator_t, + simulator_t + replan_steps, + ) + ] + else: + pending_actions = action_chunk[:replan_steps].tolist() + else: + imgs = get_libero_image(obs) + + action = np.asarray(pending_actions.pop(0), dtype=np.float32) + state = _extract_sim_state(obs) + front = np.asarray(imgs["image"], dtype=np.uint8) + wrist = np.asarray(imgs["wrist_image"], dtype=np.uint8) + if action.shape != (7,): + raise ValueError( + f"Executed action must have shape (7,), got {action.shape}." + ) + if state.shape != (8,): + raise ValueError( + f"Simulator state must have shape (8,), got {state.shape}." + ) + if front.shape != (256, 256, 3) or wrist.shape != (256, 256, 3): + raise ValueError( + "Expected two 256x256 RGB observations, got " + f"front={front.shape}, wrist={wrist.shape}." + ) + + fronts.append(np.array(front, copy=True)) + wrists.append(np.array(wrist, copy=True)) + states.append(np.array(state, copy=True)) + actions.append(np.array(action, copy=True)) + + obs, reward, done, info = env.step(action) + success = bool(done) + if success: + break + + details = { + "reason": "success" if success else "max_steps", + "env_setup_duration": env_setup_duration, + "episode_duration": time.time() - episode_started_at, + "benchmark_prompt": benchmark_prompt, + "model_prompt": model_task, + "inference_calls": inference_calls, + "frames": len(actions), + "max_steps": max_steps, + } + if not success: + return False, {}, details + + trajectory = { + "front": np.stack(fronts).astype(np.uint8, copy=False), + "wrist": np.stack(wrists).astype(np.uint8, copy=False), + "state": np.stack(states).astype(np.float32, copy=False), + "action": np.stack(actions).astype(np.float32, copy=False), + } + if not all(np.isfinite(trajectory[key]).all() for key in ("state", "action")): + raise ValueError( + "Successful trajectory contains NaN or Inf in state/action." + ) + return True, trajectory, details + finally: + close = getattr(env, "close", None) + if callable(close): + close() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _build_success_provenance( + result: dict, + job: CollectionJob, +) -> dict: + model_prompt = result.get("model_prompt") + benchmark_prompt = result.get("benchmark_prompt") + if not isinstance(model_prompt, str) or not model_prompt: + raise ValueError("Successful rollout result has no model_prompt.") + if not isinstance(benchmark_prompt, str) or not benchmark_prompt: + raise ValueError("Successful rollout result has no benchmark_prompt.") + return { + **result, + # LeRobot uses provenance['task'] as the episode's training label. + "task": model_prompt, + "dataset_task": model_prompt, + "canonical_task": job.task.canonical_task, + "benchmark_prompt": benchmark_prompt, + "bddl_file": f"{job.task.classification_name}.bddl", + } + + +def _worker_main( + worker_slot: str, + physical_gpu_id: int, + worker_args: dict, + task_queue, + status_queue, + start_event, + stop_event, +) -> None: + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(1, signal.SIGTERM) != 0: # PR_SET_PDEATHSIG + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + if os.getppid() != int(worker_args["parent_pid"]): + os._exit(1) + signal.signal(signal.SIGINT, signal.SIG_IGN) + output_dir = Path(worker_args["output_dir"]) + log_path = output_dir / "worker_logs" / f"{worker_slot}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + log_handle = log_path.open("a", encoding="utf-8", buffering=1) + sys.stdout = log_handle + sys.stderr = log_handle + + visible_devices = str(physical_gpu_id) + egl_fallback_gpu = worker_args.get("egl_fallback_gpu") + if ( + egl_fallback_gpu is not None + and int(egl_fallback_gpu) != int(physical_gpu_id) + ): + visible_devices = f"{physical_gpu_id},{int(egl_fallback_gpu)}" + os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices + os.environ["MUJOCO_GL"] = "egl" + os.environ["PYOPENGL_PLATFORM"] = "egl" + os.environ["MUJOCO_EGL_DEVICE_ID"] = "0" + os.environ["LIBERO_CONFIG_PATH"] = worker_args["libero_config_dir"] + os.environ["LIBERO_EGL_INIT_LOCK"] = "1" + egl_lock_file = worker_args["egl_lock_file"] + if worker_args["egl_lock_scope"] == "gpu": + egl_lock_file = f"{egl_lock_file}.gpu{physical_gpu_id}" + os.environ["LIBERO_EGL_INIT_LOCK_FILE"] = egl_lock_file + os.environ["LIBERO_EGL_INIT_RETRIES"] = "5" + os.environ["LIBERO_EGL_INIT_RETRY_SLEEP_SEC"] = "5" + os.environ["TOKENIZERS_PARALLELISM"] = "false" + os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1" + os.environ["DIFFSYNTH_MODEL_BASE_PATH"] = worker_args["model_base_path"] + + for path in ( + worker_args["libero_plus_root"], + str(PROJECT_ROOT), + str(PROJECT_ROOT / "src"), + str(Path(__file__).resolve().parent), + ): + if path not in sys.path: + sys.path.insert(0, path) + + try: + from eval_libero_plus_persistent import ( + _PersistentFastWAMEvaluator, + _compose_worker_config, + ) + from lerobot_rollout_writer import stage_success + + cfg = _compose_worker_config( + config_dir=worker_args["config_dir"], + task_config=worker_args["task_config"], + checkpoint=worker_args["checkpoint_load_path"], + dataset_stats=worker_args["dataset_stats"], + output_dir=worker_args["output_dir"], + physical_gpu_id=physical_gpu_id, + num_trials=1, + save_videos=False, + gripper_action_format=worker_args["gripper_action_format"], + hydra_overrides=worker_args["hydra_overrides"], + text_embedding_cache_dir=worker_args["text_embedding_cache_dir"], + ) + evaluator = _PersistentFastWAMEvaluator(cfg, output_dir) + _prepare_worker_prompt_contexts( + evaluator, + worker_args["model_task_descriptions"], + worker_slot=worker_slot, + ) + status_queue.put( + { + "type": "ready", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + } + ) + start_event.wait() + if stop_event.is_set(): + status_queue.put( + { + "type": "stopped", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + } + ) + log_handle.close() + return + except Exception: + error = traceback.format_exc() + print(error, flush=True) + status_queue.put( + { + "type": "fatal", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "error": error, + "log": str(log_path), + } + ) + log_handle.close() + return + + while True: + if stop_event.is_set(): + break + payload = task_queue.get() + if payload is None or stop_event.is_set(): + break + job = CollectionJob.from_dict(payload) + started_at = time.time() + status_queue.put( + { + "type": "started", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "job_id": job.job_id, + "task": job.task.key, + } + ) + try: + success, trajectory, details = _collect_episode(evaluator, job) + result = { + "schema_version": 1, + "job_id": job.job_id, + "job": job.to_dict(), + "env_seed": int(job.env_seed), + "policy_seed": int(job.policy_seed), + "initialization": job.initialization, + "init_state_index": int(job.init_state_index), + "success": bool(success), + "gpu_id": int(physical_gpu_id), + "checkpoint": worker_args["checkpoint"], + "completed_at": datetime.now(timezone.utc).isoformat(), + "duration": time.time() - started_at, + **details, + } + if success: + provenance = _build_success_provenance(result, job) + artifact_path = stage_success( + Path(worker_args["staging_root"]), + job.job_id, + trajectory, + provenance, + ) + result["artifact_path"] = str(artifact_path) + _atomic_json_dump(result, _result_path(output_dir, job)) + _error_path(output_dir, job).unlink(missing_ok=True) + status_queue.put( + { + "type": "done", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "job_id": job.job_id, + "task": job.task.key, + "success": bool(success), + "frames": int(details.get("frames", 0)), + "duration": time.time() - started_at, + } + ) + except Exception: + error = traceback.format_exc() + print(error, flush=True) + failure = { + "schema_version": 1, + "job_id": job.job_id, + "job": job.to_dict(), + "gpu_id": int(physical_gpu_id), + "time": datetime.now(timezone.utc).isoformat(), + "error": error, + "worker_log": str(log_path), + } + _atomic_json_dump(failure, _error_path(output_dir, job)) + status_queue.put( + { + "type": "failed", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "job_id": job.job_id, + "task": job.task.key, + "error": error.splitlines()[-1] if error.splitlines() else error, + "log": str(log_path), + "duration": time.time() - started_at, + } + ) + + status_queue.put( + { + "type": "stopped", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + } + ) + log_handle.close() + + +def _interleave_jobs(jobs: list[CollectionJob]) -> list[CollectionJob]: + buckets: dict[tuple[str, int], deque[CollectionJob]] = {} + for job in jobs: + buckets.setdefault( + (job.task.suite, job.task.robot_init_id), + deque(), + ).append(job) + scheduled: list[CollectionJob] = [] + while len(scheduled) < len(jobs): + made_progress = False + for key in sorted(buckets): + if buckets[key]: + scheduled.append(buckets[key].popleft()) + made_progress = True + if not made_progress: + raise RuntimeError("Failed to schedule all collection jobs.") + return scheduled + + +def _build_manifest( + args: argparse.Namespace, + classification_path: Path, + gpu_ids: list[int], + tasks: list[CollectionTask], + jobs: list[CollectionJob], + input_files: dict, +) -> dict: + manifest = { + "schema_version": 2, + "created_at": datetime.now(timezone.utc).isoformat(), + "checkpoint": str(Path(args.checkpoint).expanduser().resolve()), + "dataset_stats": str(Path(args.dataset_stats).expanduser().resolve()), + "reference_dataset": str(Path(args.reference_dataset).expanduser().resolve()), + "libero_plus_root": str(Path(args.libero_plus_root).expanduser().resolve()), + "classification_path": str(classification_path), + "task_config": args.task_config, + "task_source": args.task_source, + "mode": args.mode, + "selection_seed": int(args.selection_seed), + "prompt_mode": args.prompt_mode, + "prompt_embedding": _prompt_embedding_identity( + args.text_embedding_cache_dir + ), + "input_files": input_files, + "initialization": args.initialization, + "init_state_indices": ( + [] + if args.initialization == RANDOM_RESET_INITIALIZATION + else _parse_int_csv(args.init_state_indices) + ), + "gpu_ids": gpu_ids, + "gripper_action_format": args.gripper_action_format, + "hydra_overrides": list(args.override), + "tasks": [task.to_dict() for task in tasks], + "jobs": [job.to_dict() for job in jobs], + } + if args.initialization == TABLE13_EVAL_STATE0_INITIALIZATION: + manifest["seed_strategy"] = { + "environment": { + "mode": "fixed", + "seed": TABLE13_EVAL_ENV_SEED, + }, + "policy": { + "mode": "explicit_sequence", + "seeds": list(dict.fromkeys(job.policy_seed for job in jobs)), + }, + } + manifest["initial_state_strategy"] = { + "mode": "official_saved_state", + "expression": "task_suite.get_task_init_states(task_id)[0]", + "index": 0, + } + return manifest + + +def _write_or_validate_manifest( + output_dir: Path, + manifest: dict, + *, + resume: bool, +) -> None: + path = output_dir / "collection_manifest.json" + if not path.exists(): + _atomic_json_dump(manifest, path) + return + existing = _read_json(path) + comparable_existing = dict(existing) + comparable_manifest = dict(manifest) + comparable_existing.pop("created_at", None) + comparable_manifest.pop("created_at", None) + existing_schema = int(comparable_existing.get("schema_version", 1)) + if existing_schema == 1: + # Schema 1 predates prompt-embedding identity. Preserve resume support + # for existing collections, but never silently rewrite their manifest. + comparable_existing.pop("schema_version", None) + comparable_manifest.pop("schema_version", None) + comparable_manifest.pop("prompt_embedding", None) + comparable_manifest.pop("input_files", None) + print( + "[manifest-warning] resuming legacy schema 1 without a recorded " + "prompt-embedding fingerprint; keep the original online/cache mode.", + flush=True, + ) + else: + # Paths and GPU topology are provenance, not semantic identity. Schema + # 2 fingerprints the relevant input files and prompt cache, so an + # identical collection can resume after moving to another host. + for runtime_key in ( + "checkpoint", + "dataset_stats", + "reference_dataset", + "libero_plus_root", + "classification_path", + "gpu_ids", + ): + comparable_existing.pop(runtime_key, None) + comparable_manifest.pop(runtime_key, None) + if comparable_existing != comparable_manifest: + raise ValueError( + f"Existing collection manifest is incompatible with this command: {path}" + ) + if not resume: + raise FileExistsError( + f"Compatible collection already exists; pass --resume: {output_dir}" + ) + + +def _format_duration(seconds: float | None) -> str: + if seconds is None: + return "pending" + seconds = max(0, round(seconds)) + hours, remainder = divmod(seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return f"{hours:02d}:{minutes:02d}:{seconds:02d}" + + +def _attempt_summary(output_dir: Path, jobs: list[CollectionJob]) -> dict: + complete = 0 + successes = 0 + frames = 0 + for job in jobs: + path = _result_path(output_dir, job) + if not _is_complete_attempt(path, job): + continue + complete += 1 + result = _read_json(path) + if bool(result["success"]): + successes += 1 + frames += int(result.get("frames", 0)) + return { + "jobs": len(jobs), + "complete": complete, + "successes": successes, + "failures": complete - successes, + "frames": frames, + "success_rate": successes / complete if complete else 0.0, + } + + +def run(args: argparse.Namespace) -> int: + output_dir = Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + lock_path = output_dir / ".collector.lock" + lock_handle = lock_path.open("a+", encoding="utf-8") + try: + fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as error: + lock_handle.close() + raise RuntimeError( + f"Another collector already holds the output lock: {lock_path}" + ) from error + try: + lock_handle.seek(0) + lock_handle.truncate() + lock_handle.write(f"pid={os.getpid()}\n") + lock_handle.flush() + return _run_locked(args, output_dir) + finally: + fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN) + lock_handle.close() + + +def _run_locked(args: argparse.Namespace, output_dir: Path) -> int: + staging_root = output_dir / "staging" + dataset_root = output_dir / "lerobot_dataset" + libero_plus_root = Path(args.libero_plus_root).expanduser().resolve() + classification_path = ( + Path(args.classification).expanduser().resolve() + if args.classification + else default_classification_path(libero_plus_root) + ) + reference_dataset = Path(args.reference_dataset).expanduser().resolve() + input_files = _validate_inputs(args, classification_path) + + table13_tasks = _table13_robot_tasks(classification_path, reference_dataset) + if args.task_source == "table13": + candidates = table13_tasks + else: + heldout_ids = _parse_int_csv(args.heldout_robot_init_ids) + unknown = sorted(set(heldout_ids) - set(DEFAULT_HELDOUT_ROBOT_INIT_IDS)) + if unknown and not args.allow_table13_robot_init_ids: + raise ValueError( + f"Robot init ids are not in the verified Table 13-heldout set: {unknown}. " + "Pass --allow-table13-robot-init-ids if this is intentional." + ) + candidates = _make_heldout_tasks(table13_tasks, heldout_ids) + + task_keys = list(args.task_key) + tasks = _select_tasks( + candidates, + mode=args.mode, + smoke_tasks=int(args.smoke_tasks), + selection_seed=int(args.selection_seed), + task_keys=task_keys, + ) + init_state_indices = _parse_int_csv(args.init_state_indices) + jobs = _build_jobs( + tasks, + seed_pairs=_seed_pairs(args), + initialization=args.initialization, + init_state_indices=init_state_indices, + prompt_mode=args.prompt_mode, + ) + gpu_ids = _parse_gpu_ids(args.gpus) + worker_slots = [ + (f"gpu{gpu_id}-w{replica_id}", gpu_id) + for gpu_id in gpu_ids + for replica_id in range(int(args.workers_per_gpu)) + ] + libero_config_dir = _prepare_libero_runtime( + libero_plus_root, + output_dir, + ) + _preflight_text_embedding_cache( + args, + tasks, + output_dir=output_dir, + physical_gpu_id=gpu_ids[0], + ) + manifest = _build_manifest( + args, + classification_path, + gpu_ids, + tasks, + jobs, + input_files, + ) + _write_or_validate_manifest(output_dir, manifest, resume=bool(args.resume)) + print( + f"Prepared Robot rollout collection: source={args.task_source} " + f"tasks={len(tasks)} jobs={len(jobs)} gpus={gpu_ids} " + f"workers_per_gpu={args.workers_per_gpu} slots={len(worker_slots)} " + f"initialization={args.initialization} prompt={args.prompt_mode} " + f"text_cache={args.text_embedding_cache_dir or 'disabled'} " + f"checkpoint_load={args.checkpoint_load_path or args.checkpoint} " + f"egl_lock_scope={args.egl_lock_scope} " + f"output={output_dir}", + flush=True, + ) + if args.prepare_only: + return 0 + + if args.resume: + checkpoint_identity = str(Path(args.checkpoint).expanduser().resolve()) + recovered = sum( + _recover_staged_attempt( + output_dir, + staging_root, + job, + checkpoint=checkpoint_identity, + ) + for job in jobs + ) + if recovered: + print( + f"Recovered {recovered} attempt JSON file(s) from verified " + "success staging bundles.", + flush=True, + ) + + pending = [ + job + for job in jobs + if not ( + args.resume and _is_complete_attempt(_result_path(output_dir, job), job) + ) + ] + already_complete = len(jobs) - len(pending) + print( + f"Resume scan: complete={already_complete}, pending={len(pending)}", + flush=True, + ) + + failures: list[dict] = [] + if pending: + worker_args = { + "parent_pid": os.getpid(), + "output_dir": str(output_dir), + "staging_root": str(staging_root), + "config_dir": str(Path(args.config_dir).expanduser().resolve()), + "task_config": args.task_config, + "checkpoint": str(Path(args.checkpoint).expanduser().resolve()), + "checkpoint_load_path": str( + Path(args.checkpoint_load_path or args.checkpoint) + .expanduser() + .resolve() + ), + "dataset_stats": str(Path(args.dataset_stats).expanduser().resolve()), + "gripper_action_format": args.gripper_action_format, + "hydra_overrides": list(args.override), + "text_embedding_cache_dir": ( + None + if not args.text_embedding_cache_dir + else str( + Path(args.text_embedding_cache_dir).expanduser().resolve() + ) + ), + "model_task_descriptions": _model_task_descriptions( + tasks, + prompt_mode=args.prompt_mode, + ), + "libero_plus_root": str(libero_plus_root), + "libero_config_dir": str(libero_config_dir), + "egl_lock_file": os.environ.get( + "LIBERO_EGL_INIT_LOCK_FILE", + "/tmp/fastwam_libero_mujoco_egl_init.lock", + ), + "egl_lock_scope": args.egl_lock_scope, + "egl_fallback_gpu": args.egl_fallback_gpu, + "model_base_path": str(Path(args.model_base_path).expanduser().resolve()), + } + context = mp.get_context("spawn") + task_queue = context.Queue() + task_queue.cancel_join_thread() + status_queue = context.Queue() + start_event = context.Event() + stop_event = context.Event() + stop_requested = False + stop_deadline: float | None = None + previous_signal_handlers = { + signal.SIGINT: signal.getsignal(signal.SIGINT), + signal.SIGTERM: signal.getsignal(signal.SIGTERM), + } + + def request_stop(_signum, _frame) -> None: + nonlocal stop_requested, stop_deadline + if not stop_requested: + print( + "[interrupt] stopping after each worker's current attempt; " + "completed attempts remain resumable.", + flush=True, + ) + stop_deadline = time.time() + 300.0 + stop_requested = True + stop_event.set() + start_event.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + workers = [ + context.Process( + target=_worker_main, + args=( + worker_slot, + gpu_id, + worker_args, + task_queue, + status_queue, + start_event, + stop_event, + ), + name=f"libero-rollout-{worker_slot}", + ) + for worker_slot, gpu_id in worker_slots + ] + process_by_slot = { + worker_slot: worker + for (worker_slot, _), worker in zip(worker_slots, workers, strict=True) + } + for worker in workers: + worker.start() + workers_started_at = time.time() + for job in _interleave_jobs(pending): + task_queue.put(job.to_dict()) + for _ in workers: + task_queue.put(None) + + ready: set[str] = set() + completed_this_run = 0 + successes_this_run = 0 + duration_sum = 0.0 + fatal_errors: list[dict] = [] + fatal_slots: set[str] = set() + last_status_time = 0.0 + rollout_started_at: float | None = None + while any(worker.is_alive() for worker in workers): + try: + event = status_queue.get(timeout=2.0) + except queue.Empty: + event = None + if event is not None: + event_type = event["type"] + if event_type == "ready": + ready.add(str(event["worker_slot"])) + print( + f"[worker-ready] slot={event['worker_slot']} " + f"gpu={event['gpu_id']} " + f"ready={len(ready)}/{len(workers)}", + flush=True, + ) + if len(ready) == len(workers): + start_event.set() + elif event_type == "started": + if rollout_started_at is None: + rollout_started_at = time.time() + elif event_type == "done": + completed_this_run += 1 + successes_this_run += int(bool(event["success"])) + duration_sum += float(event["duration"]) + print( + f"[attempt-done] slot={event['worker_slot']} " + f"gpu={event['gpu_id']} task={event['task']} " + f"success={event['success']} frames={event['frames']} " + f"completed={already_complete + completed_this_run}/{len(jobs)}", + flush=True, + ) + elif event_type == "failed": + failures.append(event) + completed_this_run += 1 + duration_sum += float(event["duration"]) + print( + f"[attempt-error] slot={event['worker_slot']} " + f"gpu={event['gpu_id']} task={event['task']} " + f"error={event['error']} log={event['log']}", + flush=True, + ) + elif event_type == "fatal": + worker_slot = str(event["worker_slot"]) + if worker_slot not in fatal_slots: + fatal_slots.add(worker_slot) + fatal_errors.append(event) + print( + f"[worker-fatal] slot={event['worker_slot']} " + f"gpu={event['gpu_id']} log={event['log']}\n" + f"{event['error']}", + flush=True, + ) + + if not start_event.is_set() and not stop_requested: + for worker_slot, worker in process_by_slot.items(): + if ( + worker_slot not in ready + and worker_slot not in fatal_slots + and worker.exitcode is not None + ): + fatal_slots.add(worker_slot) + fatal = { + "type": "fatal", + "worker_slot": worker_slot, + "gpu_id": dict(worker_slots)[worker_slot], + "error": ( + "Worker exited before ready without a fatal event " + f"(exitcode={worker.exitcode})." + ), + "log": str( + output_dir / "worker_logs" / f"{worker_slot}.log" + ), + } + fatal_errors.append(fatal) + print( + f"[worker-fatal] slot={worker_slot} " + f"gpu={fatal['gpu_id']} log={fatal['log']}\n" + f"{fatal['error']}", + flush=True, + ) + if ( + not fatal_errors + and time.time() - workers_started_at + > float(args.worker_ready_timeout) + ): + waiting = sorted(set(process_by_slot) - ready) + fatal = { + "type": "fatal", + "worker_slot": "parent-timeout", + "gpu_id": -1, + "error": ( + "Timed out waiting for rollout workers to become ready: " + f"{waiting}" + ), + "log": str(output_dir / "worker_logs"), + } + fatal_slots.add("parent-timeout") + fatal_errors.append(fatal) + print(f"[worker-fatal] {fatal['error']}", flush=True) + + now = time.time() + if stop_deadline is not None and now >= stop_deadline: + print( + "[interrupt] graceful-stop deadline reached; terminating " + "remaining workers.", + flush=True, + ) + break + if now - last_status_time >= float(args.status_every): + average = ( + None + if completed_this_run == 0 + else duration_sum / completed_this_run + ) + remaining = max(0, len(pending) - completed_this_run) + eta = ( + None + if completed_this_run == 0 or rollout_started_at is None + else ( + (now - rollout_started_at) + * remaining + / completed_this_run + ) + ) + print( + f"[status] completed={already_complete + completed_this_run}/{len(jobs)} " + f"successes_this_run={successes_this_run} errors={len(failures)} " + f"ready={len(ready)}/{len(workers)} " + f"avg={_format_duration(average)} eta={_format_duration(eta)}", + flush=True, + ) + last_status_time = now + if fatal_errors: + start_event.set() + for worker in workers: + if worker.is_alive(): + worker.terminate() + break + + shutdown_deadline = ( + stop_deadline + if stop_deadline is not None + else time.time() + 30.0 + ) + while any(worker.is_alive() for worker in workers): + if time.time() >= shutdown_deadline: + break + for worker in workers: + if worker.is_alive(): + worker.join(timeout=0.2) + for worker in workers: + if worker.is_alive(): + worker.terminate() + worker.join(timeout=5) + if worker.is_alive(): + worker.kill() + worker.join(timeout=5) + + while True: + try: + event = status_queue.get_nowait() + except queue.Empty: + break + if event["type"] == "failed": + failures.append(event) + elif event["type"] == "fatal": + worker_slot = str(event["worker_slot"]) + if worker_slot not in fatal_slots: + fatal_slots.add(worker_slot) + fatal_errors.append(event) + + for signum, previous_handler in previous_signal_handlers.items(): + signal.signal(signum, previous_handler) + if stop_requested: + print( + "[interrupt] workers stopped cleanly; rerun the same command " + "with --resume to continue.", + flush=True, + ) + return 130 + if fatal_errors: + raise RuntimeError( + f"{len(fatal_errors)} rollout workers failed during initialization. " + "See worker logs and rerun with --resume." + ) + + summary = _attempt_summary(output_dir, jobs) + _atomic_json_dump(summary, output_dir / "collection_summary.json") + print( + "Collection attempts: " + f"complete={summary['complete']}/{summary['jobs']} " + f"success={summary['successes']} failure={summary['failures']} " + f"success_rate={summary['success_rate']:.2%} frames={summary['frames']}", + flush=True, + ) + if failures or summary["complete"] != summary["jobs"]: + raise RuntimeError( + f"Collection is incomplete: errors={len(failures)}, " + f"complete={summary['complete']}/{summary['jobs']}. " + "Fix the reported issue and rerun the same command with --resume." + ) + if summary["successes"] == 0: + raise RuntimeError( + "No successful rollout was collected, so no training dataset can be finalized. " + "Add rollout seeds/tasks and use a new output directory." + ) + + from lerobot_rollout_writer import finalize_dataset + + build_report = finalize_dataset( + staging_root, + dataset_root, + reference_dataset=reference_dataset, + ) + # finalize_dataset performs the same deep validation (including decoding + # every video) before atomically publishing the dataset. + validation_report = dict(build_report) + _atomic_json_dump( + { + "attempts": summary, + "build": build_report, + "validation": validation_report, + }, + output_dir / "final_report.json", + ) + print( + f"Finalized LeRobot dataset: {dataset_root} " + f"episodes={validation_report['total_episodes']} " + f"frames={validation_report['total_frames']}", + flush=True, + ) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", required=True) + parser.add_argument( + "--checkpoint-load-path", + default=None, + help=( + "Optional local byte-identical checkpoint used by workers for faster " + "loading. Size and SHA256 are verified; --checkpoint remains the " + "manifest/provenance identity." + ), + ) + parser.add_argument("--dataset-stats", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument( + "--reference-dataset", + required=True, + help="Existing LIBERO-Plus LeRobot dataset used to validate task/schema compatibility.", + ) + parser.add_argument("--libero-plus-root", default=str(DEFAULT_LIBERO_PLUS_ROOT)) + parser.add_argument("--classification", default=None) + parser.add_argument("--config-dir", default=str(DEFAULT_CONFIG_DIR)) + parser.add_argument("--task-config", default="libero_uncond_2cam224_1e-4") + parser.add_argument( + "--model-base-path", + required=True, + help="Directory containing the Wan model files referenced by the Hydra config.", + ) + parser.add_argument("--gpus", default="0,1,2,3,4,5,6,7") + parser.add_argument( + "--workers-per-gpu", + type=int, + default=1, + help="Number of independent rollout worker processes assigned to each GPU.", + ) + parser.add_argument( + "--text-embedding-cache-dir", + default=None, + help=( + "Directory containing exact selected-prompt T5 embedding caches. " + "When set, workers skip loading the text encoder." + ), + ) + parser.add_argument( + "--egl-lock-scope", + choices=("global", "gpu"), + default="global", + help=( + "Serialize MuJoCo EGL environment construction globally or only " + "between workers sharing one physical GPU." + ), + ) + parser.add_argument( + "--egl-fallback-gpu", + type=int, + default=None, + help=( + "Optional extra physical GPU appended to CUDA_VISIBLE_DEVICES for " + "hosts whose EGL stack requires a separate display GPU. Normally unset." + ), + ) + parser.add_argument("--mode", choices=("smoke", "full"), default="smoke") + parser.add_argument("--smoke-tasks", type=int, default=8) + parser.add_argument("--selection-seed", type=int, default=42) + parser.add_argument( + "--task-key", + action="append", + default=[], + help=( + "Exact task selector (repeatable), e.g. libero_spatial/258. " + "Overrides --mode/--smoke-tasks selection." + ), + ) + parser.add_argument( + "--task-source", + choices=("table13", "heldout"), + default="table13", + help=( + "table13 uses exact Robot test variants (test leakage if trained); " + "heldout uses robot init classes absent from Table 13." + ), + ) + parser.add_argument( + "--heldout-robot-init-ids", + default="5,111,227", + help="Comma-separated robot init classes for --task-source heldout.", + ) + parser.add_argument( + "--allow-table13-robot-init-ids", + action="store_true", + help="Allow arbitrary initstate_N values in heldout virtual tasks.", + ) + parser.add_argument( + "--rollout-seeds", + default="42,43", + help=( + "Comma-separated seeds; each value is used for both env and policy " + "except in table13_eval_state0 mode, where it is a policy-only " + "sequence kept for backward CLI convenience." + ), + ) + parser.add_argument( + "--policy-seeds", + default=None, + help=( + "Explicit comma-separated policy-only seed sequence for " + "table13_eval_state0. The environment seed remains fixed at 42." + ), + ) + parser.add_argument( + "--seed-pair", + action="append", + default=[], + metavar="ENV:POLICY", + help="Explicit env/policy seed pair (repeatable); overrides --rollout-seeds.", + ) + parser.add_argument( + "--initialization", + choices=( + RANDOM_RESET_INITIALIZATION, + FIXED_INIT_INITIALIZATION, + TABLE13_EVAL_STATE0_INITIALIZATION, + ), + default=RANDOM_RESET_INITIALIZATION, + help=( + "random_reset preserves the task's robot variant and samples a new scene " + "from env seed; fixed_init reproduces selected saved simulator states; " + "table13_eval_state0 exactly matches official Table 13 state index 0, " + "fixes env_seed=42, and varies only policy_seed." + ), + ) + parser.add_argument( + "--init-state-indices", + default="0", + help=( + "Comma-separated pruned init-state indices used by fixed_init. " + "table13_eval_state0 requires and always uses index 0." + ), + ) + parser.add_argument( + "--prompt-mode", + choices=("canonical", "benchmark"), + default="canonical", + help=( + "Language supplied to FastWAM and saved as the dataset task label. " + "benchmark uses the exact official task.language string." + ), + ) + parser.add_argument( + "--gripper-action-format", + choices=("zero_one_open_positive", "signed_open_negative"), + default="signed_open_negative", + ) + parser.add_argument("--resume", action="store_true") + parser.add_argument("--prepare-only", action="store_true") + parser.add_argument("--status-every", type=float, default=30.0) + parser.add_argument( + "--worker-ready-timeout", + type=float, + default=1200.0, + help="Fail instead of waiting forever if all rollout workers are not ready.", + ) + parser.add_argument( + "--override", + action="append", + default=[], + help="Additional Hydra override forwarded to each worker (repeatable).", + ) + return parser + + +def main() -> None: + args = build_parser().parse_args() + raise SystemExit(run(args)) + + +if __name__ == "__main__": + main() diff --git a/experiments/libero/eval_libero_plus_persistent.py b/experiments/libero/eval_libero_plus_persistent.py new file mode 100644 index 00000000..79d48768 --- /dev/null +++ b/experiments/libero/eval_libero_plus_persistent.py @@ -0,0 +1,1473 @@ +#!/usr/bin/env python3 +"""Multi-GPU LIBERO-Plus evaluator with persistent FastWAM worker replicas.""" + +from __future__ import annotations + +import argparse +import ctypes +import fcntl +import hashlib +import json +import multiprocessing as mp +import os +import queue +import signal +import sys +import time +import traceback +from collections import deque +from datetime import datetime, timezone +from pathlib import Path + +from libero_plus_eval_utils import ( + CATEGORY_LABELS, + CATEGORY_ORDER, + EXPECTED_TOTAL_TASKS, + SUITE_ORDER, + LiberoPlusTask, + default_classification_path, + failure_path, + is_complete_result, + load_task_classification, + result_path, + select_smoke_tasks, +) + + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_CONFIG_DIR = PROJECT_ROOT / "configs" +DEFAULT_LIBERO_PLUS_ROOT = PROJECT_ROOT.parent / "LIBERO-plus" +WAN_PROMPT_TEMPLATE = ( + "A video recorded from a robot's point of view executing the following " + "instruction: {task}" +) +WAN_PROMPT_CONTEXT_LEN = 128 +WAN_PROMPT_ENCODER_ID = "wan22ti2v5b" + + +def _sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_checkpoint_load_path( + checkpoint: str | Path, + checkpoint_load_path: str | Path | None, +) -> Path: + checkpoint = Path(checkpoint).expanduser().resolve() + load_path = ( + checkpoint + if checkpoint_load_path is None + else Path(checkpoint_load_path).expanduser().resolve() + ) + if not checkpoint.is_file(): + raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint}") + if not load_path.is_file(): + raise FileNotFoundError(f"Checkpoint load path does not exist: {load_path}") + checkpoint_size = checkpoint.stat().st_size + if load_path.stat().st_size != checkpoint_size: + raise ValueError( + "Checkpoint load path size differs from the manifest checkpoint: " + f"{load_path} ({load_path.stat().st_size}) != " + f"{checkpoint} ({checkpoint_size})." + ) + if load_path != checkpoint: + checkpoint_sha256 = _sha256_file(checkpoint) + load_path_sha256 = _sha256_file(load_path) + if load_path_sha256 != checkpoint_sha256: + raise ValueError( + "Checkpoint load path SHA256 differs from the manifest checkpoint: " + f"{load_path} ({load_path_sha256}) != " + f"{checkpoint} ({checkpoint_sha256})." + ) + return load_path + + +def _build_worker_slots( + gpu_ids: list[int], + workers_per_gpu: int, +) -> list[tuple[str, int]]: + if workers_per_gpu < 1: + raise ValueError("--workers-per-gpu must be at least 1.") + return [ + (f"gpu{gpu_id}-w{replica_id}", gpu_id) + for gpu_id in gpu_ids + for replica_id in range(workers_per_gpu) + ] + + +def _load_exact_task_descriptions( + tasks: list[LiberoPlusTask], +) -> list[str]: + from libero.libero import benchmark + + benchmark_dict = benchmark.get_benchmark_dict() + suite_cache = {} + descriptions: list[str] = [] + for task_spec in tasks: + if task_spec.suite not in suite_cache: + if task_spec.suite not in benchmark_dict: + raise KeyError(f"Unknown LIBERO suite {task_spec.suite!r}") + suite_cache[task_spec.suite] = benchmark_dict[task_spec.suite]() + task = suite_cache[task_spec.suite].get_task(task_spec.task_id) + if str(task.name) != task_spec.name: + raise ValueError( + "Classification/benchmark task mismatch for " + f"{task_spec.key}: classification={task_spec.name!r}, " + f"benchmark={task.name!r}." + ) + description = str(task.language) + if not description: + raise ValueError( + f"LIBERO task has an empty language instruction: {task_spec.key}" + ) + descriptions.append(description) + return descriptions + + +def _validate_text_embedding_cache( + cache_dir: str | Path, + task_descriptions: list[str], + *, + context_len: int = WAN_PROMPT_CONTEXT_LEN, +) -> dict[str, int]: + context_len = int(context_len) + if context_len < 1: + raise ValueError(f"Prompt context length must be positive, got {context_len}.") + cache_dir = Path( + os.path.expanduser(os.path.expandvars(str(cache_dir))) + ).resolve() + if not cache_dir.is_dir(): + raise NotADirectoryError( + f"Text embedding cache directory does not exist: {cache_dir}" + ) + unique_descriptions = sorted(set(task_descriptions)) + missing: list[Path] = [] + for description in unique_descriptions: + full_prompt = WAN_PROMPT_TEMPLATE.format(task=description) + digest = hashlib.sha256(full_prompt.encode("utf-8")).hexdigest() + cache_path = cache_dir / ( + f"{digest}.t5_len{context_len}." + f"{WAN_PROMPT_ENCODER_ID}.pt" + ) + if not cache_path.is_file(): + missing.append(cache_path) + if missing: + raise FileNotFoundError( + "LIBERO-Plus text embedding cache is incomplete: " + f"missing={len(missing)}/{len(unique_descriptions)} " + f"first={missing[0]}" + ) + return { + "tasks": len(task_descriptions), + "unique_prompts": len(unique_descriptions), + } + + +def _atomic_json_dump(payload: dict, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + with temporary_path.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, ensure_ascii=True, indent=2) + os.replace(temporary_path, path) + + +def _format_duration(seconds: float | None) -> str: + if seconds is None: + return "pending" + seconds = max(0, round(seconds)) + hours, remainder = divmod(seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return f"{hours:02d}:{minutes:02d}:{seconds:02d}" + + +def _interleave_tasks(tasks: list[LiberoPlusTask]) -> list[LiberoPlusTask]: + """Round-robin suite/category buckets for early load and ETA balance.""" + buckets = { + (suite, category): deque( + task + for task in tasks + if task.suite == suite and task.category == category + ) + for category in CATEGORY_ORDER + for suite in SUITE_ORDER + } + scheduled: list[LiberoPlusTask] = [] + while len(scheduled) < len(tasks): + made_progress = False + for bucket in buckets.values(): + if bucket: + scheduled.append(bucket.popleft()) + made_progress = True + if not made_progress: + raise RuntimeError("Failed to schedule all LIBERO-Plus tasks.") + return scheduled + + +def _ensure_libero_config(libero_plus_root: Path, output_dir: Path) -> Path: + config_dir = output_dir / "runtime" / "libero_config" + benchmark_root = libero_plus_root / "libero" / "libero" + payload = { + "benchmark_root": str(benchmark_root), + "bddl_files": str(benchmark_root / "bddl_files"), + "init_states": str(benchmark_root / "init_files"), + "datasets": str(output_dir / "runtime" / "libero_datasets"), + "assets": str(benchmark_root / "assets"), + } + _atomic_json_dump(payload, config_dir / "config.yaml") + Path(payload["datasets"]).mkdir(parents=True, exist_ok=True) + return config_dir + + +def _validate_runtime_inputs( + args: argparse.Namespace, + classification_path: Path, +) -> Path: + checkpoint = Path(args.checkpoint).expanduser().resolve() + checkpoint_load_path = _validate_checkpoint_load_path( + checkpoint, + args.checkpoint_load_path, + ) + dataset_stats = Path(args.dataset_stats).expanduser().resolve() + libero_plus_root = Path(args.libero_plus_root).expanduser().resolve() + config_dir = Path(args.config_dir).expanduser().resolve() + model_base_path = Path(args.model_base_path).expanduser().resolve() + required_paths = { + "checkpoint": checkpoint, + "checkpoint load path": checkpoint_load_path, + "dataset stats": dataset_stats, + "classification": classification_path, + "Hydra config dir": config_dir, + "model base path": model_base_path, + "LIBERO-Plus bddl_files": libero_plus_root / "libero" / "libero" / "bddl_files", + "LIBERO-Plus init_files": libero_plus_root / "libero" / "libero" / "init_files", + "LIBERO-Plus assets": libero_plus_root / "libero" / "libero" / "assets", + } + text_embedding_cache_dir = getattr(args, "text_embedding_cache_dir", None) + if text_embedding_cache_dir is not None: + required_paths["text embedding cache"] = ( + Path( + os.path.expanduser( + os.path.expandvars(str(text_embedding_cache_dir)) + ) + ).resolve() + ) + missing = [f"{label}: {path}" for label, path in required_paths.items() if not path.exists()] + if missing: + raise FileNotFoundError("Missing LIBERO-Plus evaluation inputs:\n- " + "\n- ".join(missing)) + if args.egl_fallback_gpu is not None and int(args.egl_fallback_gpu) < 0: + raise ValueError("--egl-fallback-gpu must be non-negative.") + return checkpoint_load_path + + +def _compose_worker_config( + *, + config_dir: str, + task_config: str, + checkpoint: str, + dataset_stats: str, + output_dir: str, + physical_gpu_id: int, + num_trials: int, + save_videos: bool, + gripper_action_format: str, + hydra_overrides: list[str], + env_seed: int | None = None, + policy_seed: int | None = None, + text_embedding_cache_dir: str | None = None, + worker_slot: str | None = None, +): + from hydra import compose, initialize_config_dir + from omegaconf import open_dict + + with initialize_config_dir(config_dir=config_dir, version_base="1.3"): + cfg = compose(config_name="sim_libero", overrides=[f"task={task_config}", *hydra_overrides]) + with open_dict(cfg): + cfg.ckpt = checkpoint + cfg.gpu_id = int(physical_gpu_id) + cfg.EVALUATION.output_dir = output_dir + cfg.EVALUATION.dataset_stats_path = dataset_stats + cfg.EVALUATION.num_trials = int(num_trials) + cfg.EVALUATION.device = "cuda" + cfg.EVALUATION.save_rollout_video = bool(save_videos) + cfg.EVALUATION.show_progress = False + cfg.EVALUATION.visualize_future_video = False + cfg.EVALUATION.gripper_action_format = gripper_action_format + if env_seed is not None: + cfg.EVALUATION.env_seed = int(env_seed) + if policy_seed is not None: + cfg.EVALUATION.policy_seed = int(policy_seed) + cfg.EVALUATION.worker_slot = ( + f"gpu{physical_gpu_id}-w0" + if worker_slot is None + else str(worker_slot) + ) + if text_embedding_cache_dir is not None: + cache_dir_value = str(text_embedding_cache_dir).strip() + if not cache_dir_value: + raise ValueError("text_embedding_cache_dir must not be empty.") + cache_dir = Path( + os.path.expanduser(os.path.expandvars(cache_dir_value)) + ).resolve() + cfg.EVALUATION.text_embedding_cache_dir = str(cache_dir) + cfg.model.load_text_encoder = False + return cfg + + +class _PersistentFastWAMEvaluator: + def __init__(self, cfg, output_dir: Path): + import logging + import random + + import numpy as np + import torch + from hydra.utils import instantiate + from omegaconf import open_dict + + from eval_libero_single import ( + CachedPromptEmbeddings, + _load_model_checkpoint, + _get_evaluation_seed, + _mixed_precision_to_model_dtype, + _resolve_dataset_stats_path, + _resolve_eval_device, + _validate_visualize_future_video_cfg, + run_single_task, + ) + from fastwam.datasets.lerobot.processors.fastwam_processor import FastWAMProcessor + from fastwam.datasets.lerobot.utils.normalizer import load_dataset_stats_from_json + from fastwam.utils.pytorch_utils import set_global_seed + from libero.libero import benchmark + + self.cfg = cfg + self.output_dir = output_dir + self.run_single_task = run_single_task + self.open_dict = open_dict + self.benchmark_dict = benchmark.get_benchmark_dict() + self.suite_cache = {} + + self.env_seed = _get_evaluation_seed(cfg, "env_seed") + self.policy_seed = _get_evaluation_seed(cfg, "policy_seed") + if self.policy_seed is not None: + set_global_seed(self.policy_seed, get_worker_init_fn=False) + _validate_visualize_future_video_cfg(cfg) + + self.model_device = _resolve_eval_device(cfg) + model_dtype = _mixed_precision_to_model_dtype(cfg.get("mixed_precision", "bf16")) + self.model = instantiate(cfg.model, model_dtype=model_dtype, device=self.model_device) + _load_model_checkpoint(self.model, str(cfg.ckpt)) + self.model = self.model.to(self.model_device).eval() + text_embedding_cache_dir = cfg.EVALUATION.get( + "text_embedding_cache_dir", None + ) + self.cached_prompt_embeddings = ( + None + if text_embedding_cache_dir is None + else CachedPromptEmbeddings( + text_embedding_cache_dir, + device=self.model_device, + context_len=int(cfg.model.get("tokenizer_max_len", 128)), + ) + ) + + dataset_stats_path = _resolve_dataset_stats_path(cfg) + dataset_stats = load_dataset_stats_from_json(str(dataset_stats_path)) + self.processor: FastWAMProcessor = instantiate(cfg.data.train.processor).eval() + self.processor.set_normalizer_from_stats(dataset_stats) + logging.info("Using dataset stats: %s", dataset_stats_path) + + action_horizon_cfg = cfg.EVALUATION.get("action_horizon", None) + self.action_horizon = ( + int(cfg.data.train.num_frames) - 1 + if action_horizon_cfg is None + else int(action_horizon_cfg) + ) + if self.action_horizon <= 0: + raise ValueError(f"EVALUATION.action_horizon must be positive, got {self.action_horizon}") + + video_size = cfg.data.train.get("video_size", [224, 224]) + if len(video_size) != 2: + raise ValueError(f"data.train.video_size must be [H, W], got {video_size}") + self.input_h = int(video_size[0]) + self.input_w = int(video_size[1]) + + # A fresh legacy evaluator starts every task from this same post-load RNG + # state. Restoring it keeps persistent workers reproducible across GPU + # assignment and --resume boundaries. + self.random_module = random + self.numpy_module = np + self.torch = torch + self.python_rng_state = random.getstate() + self.numpy_rng_state = np.random.get_state() + self.torch_cpu_rng_state = torch.random.get_rng_state() + self.torch_cuda_rng_state = torch.cuda.get_rng_state(self.model_device) + + def get_prompt_context(self, task_description: str): + if self.cached_prompt_embeddings is None: + return None + return self.cached_prompt_embeddings.get_prompt_context(task_description) + + def preload_prompt_contexts(self, task_descriptions) -> int: + if self.cached_prompt_embeddings is None: + return 0 + return self.cached_prompt_embeddings.preload_prompt_contexts( + task_descriptions + ) + + def _restore_post_load_rng_state(self) -> None: + self.random_module.setstate(self.python_rng_state) + self.numpy_module.random.set_state(self.numpy_rng_state) + self.torch.random.set_rng_state(self.torch_cpu_rng_state) + self.torch.cuda.set_rng_state( + self.torch_cuda_rng_state, + device=self.model_device, + ) + + def _get_suite(self, suite_name: str): + if suite_name not in self.suite_cache: + if suite_name not in self.benchmark_dict: + raise KeyError(f"Unknown LIBERO suite {suite_name!r}") + self.suite_cache[suite_name] = self.benchmark_dict[suite_name]() + return self.suite_cache[suite_name] + + def evaluate(self, task_spec: LiberoPlusTask) -> Path: + start_time = time.time() + self._restore_post_load_rng_state() + with self.open_dict(self.cfg): + self.cfg.EVALUATION.task_suite_name = task_spec.suite + self.cfg.EVALUATION.task_id = int(task_spec.task_id) + self.cfg.EVALUATION.category_value = task_spec.category + + task_suite = self._get_suite(task_spec.suite) + if not 0 <= task_spec.task_id < int(task_suite.n_tasks): + raise IndexError( + f"Task id {task_spec.task_id} is out of range for {task_spec.suite} " + f"(n_tasks={task_suite.n_tasks})." + ) + task = task_suite.get_task(task_spec.task_id) + if str(task.name) != task_spec.name: + raise ValueError( + f"Classification/benchmark task mismatch for {task_spec.key}: " + f"classification={task_spec.name!r}, benchmark={task.name!r}." + ) + initial_states = list(task_suite.get_task_init_states(task_spec.task_id)) + if not initial_states: + raise ValueError(f"No initial states found for {task_spec.key}.") + while len(initial_states) < int(self.cfg.EVALUATION.num_trials): + initial_states.extend( + initial_states[: int(self.cfg.EVALUATION.num_trials) - len(initial_states)] + ) + + task_video_dir = ( + self.output_dir + / "videos" + / task_spec.suite + / f"task_{task_spec.task_id:04d}" + ) + predicted_video_dir = ( + self.output_dir + / "predicted_videos" + / task_spec.suite + / f"task_{task_spec.task_id:04d}" + ) + if bool(self.cfg.EVALUATION.get("save_rollout_video", False)): + task_video_dir.mkdir(parents=True, exist_ok=True) + + results = { + "task_suite": task_spec.suite, + "task_id": task_spec.task_id, + "classification_id": task_spec.classification_id, + "classification_name": task_spec.name, + "category_value": task_spec.category, + "difficulty_level": task_spec.difficulty_level, + "env_seed": self.env_seed, + "policy_seed": self.policy_seed, + "task_description": None, + "successes": 0, + "total_episodes": int(self.cfg.EVALUATION.num_trials), + "gpu_id": int(self.cfg.gpu_id), + "worker_slot": str(self.cfg.EVALUATION.worker_slot), + "success_episodes": [], + "failure_episodes": [], + "start_time": datetime.now(timezone.utc).isoformat(), + "duration": 0.0, + } + task_results = self.run_single_task( + task=task, + initial_states=initial_states, + model=self.model, + processor=self.processor, + cfg=self.cfg, + video_dir=task_video_dir, + predicted_video_dir=predicted_video_dir, + action_horizon=self.action_horizon, + input_w=self.input_w, + input_h=self.input_h, + model_device=self.model_device, + prompt_context=self.get_prompt_context(str(task.language)), + ) + results.update(task_results) + results["duration"] = time.time() - start_time + path = result_path(self.output_dir, task_spec) + _atomic_json_dump(results, path) + failed_path = failure_path(self.output_dir, task_spec) + failed_path.unlink(missing_ok=True) + return path + + +def _worker_main( + worker_slot: str, + physical_gpu_id: int, + worker_args: dict, + task_queue, + status_queue, + start_event, + stop_event, +) -> None: + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(1, signal.SIGTERM) != 0: # PR_SET_PDEATHSIG + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + if os.getppid() != int(worker_args["parent_pid"]): + os._exit(1) + signal.signal(signal.SIGINT, signal.SIG_IGN) + + output_dir = Path(worker_args["output_dir"]) + log_path = output_dir / "worker_logs" / f"{worker_slot}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + log_handle = log_path.open("a", encoding="utf-8", buffering=1) + sys.stdout = log_handle + sys.stderr = log_handle + + visible_devices = str(physical_gpu_id) + egl_fallback_gpu = worker_args.get("egl_fallback_gpu") + if ( + egl_fallback_gpu is not None + and int(egl_fallback_gpu) != int(physical_gpu_id) + ): + visible_devices = f"{physical_gpu_id},{int(egl_fallback_gpu)}" + os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices + os.environ["MUJOCO_GL"] = "egl" + os.environ["PYOPENGL_PLATFORM"] = "egl" + # The requested physical GPU is local cuda:0 inside this worker. Hosts with + # an unusual display/EGL topology may expose one additional fallback GPU. + os.environ["MUJOCO_EGL_DEVICE_ID"] = "0" + os.environ["LIBERO_CONFIG_PATH"] = worker_args["libero_config_dir"] + os.environ["LIBERO_EGL_INIT_LOCK"] = "1" + egl_lock_file = worker_args["egl_lock_file"] + if worker_args["egl_lock_scope"] == "gpu": + egl_lock_file = f"{egl_lock_file}.gpu{physical_gpu_id}" + os.environ["LIBERO_EGL_INIT_LOCK_FILE"] = egl_lock_file + os.environ["LIBERO_EGL_INIT_RETRIES"] = "5" + os.environ["LIBERO_EGL_INIT_RETRY_SLEEP_SEC"] = "5" + os.environ["TOKENIZERS_PARALLELISM"] = "false" + os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1" + os.environ["DIFFSYNTH_MODEL_BASE_PATH"] = worker_args["model_base_path"] + + for path in ( + worker_args["libero_plus_root"], + str(PROJECT_ROOT), + str(PROJECT_ROOT / "src"), + str(Path(__file__).resolve().parent), + ): + if path not in sys.path: + sys.path.insert(0, path) + + if stop_event.is_set(): + status_queue.put( + { + "type": "stopped", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + } + ) + log_handle.close() + return + + try: + cfg = _compose_worker_config( + config_dir=worker_args["config_dir"], + task_config=worker_args["task_config"], + checkpoint=worker_args["checkpoint_load_path"], + dataset_stats=worker_args["dataset_stats"], + output_dir=worker_args["output_dir"], + physical_gpu_id=physical_gpu_id, + num_trials=worker_args["num_trials"], + save_videos=worker_args["save_videos"], + gripper_action_format=worker_args["gripper_action_format"], + hydra_overrides=worker_args["hydra_overrides"], + env_seed=worker_args.get("env_seed"), + policy_seed=worker_args.get("policy_seed"), + text_embedding_cache_dir=worker_args.get( + "text_embedding_cache_dir" + ), + worker_slot=worker_slot, + ) + evaluator = _PersistentFastWAMEvaluator(cfg, output_dir) + status_queue.put( + { + "type": "ready", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + } + ) + start_event.wait() + if stop_event.is_set(): + status_queue.put( + { + "type": "stopped", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + } + ) + log_handle.close() + return + except Exception: + error = traceback.format_exc() + print(error, flush=True) + status_queue.put( + { + "type": "fatal", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "error": error, + "log": str(log_path), + } + ) + log_handle.close() + return + + processed_tasks = 0 + while True: + if stop_event.is_set(): + break + payload = task_queue.get() + if payload is None or stop_event.is_set(): + break + task_spec = LiberoPlusTask.from_dict(payload) + task_started_at = time.time() + status_queue.put( + { + "type": "started", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "task": task_spec.key, + } + ) + try: + path = evaluator.evaluate(task_spec) + status_queue.put( + { + "type": "done", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "task": task_spec.key, + "result": str(path), + "duration": time.time() - task_started_at, + } + ) + except Exception: + error = traceback.format_exc() + print(error, flush=True) + failure = { + "task": task_spec.to_dict(), + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "time": datetime.now(timezone.utc).isoformat(), + "error": error, + "worker_log": str(log_path), + } + _atomic_json_dump(failure, failure_path(output_dir, task_spec)) + status_queue.put( + { + "type": "failed", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + "task": task_spec.key, + "error": error.splitlines()[-1] if error.splitlines() else error, + "log": str(log_path), + "duration": time.time() - task_started_at, + } + ) + + processed_tasks += 1 + max_tasks_per_worker = int(worker_args.get("max_tasks_per_worker", 0)) + if max_tasks_per_worker > 0 and processed_tasks >= max_tasks_per_worker: + break + + status_queue.put( + { + "type": "stopped", + "worker_slot": worker_slot, + "gpu_id": physical_gpu_id, + } + ) + log_handle.close() + + +def _build_manifest( + args: argparse.Namespace, + tasks: list[LiberoPlusTask], + classification_path: Path, + gpu_ids: list[int], +) -> dict: + text_embedding_cache_dir = getattr(args, "text_embedding_cache_dir", None) + return { + "schema_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + "mode": args.mode, + "checkpoint": str(Path(args.checkpoint).expanduser().resolve()), + "dataset_stats": str(Path(args.dataset_stats).expanduser().resolve()), + "libero_plus_root": str(Path(args.libero_plus_root).expanduser().resolve()), + "classification_path": str(classification_path), + "task_config": args.task_config, + "smoke_seed": None if args.mode == "full" else int(args.smoke_seed), + "env_seed": None if args.env_seed is None else int(args.env_seed), + "policy_seed": ( + None if args.policy_seed is None else int(args.policy_seed) + ), + "seed_fallback": "cfg.seed", + "category_filters": list(args.category_filters), + "num_trials": int(args.num_trials), + "gpu_ids": gpu_ids, + "save_videos": bool(args.save_videos), + "gripper_action_format": str(args.gripper_action_format), + "text_embedding_cache_dir": ( + None + if text_embedding_cache_dir is None + else str( + Path( + os.path.expanduser( + os.path.expandvars(str(text_embedding_cache_dir)) + ) + ).resolve() + ) + ), + "max_tasks_per_worker": int(args.max_tasks_per_worker), + "hydra_overrides": list(args.override), + "expected_tasks": len(tasks), + "tasks": [task.to_dict() for task in tasks], + } + + +def _write_or_validate_manifest(output_dir: Path, manifest: dict, *, resume: bool) -> None: + path = output_dir / "manifest.json" + if not path.exists(): + _atomic_json_dump(manifest, path) + return + with path.open(encoding="utf-8") as handle: + existing = json.load(handle) + compatibility_keys = ( + "mode", + "checkpoint", + "dataset_stats", + "classification_path", + "task_config", + "smoke_seed", + "env_seed", + "policy_seed", + "seed_fallback", + "category_filters", + "num_trials", + "save_videos", + "gripper_action_format", + "text_embedding_cache_dir", + "max_tasks_per_worker", + "hydra_overrides", + "tasks", + ) + legacy_defaults = { + "env_seed": None, + "policy_seed": None, + "seed_fallback": "cfg.seed", + "category_filters": [], + } + mismatches = [ + key + for key in compatibility_keys + if existing.get(key, legacy_defaults.get(key)) + != manifest.get(key, legacy_defaults.get(key)) + ] + if mismatches: + raise ValueError( + f"Existing manifest is incompatible with this run ({', '.join(mismatches)}): {path}" + ) + if not resume: + raise FileExistsError( + f"Output directory already contains a compatible run. Pass --resume or choose another path: {output_dir}" + ) + + +def _parse_gpu_ids(value: str) -> list[int]: + gpu_ids = [int(part.strip()) for part in value.split(",") if part.strip()] + if not gpu_ids: + raise ValueError("At least one GPU id is required.") + if len(set(gpu_ids)) != len(gpu_ids): + raise ValueError(f"GPU ids must be unique, got {gpu_ids}.") + return gpu_ids + + +def _normalize_category_filters(values: list[str] | None) -> list[str]: + if not values: + return [] + aliases = { + alias.casefold(): category + for category, alias in CATEGORY_LABELS.items() + } + aliases.update( + {category.casefold(): category for category in CATEGORY_ORDER} + ) + selected: set[str] = set() + unknown: list[str] = [] + for value in values: + for raw_name in str(value).split(","): + name = raw_name.strip() + if not name: + continue + category = aliases.get(name.casefold()) + if category is None: + unknown.append(name) + else: + selected.add(category) + if unknown: + choices = [ + f"{CATEGORY_LABELS[category]} ({category})" + for category in CATEGORY_ORDER + ] + raise ValueError( + f"Unknown Table 13 categories: {unknown}. " + f"Valid categories: {', '.join(choices)}." + ) + return [ + category for category in CATEGORY_ORDER if category in selected + ] + + +def run(args: argparse.Namespace) -> int: + output_dir = Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + lock_path = output_dir / ".evaluator.lock" + lock_handle = lock_path.open("a+") + try: + fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as error: + lock_handle.close() + raise RuntimeError( + f"Another evaluator already holds the output lock: {lock_path}" + ) from error + try: + lock_handle.seek(0) + lock_handle.truncate() + lock_handle.write(f"pid={os.getpid()}\n") + lock_handle.flush() + return _run_locked(args, output_dir) + finally: + fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN) + lock_handle.close() + + +def _run_locked(args: argparse.Namespace, output_dir: Path) -> int: + libero_plus_root = Path(args.libero_plus_root).expanduser().resolve() + classification_path = ( + Path(args.classification).expanduser().resolve() + if args.classification + else default_classification_path(libero_plus_root) + ) + all_tasks = load_task_classification(classification_path, require_full=True) + args.category_filters = _normalize_category_filters(args.category) + candidate_tasks = ( + all_tasks + if not args.category_filters + else [ + task + for task in all_tasks + if task.category in args.category_filters + ] + ) + tasks = ( + candidate_tasks + if args.mode == "full" + else select_smoke_tasks( + candidate_tasks, + limit=int(args.smoke_tasks), + seed=int(args.smoke_seed), + ) + ) + if ( + args.mode == "full" + and not args.category_filters + and len(tasks) != EXPECTED_TOTAL_TASKS + ): + raise ValueError(f"Full Table 13 requires {EXPECTED_TOTAL_TASKS} tasks, got {len(tasks)}.") + if int(args.num_trials) < 1: + raise ValueError(f"--num-trials must be positive; got {args.num_trials}.") + if args.mode == "full" and int(args.num_trials) != 1: + raise ValueError( + "LIBERO-Plus Table 13 full evaluation requires exactly one trial per task; " + f"got --num-trials={args.num_trials}. Multi-trial runs are allowed only in smoke mode." + ) + if args.mode == "full" and bool(args.save_videos): + raise ValueError("Full Table 13 evaluation does not save rollout videos.") + gpu_ids = _parse_gpu_ids(args.gpus) + worker_slots = _build_worker_slots(gpu_ids, int(args.workers_per_gpu)) + if int(args.max_tasks_per_worker) < 0: + raise ValueError("--max-tasks-per-worker must be non-negative (0 means unlimited).") + if args.mode == "full" and int(args.max_tasks_per_worker) != 0: + raise ValueError("Full evaluation requires --max-tasks-per-worker=0 (unlimited).") + if float(args.status_every) <= 0: + raise ValueError("--status-every must be positive.") + if float(args.worker_ready_timeout) <= 0: + raise ValueError("--worker-ready-timeout must be positive.") + worker_capacity = len(worker_slots) * int(args.max_tasks_per_worker) + if int(args.max_tasks_per_worker) > 0 and worker_capacity < len(tasks): + raise ValueError( + f"Worker capacity is too small: {len(worker_slots)} workers × " + f"{args.max_tasks_per_worker} tasks = {worker_capacity}, but {len(tasks)} tasks were selected." + ) + manifest = _build_manifest(args, tasks, classification_path, gpu_ids) + _write_or_validate_manifest(output_dir, manifest, resume=bool(args.resume)) + print( + f"Prepared LIBERO-Plus {args.mode} run: tasks={len(tasks)} " + f"num_trials={args.num_trials} " + f"smoke_seed={args.smoke_seed if args.mode == 'smoke' else 'n/a'} " + f"env_seed={args.env_seed if args.env_seed is not None else 'cfg.seed'} " + f"policy_seed={args.policy_seed if args.policy_seed is not None else 'cfg.seed'} " + f"categories={args.category_filters or 'all'} " + f"gpus={gpu_ids} workers_per_gpu={args.workers_per_gpu} " + f"workers={len(worker_slots)} egl_lock_scope={args.egl_lock_scope} " + f"checkpoint_load={args.checkpoint_load_path or args.checkpoint} " + f"output={output_dir}" + ) + if args.prepare_only: + print("prepare-only: manifest validated; no GPU workers started.") + return 0 + + text_embedding_cache_dir = getattr(args, "text_embedding_cache_dir", None) + if len(worker_slots) > len(gpu_ids) and not text_embedding_cache_dir: + raise ValueError( + "Multiple workers per GPU require --text-embedding-cache-dir; " + "loading one text encoder per worker exceeds the validated W8 " + "memory budget." + ) + checkpoint_load_path = _validate_runtime_inputs(args, classification_path) + libero_config_dir = _ensure_libero_config(libero_plus_root, output_dir) + os.environ["LIBERO_CONFIG_PATH"] = str(libero_config_dir) + if text_embedding_cache_dir: + cache_cfg = _compose_worker_config( + config_dir=str(Path(args.config_dir).expanduser().resolve()), + task_config=args.task_config, + checkpoint=str(checkpoint_load_path), + dataset_stats=str(Path(args.dataset_stats).expanduser().resolve()), + output_dir=str(output_dir), + physical_gpu_id=gpu_ids[0], + num_trials=int(args.num_trials), + save_videos=bool(args.save_videos), + gripper_action_format=str(args.gripper_action_format), + hydra_overrides=list(args.override), + env_seed=args.env_seed, + policy_seed=args.policy_seed, + text_embedding_cache_dir=text_embedding_cache_dir, + worker_slot=worker_slots[0][0], + ) + cache_context_len = int( + cache_cfg.model.get("tokenizer_max_len", WAN_PROMPT_CONTEXT_LEN) + ) + cache_report = _validate_text_embedding_cache( + text_embedding_cache_dir, + _load_exact_task_descriptions(tasks), + context_len=cache_context_len, + ) + print( + "Validated text embedding cache: " + f"tasks={cache_report['tasks']} " + f"unique_prompts={cache_report['unique_prompts']} " + f"context_len={cache_context_len} " + f"path={Path(text_embedding_cache_dir).expanduser().resolve()}", + flush=True, + ) + pending = [ + task + for task in tasks + if not ( + args.resume + and is_complete_result( + result_path(output_dir, task), + task, + num_trials=int(args.num_trials), + ) + ) + ] + already_complete = len(tasks) - len(pending) + print(f"Resume scan: complete={already_complete}, pending={len(pending)}") + if not pending: + from summarize_libero_plus import summarize + + summarize(output_dir, require_complete=True) + return 0 + + worker_args = { + "parent_pid": os.getpid(), + "output_dir": str(output_dir), + "config_dir": str(Path(args.config_dir).expanduser().resolve()), + "task_config": args.task_config, + "checkpoint": str(Path(args.checkpoint).expanduser().resolve()), + "checkpoint_load_path": str(checkpoint_load_path), + "dataset_stats": str(Path(args.dataset_stats).expanduser().resolve()), + "num_trials": int(args.num_trials), + "save_videos": bool(args.save_videos), + "gripper_action_format": str(args.gripper_action_format), + "text_embedding_cache_dir": ( + None + if text_embedding_cache_dir is None + else str( + Path( + os.path.expanduser( + os.path.expandvars(str(text_embedding_cache_dir)) + ) + ).resolve() + ) + ), + "max_tasks_per_worker": int(args.max_tasks_per_worker), + "hydra_overrides": list(args.override), + "env_seed": args.env_seed, + "policy_seed": args.policy_seed, + "libero_plus_root": str(libero_plus_root), + "libero_config_dir": str(libero_config_dir), + "egl_lock_file": os.environ.get( + "LIBERO_EGL_INIT_LOCK_FILE", + "/tmp/fastwam_libero_mujoco_egl_init.lock", + ), + "egl_lock_scope": args.egl_lock_scope, + "egl_fallback_gpu": args.egl_fallback_gpu, + "model_base_path": str(Path(args.model_base_path).expanduser().resolve()), + } + + context = mp.get_context("spawn") + task_queue = context.Queue() + # Avoid an exit-time hang if model initialization fails while the feeder + # thread still has a large full-run backlog buffered. + task_queue.cancel_join_thread() + status_queue = context.Queue() + start_event = context.Event() + stop_event = context.Event() + workers = [ + context.Process( + target=_worker_main, + args=( + worker_slot, + gpu_id, + worker_args, + task_queue, + status_queue, + start_event, + stop_event, + ), + name=f"libero-plus-{worker_slot}", + ) + for worker_slot, gpu_id in worker_slots + ] + process_by_slot = { + worker_slot: worker + for (worker_slot, _), worker in zip(worker_slots, workers, strict=True) + } + gpu_by_slot = dict(worker_slots) + + stop_requested = False + stop_deadline: float | None = None + previous_signal_handlers = { + signal.SIGINT: signal.getsignal(signal.SIGINT), + signal.SIGTERM: signal.getsignal(signal.SIGTERM), + } + + def request_stop(_signum, _frame) -> None: + nonlocal stop_requested, stop_deadline + if not stop_requested: + print( + "[interrupt] stopping after each worker's current task; " + "completed task results remain resumable.", + flush=True, + ) + stop_deadline = time.time() + 300.0 + stop_requested = True + stop_event.set() + start_event.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + for worker in workers: + worker.start() + workers_started_at = time.time() + for task in _interleave_tasks(pending): + task_queue.put(task.to_dict()) + for _ in workers: + task_queue.put(None) + + ready: set[str] = set() + stopped: set[str] = set() + completed = already_complete + finished_this_run = 0 + task_duration_sum = 0.0 + failures: list[dict] = [] + fatal_errors: list[dict] = [] + fatal_slots: set[str] = set() + last_status_time = 0.0 + rollout_started_at: float | None = None + while any(worker.is_alive() for worker in workers): + try: + event = status_queue.get(timeout=2.0) + except queue.Empty: + event = None + if event is not None: + event_type = event["type"] + if event_type == "ready": + ready.add(str(event["worker_slot"])) + print( + f"[worker-ready] slot={event['worker_slot']} " + f"gpu={event['gpu_id']} ready={len(ready)}/{len(workers)}", + flush=True, + ) + if len(ready) == len(workers): + start_event.set() + elif event_type == "started": + if rollout_started_at is None: + rollout_started_at = time.time() + elif event_type == "done": + completed += 1 + finished_this_run += 1 + task_duration_sum += float(event["duration"]) + elif event_type == "failed": + failures.append(event) + finished_this_run += 1 + task_duration_sum += float(event["duration"]) + print( + f"[task-failed] slot={event['worker_slot']} " + f"gpu={event['gpu_id']} task={event['task']} " + f"error={event['error']} log={event['log']}", + flush=True, + ) + elif event_type == "fatal": + worker_slot = str(event["worker_slot"]) + if worker_slot not in fatal_slots: + fatal_slots.add(worker_slot) + fatal_errors.append(event) + print( + f"[worker-fatal] slot={event['worker_slot']} " + f"gpu={event['gpu_id']} log={event['log']}\n" + f"{event['error']}", + flush=True, + ) + elif event_type == "stopped": + stopped.add(str(event["worker_slot"])) + + now = time.time() + if not start_event.is_set() and not stop_requested: + for worker_slot, worker in process_by_slot.items(): + if ( + worker_slot not in ready + and worker_slot not in fatal_slots + and worker.exitcode is not None + ): + fatal_slots.add(worker_slot) + fatal = { + "type": "fatal", + "worker_slot": worker_slot, + "gpu_id": gpu_by_slot[worker_slot], + "error": ( + "Worker exited before ready without a fatal event " + f"(exitcode={worker.exitcode})." + ), + "log": str( + output_dir + / "worker_logs" + / f"{worker_slot}.log" + ), + } + fatal_errors.append(fatal) + print( + f"[worker-fatal] slot={worker_slot} " + f"gpu={fatal['gpu_id']} log={fatal['log']}\n" + f"{fatal['error']}", + flush=True, + ) + if ( + not fatal_errors + and now - workers_started_at > float(args.worker_ready_timeout) + ): + waiting = sorted(set(process_by_slot) - ready) + fatal = { + "type": "fatal", + "worker_slot": "parent-timeout", + "gpu_id": -1, + "error": ( + "Timed out waiting for evaluation workers to become " + f"ready: {waiting}" + ), + "log": str(output_dir / "worker_logs"), + } + fatal_slots.add("parent-timeout") + fatal_errors.append(fatal) + print(f"[worker-fatal] {fatal['error']}", flush=True) + + for worker_slot, worker in process_by_slot.items(): + if ( + worker_slot in ready + and worker_slot not in fatal_slots + and worker.exitcode not in (None, 0) + ): + fatal_slots.add(worker_slot) + fatal = { + "type": "fatal", + "worker_slot": worker_slot, + "gpu_id": gpu_by_slot[worker_slot], + "error": ( + "Worker exited unexpectedly after ready " + f"(exitcode={worker.exitcode})." + ), + "log": str( + output_dir / "worker_logs" / f"{worker_slot}.log" + ), + } + fatal_errors.append(fatal) + print( + f"[worker-fatal] slot={worker_slot} " + f"gpu={fatal['gpu_id']} log={fatal['log']}\n" + f"{fatal['error']}", + flush=True, + ) + + if stop_deadline is not None and now >= stop_deadline: + print( + "[interrupt] graceful-stop deadline reached; terminating " + "remaining workers.", + flush=True, + ) + break + if now - last_status_time >= float(args.status_every): + alive = sum(worker.is_alive() for worker in workers) + average_task_seconds = ( + None + if finished_this_run == 0 + else task_duration_sum / finished_this_run + ) + remaining = max(0, len(pending) - finished_this_run) + eta_seconds = ( + None + if finished_this_run == 0 or rollout_started_at is None + else ( + (now - rollout_started_at) + * remaining + / finished_this_run + ) + ) + status_label = "status" if start_event.is_set() else "startup" + print( + f"[{status_label}] completed={completed}/{len(tasks)} " + f"failed={len(failures)} ready={len(ready)}/{len(workers)} " + f"alive={alive} avg_task={_format_duration(average_task_seconds)} " + f"eta={_format_duration(eta_seconds)}", + flush=True, + ) + last_status_time = now + if fatal_errors: + stop_event.set() + start_event.set() + for worker in workers: + if worker.is_alive(): + worker.terminate() + break + + start_event.set() + shutdown_deadline = ( + stop_deadline + if stop_deadline is not None + else time.time() + 30.0 + ) + while any(worker.is_alive() for worker in workers): + if time.time() >= shutdown_deadline: + break + for worker in workers: + if worker.is_alive(): + worker.join(timeout=0.2) + for worker in workers: + if worker.is_alive(): + worker.terminate() + worker.join(timeout=5) + if worker.is_alive(): + worker.kill() + worker.join(timeout=5) + + while True: + try: + event = status_queue.get_nowait() + except queue.Empty: + break + if event["type"] == "done": + completed += 1 + finished_this_run += 1 + task_duration_sum += float(event["duration"]) + elif event["type"] == "failed": + failures.append(event) + finished_this_run += 1 + task_duration_sum += float(event["duration"]) + elif event["type"] == "fatal": + worker_slot = str(event["worker_slot"]) + if worker_slot not in fatal_slots: + fatal_slots.add(worker_slot) + fatal_errors.append(event) + + for signum, previous_handler in previous_signal_handlers.items(): + signal.signal(signum, previous_handler) + if stop_requested: + print( + "[interrupt] workers stopped cleanly; rerun the same command " + "with --resume to continue.", + flush=True, + ) + return 130 + + from summarize_libero_plus import summarize + + report = summarize(output_dir, require_complete=False) + if fatal_errors or failures or not report["is_complete"]: + raise RuntimeError( + f"LIBERO-Plus evaluation incomplete: fatal_workers={len(fatal_errors)}, " + f"failed_tasks={len(failures)}, completed={report['completed_tasks']}/{report['expected_tasks']}. " + f"Rerun the same command with --resume after fixing the reported errors." + ) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", required=True, help="LightX2V/FastWAM fastwam.pt checkpoint.") + parser.add_argument( + "--checkpoint-load-path", + default=None, + help=( + "Optional local byte-identical checkpoint used by workers for " + "faster loading. Size and SHA256 are verified; --checkpoint " + "remains the manifest identity." + ), + ) + parser.add_argument("--dataset-stats", required=True, help="LIBERO-Plus dataset_stats.json.") + parser.add_argument( + "--text-embedding-cache-dir", + default=None, + help=( + "Optional Wan2.2 text embedding cache. When set, prompt contexts " + "are cached on each worker GPU and the text encoder is not loaded." + ), + ) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--libero-plus-root", default=str(DEFAULT_LIBERO_PLUS_ROOT)) + parser.add_argument("--classification", default=None) + parser.add_argument("--config-dir", default=str(DEFAULT_CONFIG_DIR)) + parser.add_argument("--task-config", default="libero_uncond_2cam224_1e-4") + parser.add_argument( + "--model-base-path", + required=True, + help="Directory containing the Wan model files referenced by the Hydra config.", + ) + parser.add_argument("--gpus", default="0,1,2,3,4,5,6,7") + parser.add_argument( + "--workers-per-gpu", + type=int, + default=1, + help="Number of independent persistent evaluator processes per GPU.", + ) + parser.add_argument( + "--egl-lock-scope", + choices=("global", "gpu"), + default="global", + help=( + "Serialize MuJoCo EGL environment construction globally or only " + "between workers sharing one physical GPU." + ), + ) + parser.add_argument( + "--egl-fallback-gpu", + type=int, + default=None, + help=( + "Optional extra physical GPU appended to CUDA_VISIBLE_DEVICES for " + "hosts whose EGL stack requires a separate display GPU. Normally unset." + ), + ) + parser.add_argument("--mode", choices=("smoke", "full"), default="smoke") + parser.add_argument("--smoke-tasks", type=int, default=14) + parser.add_argument( + "--smoke-seed", + type=int, + default=42, + help="Random seed used to sample smoke tasks from all LIBERO-Plus variants.", + ) + parser.add_argument( + "--env-seed", + type=int, + default=None, + help=( + "Environment-construction seed. Defaults to cfg.seed, preserving " + "the historical evaluator behavior." + ), + ) + parser.add_argument( + "--policy-seed", + type=int, + default=None, + help=( + "FastWAM inference seed. Defaults to cfg.seed, preserving the " + "historical evaluator behavior." + ), + ) + parser.add_argument( + "--category", + action="append", + default=[], + help=( + "Restrict evaluation to a Table 13 category (repeatable or " + "comma-separated). Canonical names and short labels such as " + "'Robot' are accepted." + ), + ) + parser.add_argument("--num-trials", type=int, default=1) + parser.add_argument("--save-videos", action="store_true") + parser.add_argument( + "--gripper-action-format", + choices=("zero_one_open_positive", "signed_open_negative"), + default="signed_open_negative", + help="Raw gripper convention after de-normalization; LIBERO-Plus uses signed_open_negative.", + ) + parser.add_argument( + "--max-tasks-per-worker", + type=int, + default=0, + help="Stop each worker after this many tasks; 0 is unlimited. Useful for queue-validation smoke runs.", + ) + parser.add_argument("--resume", action="store_true") + parser.add_argument("--prepare-only", action="store_true") + parser.add_argument("--status-every", type=float, default=30.0) + parser.add_argument( + "--worker-ready-timeout", + type=float, + default=1200.0, + help="Fail instead of waiting forever if all evaluator workers are not ready.", + ) + parser.add_argument( + "--override", + action="append", + default=[], + help="Additional Hydra override forwarded to every persistent worker (repeatable).", + ) + return parser + + +def main() -> None: + args = build_parser().parse_args() + raise SystemExit(run(args)) + + +if __name__ == "__main__": + main() diff --git a/experiments/libero/eval_libero_single.py b/experiments/libero/eval_libero_single.py index af558a4f..bc3b17b4 100644 --- a/experiments/libero/eval_libero_single.py +++ b/experiments/libero/eval_libero_single.py @@ -1,9 +1,11 @@ -import json +import hashlib import inspect +import json import logging import os import sys import time +from collections.abc import Iterable from pathlib import Path from typing import Any, Optional @@ -49,6 +51,142 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false" +PromptContext = tuple[torch.Tensor, torch.Tensor] + + +class CachedPromptEmbeddings: + """Strict loader for Wan2.2 prompt embeddings resident on the model device.""" + + ENCODER_ID = "wan22ti2v5b" + EMBEDDING_DIM = 4096 + + def __init__( + self, + cache_dir: str | Path, + *, + device: str | torch.device, + context_len: int = 128, + ): + expanded_cache_dir = os.path.expanduser(os.path.expandvars(str(cache_dir))) + self.cache_dir = Path(expanded_cache_dir).resolve() + if not self.cache_dir.is_dir(): + raise NotADirectoryError( + f"Text embedding cache directory does not exist: {self.cache_dir}" + ) + self.device = torch.device(device) + self.context_len = int(context_len) + if self.context_len <= 0: + raise ValueError(f"context_len must be positive, got {self.context_len}.") + self._device_cache: dict[str, PromptContext] = {} + self._device_cache_bytes = 0 + + @staticmethod + def _full_prompt(task_description: str) -> str: + if not isinstance(task_description, str) or not task_description: + raise ValueError("task_description must be a non-empty string.") + return DEFAULT_PROMPT.format(task=task_description) + + def _cache_path(self, full_prompt: str) -> Path: + digest = hashlib.sha256(full_prompt.encode("utf-8")).hexdigest() + return self.cache_dir / ( + f"{digest}.t5_len{self.context_len}.{self.ENCODER_ID}.pt" + ) + + def _load_prompt_context(self, full_prompt: str) -> PromptContext: + cache_path = self._cache_path(full_prompt) + if not cache_path.is_file(): + raise FileNotFoundError( + "Missing cached prompt embedding for the exact full prompt: " + f"{cache_path} (prompt={full_prompt!r})" + ) + try: + payload = torch.load(cache_path, map_location="cpu", weights_only=True) + except Exception as error: + raise RuntimeError(f"Failed to load text embedding cache: {cache_path}") from error + if not isinstance(payload, dict): + raise TypeError( + f"Cached prompt payload must be a dict, got {type(payload).__name__}: " + f"{cache_path}" + ) + expected_keys = {"context", "mask"} + actual_keys = set(payload) + if actual_keys != expected_keys: + raise ValueError( + f"Cached prompt payload keys must be exactly {sorted(expected_keys)}, " + f"got {sorted(actual_keys)}: {cache_path}" + ) + + context = payload["context"] + context_mask = payload["mask"] + if not isinstance(context, torch.Tensor) or not isinstance( + context_mask, torch.Tensor + ): + raise TypeError( + "Cached `context` and `mask` must both be torch.Tensor instances: " + f"{cache_path}" + ) + expected_context_shape = (self.context_len, self.EMBEDDING_DIM) + expected_mask_shape = (self.context_len,) + if tuple(context.shape) != expected_context_shape: + raise ValueError( + f"Cached context shape must be {expected_context_shape}, " + f"got {tuple(context.shape)}: {cache_path}" + ) + if tuple(context_mask.shape) != expected_mask_shape: + raise ValueError( + f"Cached mask shape must be {expected_mask_shape}, " + f"got {tuple(context_mask.shape)}: {cache_path}" + ) + if context.dtype != torch.bfloat16: + raise TypeError( + f"Cached context dtype must be torch.bfloat16, got {context.dtype}: " + f"{cache_path}" + ) + if context_mask.dtype != torch.bool: + raise TypeError( + f"Cached mask dtype must be torch.bool, got {context_mask.dtype}: " + f"{cache_path}" + ) + if not bool(torch.isfinite(context).all().item()): + raise ValueError(f"Cached context contains non-finite values: {cache_path}") + + # Match FastWAM.encode_prompt and the training dataset exactly: encoder + # outputs outside the tokenizer mask are zeroed, then cross-attention + # receives an all-valid mask. + context = context.contiguous().clone() + context[~context_mask] = 0 + context_mask = torch.ones_like(context_mask) + context = context.to(device=self.device, non_blocking=True) + context_mask = context_mask.to(device=self.device, non_blocking=True) + return context, context_mask + + def get_prompt_context(self, task_description: str) -> PromptContext: + full_prompt = self._full_prompt(task_description) + cached = self._device_cache.get(full_prompt) + if cached is not None: + return cached + + context, context_mask = self._load_prompt_context(full_prompt) + self._device_cache[full_prompt] = (context, context_mask) + self._device_cache_bytes += ( + context.numel() * context.element_size() + + context_mask.numel() * context_mask.element_size() + ) + logging.info( + "Cached prompt embedding on %s: prompts=%d memory=%.2f MiB", + self.device, + len(self._device_cache), + self._device_cache_bytes / (1024 * 1024), + ) + return context, context_mask + + def preload_prompt_contexts(self, task_descriptions: Iterable[str]) -> int: + before = len(self._device_cache) + for task_description in dict.fromkeys(task_descriptions): + self.get_prompt_context(task_description) + return len(self._device_cache) - before + + class NumpyEncoder(json.JSONEncoder): def default(self, obj): if isinstance(obj, np.integer): @@ -117,39 +255,6 @@ def _resolve_dataset_stats_path(cfg: DictConfig) -> Path: def _load_model_checkpoint(model: torch.nn.Module, ckpt: str) -> None: model.load_checkpoint(ckpt) logging.info("Loaded checkpoint via model.load_checkpoint: %s", ckpt) - return - - # deprecated legacy checkpoint loading - payload = torch.load(ckpt, map_location="cpu") - if not isinstance(payload, dict): - raise ValueError(f"Legacy checkpoint payload must be dict, got: {type(payload)}") - - if "mot" in payload and hasattr(model, "mot"): - missing, unexpected = model.mot.load_state_dict(payload["mot"], strict=False) - logging.warning( - "Loaded fallback `mot` state_dict with strict=False. Missing=%d Unexpected=%d", - len(missing), - len(unexpected), - ) - return - - state_dict = None - for key in ("model_state_dict", "state_dict", "model"): - value = payload.get(key) - if isinstance(value, dict): - state_dict = value - break - if state_dict is None and all(torch.is_tensor(v) for v in payload.values()): - state_dict = payload - if state_dict is None: - raise ValueError(f"Cannot parse legacy checkpoint keys from: {ckpt}") - - missing, unexpected = model.load_state_dict(state_dict, strict=False) - logging.warning( - "Loaded fallback model state_dict with strict=False. Missing=%d Unexpected=%d", - len(missing), - len(unexpected), - ) def _center_crop_resize(image: np.ndarray, width: int, height: int) -> np.ndarray: @@ -275,6 +380,31 @@ def _denormalize_action(action: torch.Tensor, processor: FastWAMProcessor) -> np return denorm.numpy() +def _postprocess_gripper_action(action: np.ndarray, cfg: DictConfig) -> np.ndarray: + """Convert the dataset's raw gripper convention to LIBERO's signed convention.""" + gripper_action_format = str( + cfg.EVALUATION.get("gripper_action_format", "zero_one_open_positive") + ) + if gripper_action_format == "zero_one_open_positive": + # Released FastWAM LIBERO data stores 0=close, 1=open. LIBERO expects + # -1=open, +1=close. + action[..., -1] = action[..., -1] * 2 - 1 + action = invert_gripper_action(action) + elif gripper_action_format == "signed_open_negative": + # LIBERO-Plus already stores -1=open, +1=close, matching the simulator. + pass + else: + raise ValueError( + "Unsupported EVALUATION.gripper_action_format=" + f"{gripper_action_format!r}; expected 'zero_one_open_positive' " + "or 'signed_open_negative'." + ) + + if bool(cfg.EVALUATION.get("binarize_gripper", False)): + action[..., -1] = np.sign(action[..., -1]) + return action + + def _get_num_video_frames(cfg: DictConfig) -> int: return (int(cfg.data.train.num_frames) - 1) // int(cfg.data.train.action_video_freq_ratio) + 1 @@ -367,6 +497,7 @@ def _predict_action_chunk( input_w: int, input_h: int, model_device: str, + prompt_context: Optional[PromptContext] = None, ) -> tuple[np.ndarray, dict, Optional[list[Image.Image]]]: num_inference_steps_cfg = cfg.EVALUATION.get("num_inference_steps", None) if num_inference_steps_cfg is None: @@ -387,7 +518,6 @@ def _predict_action_chunk( ) infer_kwargs = { - "prompt": prompt, "input_image": image, "action_horizon": action_horizon, "negative_prompt": str(cfg.EVALUATION.get("negative_prompt", "")), @@ -399,10 +529,19 @@ def _predict_action_chunk( if cfg.EVALUATION.get("sigma_shift") is None else float(cfg.EVALUATION.get("sigma_shift")) ), - "seed": None if cfg.get("seed") is None else int(cfg.seed), + "seed": _get_evaluation_seed(cfg, "policy_seed"), "rand_device": str(cfg.EVALUATION.get("rand_device", "cpu")), "tiled": bool(cfg.EVALUATION.get("tiled", False)), } + if prompt_context is None: + infer_kwargs["prompt"] = prompt + else: + context, context_mask = prompt_context + infer_kwargs.update( + prompt=None, + context=context, + context_mask=context_mask, + ) visualize_future_video = bool(cfg.EVALUATION.get("visualize_future_video", False)) predicted_future_frames = None if visualize_future_video: @@ -420,12 +559,7 @@ def _predict_action_chunk( action = _denormalize_action(action, processor)[0] # [T, D] - # The dataloader flips the sign of the gripper action to align with other datasets - # (0 = close, 1 = open), so flip it back (-1 = open, +1 = close) before executing the action - action[..., -1] = action[..., -1] * 2 - 1 - action = invert_gripper_action(action) - if bool(cfg.EVALUATION.get("binarize_gripper", False)): - action[..., -1] = np.sign(action[..., -1]) + action = _postprocess_gripper_action(action, cfg) return action, imgs, predicted_future_frames @@ -442,6 +576,14 @@ def _get_max_steps(task_suite_name: str) -> int: return suite_steps[task_suite_name] +def _get_evaluation_seed(cfg: DictConfig, name: str) -> Optional[int]: + """Resolve an evaluation seed while preserving the legacy cfg.seed fallback.""" + value = cfg.EVALUATION.get(name, None) + if value is None: + value = cfg.get("seed", None) + return None if value is None else int(value) + + def run_single_episode( env, initial_state, @@ -455,12 +597,16 @@ def run_single_episode( input_w: int, input_h: int, model_device: str, + prompt_context: Optional[PromptContext] = None, ) -> tuple[bool, list, list[dict[str, Any]], Optional[float]]: max_steps = _get_max_steps(cfg.EVALUATION.task_suite_name) replan_steps = int(cfg.EVALUATION.get("replan_steps", 5)) num_steps_wait = int(cfg.EVALUATION.get("num_steps_wait", 5)) use_action_ensembler = bool(cfg.EVALUATION.get("use_action_ensembler", False)) visualize_future_video = bool(cfg.EVALUATION.get("visualize_future_video", False)) + collect_replay = visualize_future_video or bool( + cfg.EVALUATION.get("save_rollout_video", True) + ) capture_steps = set(_get_future_frame_capture_steps(cfg)[1:]) env.reset() @@ -479,7 +625,11 @@ def run_single_episode( t = 0 done = False - pbar = tqdm(total=max_steps + num_steps_wait, desc=f"Episode {episode_idx + 1}") + pbar = tqdm( + total=max_steps + num_steps_wait, + desc=f"Episode {episode_idx + 1}", + disable=not bool(cfg.EVALUATION.get("show_progress", True)), + ) while t < max_steps + num_steps_wait: pbar.update(1) if t < num_steps_wait: @@ -498,6 +648,7 @@ def run_single_episode( input_w=input_w, input_h=input_h, model_device=model_device, + prompt_context=prompt_context, ) if predicted_future_frames is not None: current_replan_idx += 1 @@ -514,10 +665,12 @@ def run_single_episode( pending_actions = [ensembler.get_action(ts).tolist() for ts in range(t, t + replan_steps)] else: pending_actions = action_chunk[:replan_steps].tolist() - replay_images.append(imgs.copy()) + if collect_replay: + replay_images.append(imgs.copy()) else: imgs = get_libero_image(obs) - replay_images.append(imgs.copy()) + if collect_replay: + replay_images.append(imgs.copy()) obs, _, done, _ = env.step(pending_actions.pop(0)) if visualize_future_video and current_predicted_future_clip is not None: @@ -594,85 +747,104 @@ def run_single_task( input_w: int, input_h: int, model_device: str, + prompt_context: Optional[PromptContext] = None, ) -> dict: - env, task_description = get_libero_env(task, LIBERO_ENV_RESOLUTION, cfg.get("seed")) - visualize_future_video = bool(cfg.EVALUATION.get("visualize_future_video", False)) - results = { - "successes": 0, - "failure_episodes": [], - "success_episodes": [], - "task_description": task_description, - } - if visualize_future_video: - results["episode_future_video_psnr"] = [] - results["future_video_psnr_mean"] = None - - for trial_idx in range(int(cfg.EVALUATION.num_trials)): - success, replay_images, predicted_future_video_clips, episode_mean_psnr = run_single_episode( - env=env, - initial_state=initial_states[trial_idx], - task_description=task_description, - model=model, - processor=processor, - cfg=cfg, - episode_idx=trial_idx, - action_horizon=action_horizon, - input_w=input_w, - input_h=input_h, - model_device=model_device, - ) - if success: - results["successes"] += 1 - results["success_episodes"].append(trial_idx) - else: - results["failure_episodes"].append(trial_idx) - if visualize_future_video: - results["episode_future_video_psnr"].append(episode_mean_psnr) - - save_rollout_video( - video_dir, - replay_images, - f"task{cfg.EVALUATION.task_id}_trial{trial_idx}", - success=success, - task_description=task_description, - ) + env_setup_started_at = time.time() + env, task_description = get_libero_env( + task, + LIBERO_ENV_RESOLUTION, + _get_evaluation_seed(cfg, "env_seed"), + ) + env_setup_duration = time.time() - env_setup_started_at + try: + visualize_future_video = bool(cfg.EVALUATION.get("visualize_future_video", False)) + save_rollouts = bool(cfg.EVALUATION.get("save_rollout_video", True)) + results = { + "successes": 0, + "failure_episodes": [], + "success_episodes": [], + "task_description": task_description, + "env_setup_duration": env_setup_duration, + "episode_durations": [], + } if visualize_future_video: - if len(predicted_future_video_clips) == 0: - logging.warning( - "No predicted future frames collected for task %s trial %s.", - cfg.EVALUATION.task_id, - trial_idx, - ) + results["episode_future_video_psnr"] = [] + results["future_video_psnr_mean"] = None + + for trial_idx in range(int(cfg.EVALUATION.num_trials)): + episode_started_at = time.time() + success, replay_images, predicted_future_video_clips, episode_mean_psnr = run_single_episode( + env=env, + initial_state=initial_states[trial_idx], + task_description=task_description, + model=model, + processor=processor, + cfg=cfg, + episode_idx=trial_idx, + action_horizon=action_horizon, + input_w=input_w, + input_h=input_h, + model_device=model_device, + prompt_context=prompt_context, + ) + results["episode_durations"].append(time.time() - episode_started_at) + if success: + results["successes"] += 1 + results["success_episodes"].append(trial_idx) else: - all_gt_frames = [] - all_pred_frames = [] - for clip in predicted_future_video_clips: - all_gt_frames.extend(clip["gt_frames"]) - all_pred_frames.extend(clip["pred_frames"]) + results["failure_episodes"].append(trial_idx) + if visualize_future_video: + results["episode_future_video_psnr"].append(episode_mean_psnr) + + if save_rollouts: + save_rollout_video( + video_dir, + replay_images, + f"task{cfg.EVALUATION.task_id}_trial{trial_idx}", + success=success, + task_description=task_description, + ) + if visualize_future_video: + if len(predicted_future_video_clips) == 0: + logging.warning( + "No predicted future frames collected for task %s trial %s.", + cfg.EVALUATION.task_id, + trial_idx, + ) + else: + all_gt_frames = [] + all_pred_frames = [] + for clip in predicted_future_video_clips: + all_gt_frames.extend(clip["gt_frames"]) + all_pred_frames.extend(clip["pred_frames"]) + save_prediction_video( + predicted_video_dir, + clip["gt_frames"], + clip["pred_frames"], + f"task{cfg.EVALUATION.task_id}_trial{trial_idx}", + clip["replan_idx"], + success=success, + task_description=task_description, + ) save_prediction_video( predicted_video_dir, - clip["gt_frames"], - clip["pred_frames"], + all_gt_frames, + all_pred_frames, f"task{cfg.EVALUATION.task_id}_trial{trial_idx}", - clip["replan_idx"], + "all", success=success, task_description=task_description, ) - save_prediction_video( - predicted_video_dir, - all_gt_frames, - all_pred_frames, - f"task{cfg.EVALUATION.task_id}_trial{trial_idx}", - "all", - success=success, - task_description=task_description, - ) - if visualize_future_video: - valid_episode_psnr = [x for x in results["episode_future_video_psnr"] if x is not None] - if len(valid_episode_psnr) > 0: - results["future_video_psnr_mean"] = float(np.mean(valid_episode_psnr)) - return results + if visualize_future_video: + valid_episode_psnr = [x for x in results["episode_future_video_psnr"] if x is not None] + if len(valid_episode_psnr) > 0: + results["future_video_psnr_mean"] = float(np.mean(valid_episode_psnr)) + return results + finally: + close = getattr(env, "close", None) + if callable(close): + close() @hydra.main(version_base="1.3", config_path="../../configs", config_name="sim_libero.yaml") @@ -681,8 +853,9 @@ def eval_single_process(cfg: DictConfig): partial_state = PartialState() partial_state.config = cfg - if cfg.get("seed") is not None: - set_global_seed(int(cfg.seed), get_worker_init_fn=False) + policy_seed = _get_evaluation_seed(cfg, "policy_seed") + if policy_seed is not None: + set_global_seed(policy_seed, get_worker_init_fn=False) if cfg.ckpt is None: raise ValueError("cfg.ckpt must not be None.") @@ -720,8 +893,6 @@ def eval_single_process(cfg: DictConfig): raise ValueError(f"data.train.video_size must be [H, W], got {video_size}") input_h = int(video_size[0]) input_w = int(video_size[1]) - concat_multi_camera = cfg.data.train.get("concat_multi_camera", None) - shape_meta_images = [meta["shape"] for meta in processor.shape_meta["images"]] local_log_dir = Path(cfg.EVALUATION.output_dir) local_log_dir.mkdir(parents=True, exist_ok=True) diff --git a/experiments/libero/lerobot_rollout_writer.py b/experiments/libero/lerobot_rollout_writer.py new file mode 100644 index 00000000..6151b117 --- /dev/null +++ b/experiments/libero/lerobot_rollout_writer.py @@ -0,0 +1,1174 @@ +#!/usr/bin/env python3 +"""Crash-safe writer for successful LIBERO rollouts in LeRobot v2.1 format. + +Workers call :func:`stage_success` after an episode succeeds. The main process +then calls :func:`finalize_dataset` to deterministically assemble all completed +staging directories into a standalone dataset. Failed/incomplete worker +directories are never considered by the finalizer. + +The public API deliberately keeps collection and publication separate:: + + stage_success(staging_root, job_id, trajectory, provenance) + summary = finalize_dataset(staging_root, dataset_root, + reference_dataset=base_dataset) + report = validate_dataset(dataset_root, reference_dataset=base_dataset) + +``trajectory`` must contain ``front``, ``wrist``, ``state``, and ``action``. +``provenance`` must contain the exact training prompt in ``task``. Every value +in ``provenance`` is preserved in ``meta/provenance.jsonl``. +""" + +from __future__ import annotations + +import ctypes +import errno +import hashlib +import json +import math +import os +import re +import shutil +import tempfile +import uuid +from pathlib import Path +from typing import Any, Mapping + +import av +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + + +FPS = 20 +IMAGE_HEIGHT = 256 +IMAGE_WIDTH = 256 +CHUNKS_SIZE = 1000 +STAGING_VERSION = 1 +CODEBASE_VERSION = "v2.1" + +FRONT_KEY = "observation.images.front" +WRIST_KEY = "observation.images.wrist" +STATE_KEY = "observation.state" +ACTION_KEY = "action" +VIDEO_KEYS = (FRONT_KEY, WRIST_KEY) + +DATA_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" +VIDEO_PATH = ( + "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4" +) + +_JOB_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,191}$") +_STATE_TYPE = pa.list_(pa.field("element", pa.float32()), 8) +_ACTION_TYPE = pa.list_(pa.field("element", pa.float32()), 7) +PARQUET_SCHEMA = pa.schema( + [ + pa.field(STATE_KEY, _STATE_TYPE), + pa.field(ACTION_KEY, _ACTION_TYPE), + pa.field("timestamp", pa.float32()), + pa.field("frame_index", pa.int64()), + pa.field("episode_index", pa.int64()), + pa.field("index", pa.int64()), + pa.field("task_index", pa.int64()), + ], + metadata={ + b"huggingface": json.dumps( + { + "info": { + "features": { + STATE_KEY: { + "feature": {"dtype": "float32", "_type": "Value"}, + "length": 8, + "_type": "Sequence", + }, + ACTION_KEY: { + "feature": {"dtype": "float32", "_type": "Value"}, + "length": 7, + "_type": "Sequence", + }, + "timestamp": {"dtype": "float32", "_type": "Value"}, + "frame_index": {"dtype": "int64", "_type": "Value"}, + "episode_index": {"dtype": "int64", "_type": "Value"}, + "index": {"dtype": "int64", "_type": "Value"}, + "task_index": {"dtype": "int64", "_type": "Value"}, + } + } + } + ).encode("utf-8") + }, +) + + +def _json_default(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, np.generic): + return value.item() + if isinstance(value, np.ndarray): + return value.tolist() + raise TypeError(f"{type(value).__name__} is not JSON serializable") + + +def _normalize_json(value: Any) -> Any: + """Return a JSON-safe copy and reject NaN/Infinity early.""" + encoded = json.dumps( + value, + default=_json_default, + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ) + return json.loads(encoded) + + +def _canonical_json(value: Any) -> bytes: + return json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + + +def _write_json(path: Path, value: Any, *, indent: int | None = 2) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + json.dump( + value, + handle, + ensure_ascii=False, + allow_nan=False, + indent=indent, + sort_keys=False, + ) + handle.write("\n") + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write( + json.dumps( + row, + ensure_ascii=False, + allow_nan=False, + sort_keys=False, + separators=(",", ":"), + ) + ) + handle.write("\n") + + +def _read_json(path: Path) -> Any: + with path.open(encoding="utf-8") as handle: + return json.load(handle) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + with path.open(encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(4 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_job_id(job_id: str) -> str: + if not isinstance(job_id, str) or not _JOB_ID_RE.fullmatch(job_id): + raise ValueError( + f"job_id must match [A-Za-z0-9][A-Za-z0-9_.-]{{0,191}}; got {job_id!r}" + ) + return job_id + + +def _validated_trajectory( + trajectory: Mapping[str, np.ndarray], +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + missing = [ + key for key in ("front", "wrist", "state", "action") if key not in trajectory + ] + if missing: + raise ValueError(f"trajectory is missing fields: {missing}") + + front = np.asarray(trajectory["front"]) + wrist = np.asarray(trajectory["wrist"]) + state = np.asarray(trajectory["state"]) + action = np.asarray(trajectory["action"]) + + expected_images = (None, IMAGE_HEIGHT, IMAGE_WIDTH, 3) + for name, frames in (("front", front), ("wrist", wrist)): + if frames.ndim != 4 or frames.shape[1:] != expected_images[1:]: + raise ValueError( + f"{name} must have shape [N, {IMAGE_HEIGHT}, {IMAGE_WIDTH}, 3], " + f"got {frames.shape}" + ) + if frames.dtype != np.uint8: + raise TypeError(f"{name} must be uint8, got {frames.dtype}") + + if state.ndim != 2 or state.shape[1] != 8: + raise ValueError(f"state must have shape [N, 8], got {state.shape}") + if action.ndim != 2 or action.shape[1] != 7: + raise ValueError(f"action must have shape [N, 7], got {action.shape}") + if state.dtype != np.float32: + raise TypeError(f"state must be float32, got {state.dtype}") + if action.dtype != np.float32: + raise TypeError(f"action must be float32, got {action.dtype}") + + lengths = {front.shape[0], wrist.shape[0], state.shape[0], action.shape[0]} + if len(lengths) != 1: + raise ValueError( + "trajectory arrays must have the same first dimension, got " + f"front={len(front)}, wrist={len(wrist)}, state={len(state)}, " + f"action={len(action)}" + ) + if not lengths or next(iter(lengths)) < 1: + raise ValueError("trajectory must contain at least one frame") + if not np.isfinite(state).all(): + raise ValueError("state contains NaN or Infinity") + if not np.isfinite(action).all(): + raise ValueError("action contains NaN or Infinity") + + return ( + np.ascontiguousarray(front), + np.ascontiguousarray(wrist), + np.ascontiguousarray(state), + np.ascontiguousarray(action), + ) + + +def _episode_digest( + job_id: str, + task: str, + provenance: dict[str, Any], + arrays: tuple[np.ndarray, ...], +) -> str: + digest = hashlib.sha256() + digest.update(job_id.encode("utf-8")) + digest.update(b"\0") + digest.update(task.encode("utf-8")) + digest.update(b"\0") + digest.update(_canonical_json(provenance)) + for name, array in zip(("front", "wrist", "state", "action"), arrays): + digest.update(b"\0" + name.encode("ascii") + b"\0") + digest.update(str(array.dtype).encode("ascii")) + digest.update(_canonical_json(list(array.shape))) + digest.update(memoryview(array).cast("B")) + return digest.hexdigest() + + +def _scalar_stats(values: np.ndarray) -> dict[str, list[Any]]: + values = np.asarray(values) + if values.ndim == 1: + values = values[:, None] + numeric = values.astype(np.float64, copy=False) + min_values = values.min(axis=0) + max_values = values.max(axis=0) + return { + "min": min_values.tolist(), + "max": max_values.tolist(), + "mean": numeric.mean(axis=0).tolist(), + "std": numeric.std(axis=0).tolist(), + "count": [int(values.shape[0])], + } + + +def _image_stats(frames: np.ndarray, max_samples: int = 100) -> dict[str, list[Any]]: + if len(frames) > max_samples: + sample_indices = np.linspace(0, len(frames) - 1, max_samples) + sample_indices = np.rint(sample_indices).astype(np.int64) + frames = frames[sample_indices] + + channel_min = np.full(3, np.inf, dtype=np.float64) + channel_max = np.full(3, -np.inf, dtype=np.float64) + channel_sum = np.zeros(3, dtype=np.float64) + channel_sum_sq = np.zeros(3, dtype=np.float64) + pixel_count = 0 + for frame in frames: + normalized = frame.astype(np.float64) / 255.0 + channel_min = np.minimum(channel_min, normalized.min(axis=(0, 1))) + channel_max = np.maximum(channel_max, normalized.max(axis=(0, 1))) + channel_sum += normalized.sum(axis=(0, 1)) + channel_sum_sq += np.square(normalized).sum(axis=(0, 1)) + pixel_count += normalized.shape[0] * normalized.shape[1] + + mean = channel_sum / pixel_count + variance = np.maximum(channel_sum_sq / pixel_count - np.square(mean), 0.0) + std = np.sqrt(variance) + + def nested(values: np.ndarray) -> list[list[list[float]]]: + return values[:, None, None].tolist() + + return { + "min": nested(channel_min), + "max": nested(channel_max), + "mean": nested(mean), + "std": nested(std), + "count": [int(len(frames))], + } + + +def _encode_h264(frames: np.ndarray, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with av.open(str(path), mode="w", format="mp4") as container: + stream = container.add_stream("libx264", rate=FPS) + stream.width = IMAGE_WIDTH + stream.height = IMAGE_HEIGHT + stream.pix_fmt = "yuv420p" + stream.options = { + "crf": "23", + "preset": "veryfast", + "g": str(FPS), + } + for frame in frames: + video_frame = av.VideoFrame.from_ndarray(frame, format="rgb24") + for packet in stream.encode(video_frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + + +def _probe_video(path: Path, *, count_frames: bool) -> dict[str, Any]: + with av.open(str(path), mode="r") as container: + if not container.streams.video: + raise ValueError(f"{path} has no video stream") + stream = container.streams.video[0] + codec = stream.codec_context + rate = stream.average_rate or stream.base_rate + result: dict[str, Any] = { + "video.height": int(codec.height), + "video.width": int(codec.width), + "video.codec": str(codec.name), + "video.pix_fmt": str(codec.pix_fmt), + "video.is_depth_map": False, + "video.fps": float(rate) if rate is not None else None, + "video.channels": 3, + "has_audio": bool(container.streams.audio), + } + if count_frames: + result["decoded_frames"] = sum(1 for _ in container.decode(stream)) + return result + + +def _file_record(path: Path, root: Path) -> dict[str, Any]: + return { + "path": path.relative_to(root).as_posix(), + "size": path.stat().st_size, + "sha256": _sha256_file(path), + } + + +def _verify_file_record(root: Path, record: Mapping[str, Any]) -> Path: + relative = Path(str(record["path"])) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"unsafe staged relative path: {relative}") + path = root / relative + if not path.is_file(): + raise FileNotFoundError(path) + if path.stat().st_size != int(record["size"]): + raise ValueError(f"staged file size changed: {path}") + if _sha256_file(path) != record["sha256"]: + raise ValueError(f"staged file checksum changed: {path}") + return path + + +def stage_success( + staging_root: Path, + job_id: str, + trajectory: dict[str, np.ndarray], + provenance: dict, +) -> Path: + """Atomically stage one successful rollout. + + Repeating the same ``job_id`` with byte-identical trajectory/provenance is + idempotent. Reusing a ``job_id`` for different content raises an error, + preventing a resumed collector from silently replacing a successful sample. + """ + staging_root = Path(staging_root).expanduser().resolve() + job_id = _validate_job_id(job_id) + if not isinstance(provenance, Mapping): + raise TypeError("provenance must be a mapping") + normalized_provenance = _normalize_json(dict(provenance)) + task = normalized_provenance.get("task") + if not isinstance(task, str) or not task.strip(): + raise ValueError("provenance['task'] must be a non-empty string") + if normalized_provenance.get("success") is False: + raise ValueError("stage_success cannot stage provenance with success=false") + + arrays = _validated_trajectory(trajectory) + front, wrist, state, action = arrays + content_digest = _episode_digest(job_id, task, normalized_provenance, arrays) + + success_root = staging_root / "success" + temp_root = staging_root / ".tmp" + success_root.mkdir(parents=True, exist_ok=True) + temp_root.mkdir(parents=True, exist_ok=True) + destination = success_root / job_id + + if destination.exists(): + manifest = _read_json(destination / "manifest.json") + if manifest.get("content_digest") == content_digest: + for record in manifest.get("files", []): + _verify_file_record(destination, record) + return destination + raise FileExistsError( + f"staged job_id {job_id!r} already exists with different content" + ) + + temporary = Path(tempfile.mkdtemp(prefix=f"{job_id}.", dir=temp_root)) + try: + numeric_path = temporary / "trajectory.npz" + np.savez( + numeric_path, + state=state, + action=action, + ) + front_path = temporary / "front.mp4" + wrist_path = temporary / "wrist.mp4" + _encode_h264(front, front_path) + _encode_h264(wrist, wrist_path) + + front_probe = _probe_video(front_path, count_frames=True) + wrist_probe = _probe_video(wrist_path, count_frames=True) + for name, probe in (("front", front_probe), ("wrist", wrist_probe)): + if probe["decoded_frames"] != len(state): + raise RuntimeError( + f"{name} video encoded {probe['decoded_frames']} frames; " + f"expected {len(state)}" + ) + if probe["video.codec"] != "h264": + raise RuntimeError( + f"{name} video codec is {probe['video.codec']!r}, expected h264" + ) + + provenance_path = temporary / "provenance.json" + _write_json(provenance_path, normalized_provenance) + stats = { + FRONT_KEY: _image_stats(front), + WRIST_KEY: _image_stats(wrist), + STATE_KEY: _scalar_stats(state), + ACTION_KEY: _scalar_stats(action), + } + files = [ + _file_record(numeric_path, temporary), + _file_record(front_path, temporary), + _file_record(wrist_path, temporary), + _file_record(provenance_path, temporary), + ] + manifest = { + "staging_version": STAGING_VERSION, + "job_id": job_id, + "content_digest": content_digest, + "task": task, + "num_frames": int(len(state)), + "fps": FPS, + "image_shape": [IMAGE_HEIGHT, IMAGE_WIDTH, 3], + "state_dim": 8, + "action_dim": 7, + "stats": stats, + "files": files, + } + _write_json(temporary / "manifest.json", manifest) + + try: + os.rename(temporary, destination) + except OSError as error: + if error.errno not in (errno.EEXIST, errno.ENOTEMPTY): + raise + existing = _read_json(destination / "manifest.json") + if existing.get("content_digest") != content_digest: + raise FileExistsError( + f"concurrent writer used job_id {job_id!r} for different content" + ) from error + shutil.rmtree(temporary) + return destination + except Exception: + if temporary.exists(): + shutil.rmtree(temporary, ignore_errors=True) + raise + + +def _load_staged(staging_root: Path) -> list[dict[str, Any]]: + success_root = staging_root / "success" + if not success_root.is_dir(): + raise FileNotFoundError(f"no success staging directory: {success_root}") + + staged: list[dict[str, Any]] = [] + for episode_dir in sorted(path for path in success_root.iterdir() if path.is_dir()): + manifest_path = episode_dir / "manifest.json" + if not manifest_path.is_file(): + # Defensive: only a complete atomic directory is eligible. + continue + manifest = _read_json(manifest_path) + if manifest.get("staging_version") != STAGING_VERSION: + raise ValueError( + f"unsupported staging version in {manifest_path}: " + f"{manifest.get('staging_version')}" + ) + if manifest.get("job_id") != episode_dir.name: + raise ValueError(f"job_id/path mismatch in {manifest_path}") + if int(manifest.get("fps", -1)) != FPS: + raise ValueError(f"unexpected fps in {manifest_path}") + if list(manifest.get("image_shape", [])) != [IMAGE_HEIGHT, IMAGE_WIDTH, 3]: + raise ValueError(f"unexpected image shape in {manifest_path}") + if int(manifest.get("state_dim", -1)) != 8: + raise ValueError(f"unexpected state dimension in {manifest_path}") + if int(manifest.get("action_dim", -1)) != 7: + raise ValueError(f"unexpected action dimension in {manifest_path}") + + paths = { + Path(str(record["path"])).name: _verify_file_record(episode_dir, record) + for record in manifest["files"] + } + required = {"trajectory.npz", "front.mp4", "wrist.mp4", "provenance.json"} + if set(paths) != required: + raise ValueError( + f"{manifest_path} has staged files {sorted(paths)}, " + f"expected {sorted(required)}" + ) + provenance = _read_json(paths["provenance.json"]) + task = provenance.get("task") + if task != manifest.get("task"): + raise ValueError(f"task/provenance mismatch in {manifest_path}") + staged.append( + { + "directory": episode_dir, + "manifest": manifest, + "provenance": provenance, + "paths": paths, + } + ) + + if not staged: + raise ValueError(f"no complete successful episodes under {success_root}") + # job_id is the collector's stable attempt identifier and is the explicit + # deterministic ordering contract for retries/resume. + staged.sort(key=lambda item: item["manifest"]["job_id"]) + return staged + + +def _reference_info(reference_dataset: Path | None) -> dict[str, Any] | None: + if reference_dataset is None: + return None + root = Path(reference_dataset).expanduser().resolve() + info_path = root / "meta" / "info.json" + if not info_path.is_file(): + raise FileNotFoundError(info_path) + info = _read_json(info_path) + _check_feature_contract(info, source=str(root)) + return info + + +def _check_feature_contract(info: Mapping[str, Any], *, source: str) -> None: + if str(info.get("codebase_version")) != CODEBASE_VERSION: + raise ValueError( + f"{source} codebase_version is {info.get('codebase_version')!r}, " + f"expected {CODEBASE_VERSION!r}" + ) + if int(info.get("fps", -1)) != FPS: + raise ValueError(f"{source} fps is {info.get('fps')!r}, expected {FPS}") + features = info.get("features", {}) + expected = { + FRONT_KEY: ("video", [IMAGE_HEIGHT, IMAGE_WIDTH, 3]), + WRIST_KEY: ("video", [IMAGE_HEIGHT, IMAGE_WIDTH, 3]), + STATE_KEY: ("float32", [8]), + ACTION_KEY: ("float32", [7]), + "timestamp": ("float32", [1]), + "frame_index": ("int64", [1]), + "episode_index": ("int64", [1]), + "index": ("int64", [1]), + "task_index": ("int64", [1]), + } + for key, (dtype, shape) in expected.items(): + feature = features.get(key) + if not isinstance(feature, Mapping): + raise ValueError(f"{source} is missing feature {key!r}") + if feature.get("dtype") != dtype or list(feature.get("shape", [])) != shape: + raise ValueError( + f"{source} feature {key!r} is " + f"dtype={feature.get('dtype')!r}, shape={feature.get('shape')!r}; " + f"expected dtype={dtype!r}, shape={shape!r}" + ) + + +def _feature_names( + reference_info: dict[str, Any] | None, key: str, fallback: Any +) -> Any: + if reference_info is None: + return fallback + return reference_info["features"][key].get("names", fallback) + + +def _make_info( + *, + episode_count: int, + frame_count: int, + chunks_size: int, + video_info: dict[str, Any], + reference_info: dict[str, Any] | None, +) -> dict[str, Any]: + image_names = ["height", "width", "channel"] + features = { + FRONT_KEY: { + "dtype": "video", + "shape": [IMAGE_HEIGHT, IMAGE_WIDTH, 3], + "names": _feature_names(reference_info, FRONT_KEY, image_names), + "info": dict(video_info), + }, + WRIST_KEY: { + "dtype": "video", + "shape": [IMAGE_HEIGHT, IMAGE_WIDTH, 3], + "names": _feature_names(reference_info, WRIST_KEY, image_names), + "info": dict(video_info), + }, + STATE_KEY: { + "dtype": "float32", + "shape": [8], + "names": _feature_names( + reference_info, STATE_KEY, [f"state_{index}" for index in range(8)] + ), + }, + ACTION_KEY: { + "dtype": "float32", + "shape": [7], + "names": _feature_names( + reference_info, ACTION_KEY, [f"action_{index}" for index in range(7)] + ), + }, + "timestamp": {"dtype": "float32", "shape": [1], "names": None}, + "frame_index": {"dtype": "int64", "shape": [1], "names": None}, + "episode_index": {"dtype": "int64", "shape": [1], "names": None}, + "index": {"dtype": "int64", "shape": [1], "names": None}, + "task_index": {"dtype": "int64", "shape": [1], "names": None}, + } + return { + "codebase_version": CODEBASE_VERSION, + "robot_type": ( + reference_info.get("robot_type", "panda") + if reference_info is not None + else "panda" + ), + "total_episodes": episode_count, + "total_frames": frame_count, + "total_tasks": 0, # filled after task indexing + "total_videos": episode_count * len(VIDEO_KEYS), + "total_chunks": math.ceil(episode_count / chunks_size), + "chunks_size": chunks_size, + "fps": FPS, + "splits": {"train": f"0:{episode_count}"}, + "data_path": DATA_PATH, + "video_path": VIDEO_PATH, + "features": features, + } + + +def _fixed_list_array(values: np.ndarray, list_type: pa.DataType) -> pa.Array: + flattened = pa.array(values.reshape(-1), type=pa.float32()) + array = pa.FixedSizeListArray.from_arrays(flattened, values.shape[1]) + return array.cast(list_type) + + +def _write_episode_parquet( + path: Path, + *, + state: np.ndarray, + action: np.ndarray, + episode_index: int, + global_start: int, + task_index: int, +) -> None: + length = len(state) + table = pa.Table.from_arrays( + [ + _fixed_list_array(state, _STATE_TYPE), + _fixed_list_array(action, _ACTION_TYPE), + pa.array(np.arange(length, dtype=np.float32) / np.float32(FPS)), + pa.array(np.arange(length, dtype=np.int64)), + pa.array(np.full(length, episode_index, dtype=np.int64)), + pa.array(np.arange(global_start, global_start + length, dtype=np.int64)), + pa.array(np.full(length, task_index, dtype=np.int64)), + ], + schema=PARQUET_SCHEMA, + ) + path.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table, path, compression="snappy") + + +def _copy_file(source: Path, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + + +def _rename_exchange(first: Path, second: Path) -> bool: + """Atomically exchange two directories when Linux renameat2 is available.""" + try: + libc = ctypes.CDLL(None, use_errno=True) + renameat2 = libc.renameat2 + except (AttributeError, OSError): + return False + renameat2.argtypes = [ + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_uint, + ] + renameat2.restype = ctypes.c_int + at_fdcwd = -100 + rename_exchange = 2 + result = renameat2( + at_fdcwd, + os.fsencode(first), + at_fdcwd, + os.fsencode(second), + rename_exchange, + ) + if result == 0: + return True + error = ctypes.get_errno() + if error in (errno.ENOSYS, errno.EINVAL, errno.ENOTSUP, errno.EXDEV): + return False + raise OSError(error, os.strerror(error), f"{first} <-> {second}") + + +def _publish_directory(build_root: Path, dataset_root: Path) -> None: + if not dataset_root.exists(): + os.rename(build_root, dataset_root) + return + if not dataset_root.is_dir(): + raise NotADirectoryError(dataset_root) + + # On Linux this makes an existing dataset switch to the fully validated + # replacement in one namespace operation. build_root then names the old + # dataset and can be removed without affecting readers of the new root. + if _rename_exchange(build_root, dataset_root): + shutil.rmtree(build_root, ignore_errors=True) + return + + # Portable fallback: keep the old tree recoverable until the new rename + # succeeds, and restore it if publication fails. + backup = dataset_root.parent / (f".{dataset_root.name}.backup-{uuid.uuid4().hex}") + os.rename(dataset_root, backup) + try: + os.rename(build_root, dataset_root) + except Exception: + os.rename(backup, dataset_root) + raise + shutil.rmtree(backup, ignore_errors=True) + + +def finalize_dataset( + staging_root: Path, + dataset_root: Path, + *, + reference_dataset: Path | None = None, +) -> dict: + """Deterministically build and atomically publish a LeRobot v2.1 dataset.""" + staging_root = Path(staging_root).expanduser().resolve() + dataset_root = Path(dataset_root).expanduser().resolve() + if dataset_root == staging_root or staging_root in dataset_root.parents: + raise ValueError("dataset_root must not be staging_root or inside it") + if dataset_root in staging_root.parents: + raise ValueError("staging_root must not be inside dataset_root") + + staged = _load_staged(staging_root) + reference_info = _reference_info(reference_dataset) + chunks_size = ( + int(reference_info.get("chunks_size", CHUNKS_SIZE)) + if reference_info is not None + else CHUNKS_SIZE + ) + if chunks_size < 1: + raise ValueError(f"invalid chunks_size: {chunks_size}") + + tasks = sorted({item["manifest"]["task"] for item in staged}) + task_indices = {task: index for index, task in enumerate(tasks)} + dataset_root.parent.mkdir(parents=True, exist_ok=True) + build_root = Path( + tempfile.mkdtemp( + prefix=f".{dataset_root.name}.building-", + dir=dataset_root.parent, + ) + ) + try: + episodes: list[dict[str, Any]] = [] + episode_stats: list[dict[str, Any]] = [] + provenance_rows: list[dict[str, Any]] = [] + global_index = 0 + first_video_info: dict[str, Any] | None = None + + for episode_index, item in enumerate(staged): + manifest = item["manifest"] + task = manifest["task"] + task_index = task_indices[task] + length = int(manifest["num_frames"]) + with np.load( + item["paths"]["trajectory.npz"], allow_pickle=False + ) as archive: + state = np.asarray(archive["state"]) + action = np.asarray(archive["action"]) + # Re-validate numeric staging, including exact float32 dtypes. + if state.shape != (length, 8) or state.dtype != np.float32: + raise ValueError( + f"invalid staged state for {manifest['job_id']}: " + f"{state.shape} {state.dtype}" + ) + if action.shape != (length, 7) or action.dtype != np.float32: + raise ValueError( + f"invalid staged action for {manifest['job_id']}: " + f"{action.shape} {action.dtype}" + ) + if not np.isfinite(state).all() or not np.isfinite(action).all(): + raise ValueError(f"non-finite numeric data for {manifest['job_id']}") + + episode_chunk = episode_index // chunks_size + parquet_path = build_root / DATA_PATH.format( + episode_chunk=episode_chunk, + episode_index=episode_index, + ) + _write_episode_parquet( + parquet_path, + state=state, + action=action, + episode_index=episode_index, + global_start=global_index, + task_index=task_index, + ) + + for video_key, staged_name in ( + (FRONT_KEY, "front.mp4"), + (WRIST_KEY, "wrist.mp4"), + ): + video_path = build_root / VIDEO_PATH.format( + episode_chunk=episode_chunk, + episode_index=episode_index, + video_key=video_key, + ) + _copy_file(item["paths"][staged_name], video_path) + if first_video_info is None: + first_video_info = _probe_video(video_path, count_frames=False) + + timestamps = np.arange(length, dtype=np.float32) / np.float32(FPS) + frame_indices = np.arange(length, dtype=np.int64) + episode_indices = np.full(length, episode_index, dtype=np.int64) + indices = np.arange(global_index, global_index + length, dtype=np.int64) + task_index_values = np.full(length, task_index, dtype=np.int64) + stats = dict(manifest["stats"]) + stats.update( + { + "timestamp": _scalar_stats(timestamps), + "frame_index": _scalar_stats(frame_indices), + "episode_index": _scalar_stats(episode_indices), + "index": _scalar_stats(indices), + "task_index": _scalar_stats(task_index_values), + } + ) + episodes.append( + { + "episode_index": episode_index, + "tasks": [task], + "length": length, + } + ) + episode_stats.append({"episode_index": episode_index, "stats": stats}) + provenance_rows.append( + { + "episode_index": episode_index, + "job_id": manifest["job_id"], + "content_digest": manifest["content_digest"], + "task": task, + "provenance": item["provenance"], + } + ) + global_index += length + + assert first_video_info is not None + if first_video_info["video.codec"] != "h264": + raise ValueError( + f"published video codec is {first_video_info['video.codec']!r}, " + "expected h264" + ) + first_video_info.pop("decoded_frames", None) + fps_value = first_video_info["video.fps"] + if fps_value is None or not math.isclose(float(fps_value), FPS, abs_tol=1e-6): + raise ValueError(f"published video fps is {fps_value!r}, expected {FPS}") + # LeRobot metadata represents integral rates as integers. + first_video_info["video.fps"] = FPS + + info = _make_info( + episode_count=len(episodes), + frame_count=global_index, + chunks_size=chunks_size, + video_info=first_video_info, + reference_info=reference_info, + ) + info["total_tasks"] = len(tasks) + meta_root = build_root / "meta" + _write_json(meta_root / "info.json", info) + _write_jsonl( + meta_root / "tasks.jsonl", + [{"task_index": index, "task": task} for index, task in enumerate(tasks)], + ) + _write_jsonl(meta_root / "episodes.jsonl", episodes) + _write_jsonl(meta_root / "episodes_stats.jsonl", episode_stats) + _write_jsonl(meta_root / "provenance.jsonl", provenance_rows) + + report = validate_dataset( + build_root, + reference_dataset=reference_dataset, + ) + _publish_directory(build_root, dataset_root) + report["dataset_root"] = str(dataset_root) + report["staging_root"] = str(staging_root) + return report + except Exception: + if build_root.exists(): + shutil.rmtree(build_root, ignore_errors=True) + raise + + +def _validate_parquet_schema(path: Path) -> pa.Schema: + schema = pq.read_schema(path) + if not schema.equals(PARQUET_SCHEMA, check_metadata=True): + raise ValueError( + f"{path} has incompatible Parquet schema:\n{schema}\n" + f"expected:\n{PARQUET_SCHEMA}" + ) + return schema + + +def validate_dataset( + dataset_root: Path, + reference_dataset: Path | None = None, +) -> dict: + """Deep-validate metadata, Parquet values, and every encoded video. + + Returns a compact summary on success and raises on the first contract + violation. This intentionally decodes all videos so truncated MP4 files + cannot pass validation. + """ + dataset_root = Path(dataset_root).expanduser().resolve() + info = _read_json(dataset_root / "meta" / "info.json") + _check_feature_contract(info, source=str(dataset_root)) + reference_info = _reference_info(reference_dataset) + + if info.get("data_path") != DATA_PATH: + raise ValueError(f"unexpected data_path: {info.get('data_path')!r}") + if info.get("video_path") != VIDEO_PATH: + raise ValueError(f"unexpected video_path: {info.get('video_path')!r}") + if reference_info is not None: + for key in ( + FRONT_KEY, + WRIST_KEY, + STATE_KEY, + ACTION_KEY, + "timestamp", + "frame_index", + "episode_index", + "index", + "task_index", + ): + actual = info["features"][key] + reference = reference_info["features"][key] + for attribute in ("dtype", "shape", "names"): + if actual.get(attribute) != reference.get(attribute): + raise ValueError( + f"feature {key!r} {attribute} differs from reference: " + f"{actual.get(attribute)!r} != {reference.get(attribute)!r}" + ) + + tasks = _read_jsonl(dataset_root / "meta" / "tasks.jsonl") + episodes = _read_jsonl(dataset_root / "meta" / "episodes.jsonl") + episode_stats = _read_jsonl(dataset_root / "meta" / "episodes_stats.jsonl") + provenance = _read_jsonl(dataset_root / "meta" / "provenance.jsonl") + + episode_count = int(info["total_episodes"]) + frame_count = int(info["total_frames"]) + task_count = int(info["total_tasks"]) + chunks_size = int(info["chunks_size"]) + if episode_count < 1 or frame_count < 1 or task_count < 1: + raise ValueError("dataset totals must all be positive") + if chunks_size < 1: + raise ValueError("chunks_size must be positive") + if len(episodes) != episode_count: + raise ValueError("episodes.jsonl length does not match total_episodes") + if len(episode_stats) != episode_count: + raise ValueError("episodes_stats.jsonl length does not match total_episodes") + if len(provenance) != episode_count: + raise ValueError("provenance.jsonl length does not match total_episodes") + if len(tasks) != task_count: + raise ValueError("tasks.jsonl length does not match total_tasks") + if int(info["total_videos"]) != episode_count * len(VIDEO_KEYS): + raise ValueError("total_videos is inconsistent") + if int(info["total_chunks"]) != math.ceil(episode_count / chunks_size): + raise ValueError("total_chunks is inconsistent") + if info.get("splits") != {"train": f"0:{episode_count}"}: + raise ValueError("train split is inconsistent") + + expected_episode_indices = list(range(episode_count)) + for name, rows in ( + ("episodes", episodes), + ("episodes_stats", episode_stats), + ("provenance", provenance), + ): + indices = [int(row["episode_index"]) for row in rows] + if indices != expected_episode_indices: + raise ValueError(f"{name} episode indices are not contiguous") + task_map = {int(row["task_index"]): row["task"] for row in tasks} + if sorted(task_map) != list(range(task_count)) or len(task_map) != len(tasks): + raise ValueError("task indices are not unique and contiguous") + if len(set(task_map.values())) != task_count: + raise ValueError("task strings are not unique") + + global_start = 0 + video_codec: str | None = None + for episode_index, episode in enumerate(episodes): + length = int(episode["length"]) + if length < 1: + raise ValueError(f"episode {episode_index} has no frames") + if not isinstance(episode.get("tasks"), list) or len(episode["tasks"]) != 1: + raise ValueError(f"episode {episode_index} must contain exactly one task") + + episode_chunk = episode_index // chunks_size + parquet_path = dataset_root / DATA_PATH.format( + episode_chunk=episode_chunk, + episode_index=episode_index, + ) + if not parquet_path.is_file(): + raise FileNotFoundError(parquet_path) + _validate_parquet_schema(parquet_path) + table = pq.read_table(parquet_path) + if table.num_rows != length: + raise ValueError( + f"{parquet_path} has {table.num_rows} rows; expected {length}" + ) + + state = np.asarray( + table[STATE_KEY].combine_chunks().to_pylist(), dtype=np.float32 + ) + action = np.asarray( + table[ACTION_KEY].combine_chunks().to_pylist(), dtype=np.float32 + ) + if state.shape != (length, 8) or action.shape != (length, 7): + raise ValueError(f"numeric dimensions are invalid in {parquet_path}") + if not np.isfinite(state).all() or not np.isfinite(action).all(): + raise ValueError(f"non-finite numeric values in {parquet_path}") + + expected_columns = { + "timestamp": np.arange(length, dtype=np.float32) / np.float32(FPS), + "frame_index": np.arange(length, dtype=np.int64), + "episode_index": np.full(length, episode_index, dtype=np.int64), + "index": np.arange(global_start, global_start + length, dtype=np.int64), + } + for column_name, expected in expected_columns.items(): + actual = table[column_name].combine_chunks().to_numpy(zero_copy_only=False) + if column_name == "timestamp": + equal = np.allclose(actual, expected, atol=1e-6, rtol=0) + else: + equal = np.array_equal(actual, expected) + if not equal: + raise ValueError(f"{column_name} values are invalid in {parquet_path}") + + task_values = ( + table["task_index"].combine_chunks().to_numpy(zero_copy_only=False) + ) + unique_task_values = np.unique(task_values) + if len(unique_task_values) != 1: + raise ValueError(f"multiple task indices in {parquet_path}") + task_index = int(unique_task_values[0]) + if task_index not in task_map: + raise ValueError(f"unknown task index {task_index} in {parquet_path}") + if episode["tasks"][0] != task_map[task_index]: + raise ValueError(f"episode task does not match Parquet in {parquet_path}") + + stats = episode_stats[episode_index].get("stats", {}) + required_stats = { + FRONT_KEY, + WRIST_KEY, + STATE_KEY, + ACTION_KEY, + "timestamp", + "frame_index", + "episode_index", + "index", + "task_index", + } + if set(stats) != required_stats: + raise ValueError( + f"episode {episode_index} stats keys are {sorted(stats)}, " + f"expected {sorted(required_stats)}" + ) + if provenance[episode_index].get("task") != episode["tasks"][0]: + raise ValueError(f"episode {episode_index} provenance task mismatch") + + for video_key in VIDEO_KEYS: + video_path = dataset_root / VIDEO_PATH.format( + episode_chunk=episode_chunk, + episode_index=episode_index, + video_key=video_key, + ) + if not video_path.is_file(): + raise FileNotFoundError(video_path) + probe = _probe_video(video_path, count_frames=True) + if ( + probe["video.height"] != IMAGE_HEIGHT + or probe["video.width"] != IMAGE_WIDTH + ): + raise ValueError(f"video dimensions are invalid: {video_path}") + if probe["video.codec"] != "h264": + raise ValueError(f"video is not H.264: {video_path}") + if probe["decoded_frames"] != length: + raise ValueError( + f"{video_path} decodes {probe['decoded_frames']} frames; " + f"expected {length}" + ) + if probe["has_audio"]: + raise ValueError(f"unexpected audio stream: {video_path}") + if probe["video.fps"] is None or not math.isclose( + float(probe["video.fps"]), FPS, abs_tol=1e-6 + ): + raise ValueError(f"video fps is not {FPS}: {video_path}") + if video_codec is None: + video_codec = probe["video.codec"] + elif video_codec != probe["video.codec"]: + raise ValueError("dataset contains mixed video codecs") + global_start += length + + if global_start != frame_count: + raise ValueError( + f"episode lengths sum to {global_start}; total_frames is {frame_count}" + ) + video_metadata = info["features"][FRONT_KEY]["info"] + if video_metadata.get("video.codec") != "h264": + raise ValueError("info.json does not declare H.264 video") + if int(video_metadata.get("video.fps", -1)) != FPS: + raise ValueError("info.json video fps is inconsistent") + + return { + "dataset_root": str(dataset_root), + "total_episodes": episode_count, + "total_frames": frame_count, + "total_tasks": task_count, + "total_videos": episode_count * len(VIDEO_KEYS), + "video_codec": video_codec, + "fps": FPS, + "valid": True, + } + + +__all__ = [ + "ACTION_KEY", + "FPS", + "FRONT_KEY", + "PARQUET_SCHEMA", + "STATE_KEY", + "WRIST_KEY", + "finalize_dataset", + "stage_success", + "validate_dataset", +] diff --git a/experiments/libero/libero_plus_eval_utils.py b/experiments/libero/libero_plus_eval_utils.py new file mode 100644 index 00000000..9c28286e --- /dev/null +++ b/experiments/libero/libero_plus_eval_utils.py @@ -0,0 +1,231 @@ +"""Shared, dependency-light helpers for LIBERO-Plus Table 13 evaluation.""" + +from __future__ import annotations + +import json +import random +from dataclasses import asdict, dataclass +from pathlib import Path + + +SUITE_ORDER = ( + "libero_spatial", + "libero_object", + "libero_goal", + "libero_10", +) + +# Match the display order in Kairos Table 13. +CATEGORY_ORDER = ( + "Camera Viewpoints", + "Robot Initial States", + "Language Instructions", + "Light Conditions", + "Background Textures", + "Sensor Noise", + "Objects Layout", +) + +CATEGORY_LABELS = { + "Camera Viewpoints": "Camera", + "Robot Initial States": "Robot", + "Language Instructions": "Language", + "Light Conditions": "Light", + "Background Textures": "Background", + "Sensor Noise": "Noise", + "Objects Layout": "Layout", +} + +EXPECTED_SUITE_COUNTS = { + "libero_spatial": 2402, + "libero_object": 2518, + "libero_goal": 2591, + "libero_10": 2519, +} +EXPECTED_TOTAL_TASKS = 10030 + + +@dataclass(frozen=True) +class LiberoPlusTask: + suite: str + task_id: int + classification_id: int + name: str + category: str + difficulty_level: int | None + + @property + def key(self) -> str: + return f"{self.suite}/{self.task_id}" + + def to_dict(self) -> dict: + return asdict(self) + + @classmethod + def from_dict(cls, value: dict) -> "LiberoPlusTask": + return cls( + suite=str(value["suite"]), + task_id=int(value["task_id"]), + classification_id=int(value["classification_id"]), + name=str(value["name"]), + category=str(value["category"]), + difficulty_level=( + None + if value.get("difficulty_level") is None + else int(value["difficulty_level"]) + ), + ) + + +def default_classification_path(libero_plus_root: str | Path) -> Path: + return ( + Path(libero_plus_root).expanduser().resolve() + / "libero" + / "libero" + / "benchmark" + / "task_classification.json" + ) + + +def load_task_classification( + classification_path: str | Path, + *, + require_full: bool = True, +) -> list[LiberoPlusTask]: + path = Path(classification_path).expanduser().resolve() + with path.open(encoding="utf-8") as handle: + raw = json.load(handle) + + if not isinstance(raw, dict): + raise TypeError(f"Expected a JSON object in {path}, got {type(raw).__name__}.") + missing_suites = [suite for suite in SUITE_ORDER if suite not in raw] + if missing_suites: + raise ValueError(f"Classification file is missing suites: {missing_suites}") + + tasks: list[LiberoPlusTask] = [] + seen_keys: set[str] = set() + for suite in SUITE_ORDER: + entries = raw[suite] + if not isinstance(entries, list): + raise TypeError(f"Classification entry {suite!r} must be a list.") + if require_full and len(entries) != EXPECTED_SUITE_COUNTS[suite]: + raise ValueError( + f"Unexpected task count for {suite}: {len(entries)}; " + f"expected {EXPECTED_SUITE_COUNTS[suite]}." + ) + + for position, entry in enumerate(entries): + classification_id = int(entry["id"]) + expected_id = position + 1 + if classification_id != expected_id: + raise ValueError( + f"Non-contiguous classification ids in {suite}: " + f"position={position}, id={classification_id}, expected={expected_id}." + ) + category = str(entry["category"]) + if category not in CATEGORY_ORDER: + raise ValueError(f"Unknown LIBERO-Plus category {category!r} in {suite}/{classification_id}.") + task = LiberoPlusTask( + suite=suite, + task_id=classification_id - 1, + classification_id=classification_id, + name=str(entry["name"]), + category=category, + difficulty_level=( + None + if entry.get("difficulty_level") is None + else int(entry["difficulty_level"]) + ), + ) + if task.key in seen_keys: + raise ValueError(f"Duplicate LIBERO-Plus task key: {task.key}") + seen_keys.add(task.key) + tasks.append(task) + + if require_full and len(tasks) != EXPECTED_TOTAL_TASKS: + raise ValueError( + f"Unexpected total LIBERO-Plus task count: {len(tasks)}; " + f"expected {EXPECTED_TOTAL_TASKS}." + ) + return tasks + + +def select_smoke_tasks( + tasks: list[LiberoPlusTask], + limit: int = 8, + *, + seed: int = 42, +) -> list[LiberoPlusTask]: + """Randomly sample a reproducible subset balanced over present categories.""" + if limit < 1: + raise ValueError(f"Smoke task limit must be positive, got {limit}.") + if limit >= len(tasks): + return list(tasks) + + rng = random.Random(seed) + category_buckets = { + category: [task for task in tasks if task.category == category] + for category in CATEGORY_ORDER + } + active_categories = [ + category for category in CATEGORY_ORDER if category_buckets[category] + ] + if not active_categories: + raise ValueError("Cannot select smoke tasks from an empty task list.") + base_quota, remainder = divmod(limit, len(active_categories)) + quotas = {category: base_quota for category in active_categories} + for category in rng.sample(active_categories, k=remainder): + quotas[category] += 1 + + selected: list[LiberoPlusTask] = [] + for category in active_categories: + quota = quotas[category] + bucket = category_buckets[category] + if quota > len(bucket): + raise ValueError( + f"Smoke quota for {category!r} is {quota}, but only {len(bucket)} tasks exist." + ) + selected.extend(rng.sample(bucket, k=quota)) + rng.shuffle(selected) + return selected + + +def result_path(output_dir: str | Path, task: LiberoPlusTask) -> Path: + return ( + Path(output_dir) + / "results" + / task.suite + / f"task_{task.task_id:04d}.json" + ) + + +def failure_path(output_dir: str | Path, task: LiberoPlusTask) -> Path: + return ( + Path(output_dir) + / "failures" + / task.suite + / f"task_{task.task_id:04d}.json" + ) + + +def is_complete_result( + path: str | Path, + task: LiberoPlusTask, + *, + num_trials: int, +) -> bool: + path = Path(path) + if not path.is_file(): + return False + try: + with path.open(encoding="utf-8") as handle: + result = json.load(handle) + return ( + result.get("task_suite") == task.suite + and int(result.get("task_id", -1)) == task.task_id + and result.get("category_value") == task.category + and int(result.get("total_episodes", -1)) == int(num_trials) + and 0 <= int(result.get("successes", -1)) <= int(num_trials) + ) + except (OSError, TypeError, ValueError, json.JSONDecodeError): + return False diff --git a/experiments/libero/libero_utils.py b/experiments/libero/libero_utils.py index 68edfd81..f6405fec 100644 --- a/experiments/libero/libero_utils.py +++ b/experiments/libero/libero_utils.py @@ -1,8 +1,11 @@ """Utils for evaluating policies in LIBERO simulation environments.""" +import logging import math +import os import time import pathlib +from contextlib import contextmanager import imageio from PIL import Image, ImageDraw @@ -16,6 +19,61 @@ LIBERO_ENV_RESOLUTION = 256 # resolution used to render training data +def _env_truthy(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +@contextmanager +def _egl_init_lock(): + if not _env_truthy("LIBERO_EGL_INIT_LOCK", default=True): + yield + return + + import fcntl + + lock_file = os.environ.get("LIBERO_EGL_INIT_LOCK_FILE", "/tmp/libero_mujoco_egl_init.lock") + pathlib.Path(lock_file).parent.mkdir(parents=True, exist_ok=True) + with open(lock_file, "w", encoding="utf-8") as handle: + fcntl.flock(handle, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(handle, fcntl.LOCK_UN) + + +def _create_offscreen_env_with_retry(env_args: dict): + max_attempts = max(1, int(os.environ.get("LIBERO_EGL_INIT_RETRIES", "5"))) + retry_sleep = float(os.environ.get("LIBERO_EGL_INIT_RETRY_SLEEP_SEC", "5")) + last_error = None + + for attempt in range(1, max_attempts + 1): + try: + with _egl_init_lock(): + return OffScreenRenderEnv(**env_args) + except RuntimeError as error: + last_error = error + message = str(error) + is_egl_error = ( + "MUJOCO_EGL_DEVICE_ID environment variable" in message + or "Cannot initialize a EGL device display" in message + ) + if not is_egl_error or attempt == max_attempts: + raise + logging.warning( + "LIBERO EGL init failed on attempt %s/%s; retrying in %.1fs: %s", + attempt, + max_attempts, + retry_sleep, + error, + ) + time.sleep(retry_sleep) + + raise RuntimeError("LIBERO EGL init failed") from last_error + + def get_libero_env(task, resolution, seed, env_num=1): """Initializes and returns the LIBERO environment, along with the task description.""" task_description = task.language @@ -25,14 +83,15 @@ def get_libero_env(task, resolution, seed, env_num=1): / task.bddl_file ) env_args = { - "bddl_file_name": task_bddl_file, + # LIBERO-Plus still performs substring checks on this argument. + "bddl_file_name": str(task_bddl_file), "camera_heights": resolution, "camera_widths": resolution, } if env_num > 1: env = SubprocVectorEnv([lambda: OffScreenRenderEnv(**env_args) for _ in range(env_num)]) else: - env = OffScreenRenderEnv(**env_args) + env = _create_offscreen_env_with_retry(env_args) env.seed( seed ) # IMPORTANT: seed seems to affect object positions even when using fixed initial state diff --git a/experiments/libero/retry_libero_plus_failed_rollouts.py b/experiments/libero/retry_libero_plus_failed_rollouts.py new file mode 100644 index 00000000..ec2c3a06 --- /dev/null +++ b/experiments/libero/retry_libero_plus_failed_rollouts.py @@ -0,0 +1,1859 @@ +"""Adaptively retry failed LIBERO-Plus Robot variants in an existing collection. + +The base collector has an immutable, static job manifest. This companion +runner leaves that manifest untouched and creates a separate retry campaign: + +* by default, only variants with zero successes in the base manifest are + targeted; ``--target-successes-per-variant`` can require repeated successes; +* each variant has at most one outstanding rollout; +* a failure or insufficient success count schedules another deterministic job; +* a variant retires once its base + retry successes reach the requested target; +* per-variant attempts and active campaign runtime are bounded; +* attempts and success staging are appended to the existing collection root; +* finalization rebuilds and atomically publishes the union of old and new + successful episodes. + +The campaign manifest freezes its target set, input fingerprints, seed +derivation, and limits. Resume reconstructs task state from atomic attempt +files, while ``state.json`` only accounts for active wall-clock runtime. +""" + +from __future__ import annotations + +import argparse +import fcntl +import hashlib +import json +import multiprocessing as mp +import os +import queue +import re +import signal +import time +from collections import Counter +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import collect_libero_plus_self_rollouts as base + + +CAMPAIGN_SCHEMA_VERSION = 1 +STATE_SCHEMA_VERSION = 1 +DEFAULT_CAMPAIGN_NAME = "retry_zero_success_robot_v1" +DEFAULT_RETRY_SEED_SALT = "fastwam-libero-plus-robot-retry-v1" +MAX_GRACEFUL_STOP_SECONDS = 300.0 +SEED_MODULUS = 2**31 - 1 + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _campaign_dir(output_dir: Path, name: str) -> Path: + if ( + not name + or name in {".", ".."} + or Path(name).name != name + or re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", name) is None + ): + raise ValueError( + "--campaign-name must contain only letters, digits, '.', '_', or '-'." + ) + return output_dir / "retry_campaigns" / name + + +def _task_identity(task: base.CollectionTask) -> str: + return f"{task.suite}/{task.task_id}" + + +def _uses_global_additional_target(args: argparse.Namespace) -> bool: + below = getattr(args, "select_global_successes_below", None) + additional = getattr( + args, + "target_additional_successes_per_variant", + None, + ) + if (below is None) != (additional is None): + raise ValueError( + "--select-global-successes-below and " + "--target-additional-successes-per-variant must be used together." + ) + return below is not None + + +def _effective_target_successes(args: argparse.Namespace) -> int: + if _uses_global_additional_target(args): + return int(args.target_additional_successes_per_variant) + return int(args.target_successes_per_variant) + + +def _normalize_wall_deadline(value: str | None) -> str | None: + if value is None: + return None + try: + deadline = datetime.fromisoformat(str(value)) + except ValueError as error: + raise ValueError( + "--wall-deadline must be an ISO-8601 timestamp with timezone." + ) from error + if deadline.tzinfo is None or deadline.utcoffset() is None: + raise ValueError("--wall-deadline must include an explicit timezone.") + return deadline.astimezone(timezone.utc).isoformat() + + +def _wall_deadline_timestamp(value: str | None) -> float | None: + normalized = _normalize_wall_deadline(value) + if normalized is None: + return None + return datetime.fromisoformat(normalized).timestamp() + + +def _derive_seed( + task: base.CollectionTask, + *, + retry_number: int, + channel: str, + salt: str, +) -> int: + payload = { + "channel": channel, + "retry_number": int(retry_number), + "salt": salt, + "task": task.to_dict(), + } + digest = hashlib.sha256( + json.dumps(payload, ensure_ascii=True, sort_keys=True).encode("utf-8") + ).digest() + value = int.from_bytes(digest[:8], "big") % SEED_MODULUS + if value in {0, 42, 43}: + value = (value + 104729) % SEED_MODULUS + return max(1, value) + + +def _retry_job( + task: base.CollectionTask, + *, + retry_index: int, + retry_index_start: int, + seed_salt: str, + initialization: str, + init_state_index: int, + prompt_mode: str, +) -> base.CollectionJob: + retry_number = int(retry_index_start) + int(retry_index) + if initialization == base.TABLE13_EVAL_STATE0_INITIALIZATION: + env_seed = base.TABLE13_EVAL_ENV_SEED + policy_seed = retry_number + else: + env_seed = _derive_seed( + task, + retry_number=retry_number, + channel="environment", + salt=seed_salt, + ) + policy_seed = _derive_seed( + task, + retry_number=retry_number, + channel="policy", + salt=seed_salt, + ) + return base.CollectionJob( + task=task, + env_seed=env_seed, + policy_seed=policy_seed, + initialization=initialization, + init_state_index=init_state_index, + prompt_mode=prompt_mode, + ) + + +def _retry_init_state_index(initialization: str) -> int: + if initialization == base.RANDOM_RESET_INITIALIZATION: + return -1 + if initialization == base.TABLE13_EVAL_STATE0_INITIALIZATION: + return 0 + raise ValueError(f"Unsupported adaptive retry initialization: {initialization}") + + +def _read_base_collection( + output_dir: Path, +) -> tuple[dict[str, Any], list[base.CollectionTask], list[base.CollectionJob]]: + manifest_path = output_dir / "collection_manifest.json" + if not manifest_path.is_file(): + raise FileNotFoundError( + f"Existing base collection manifest not found: {manifest_path}" + ) + manifest = base._read_json(manifest_path) + tasks = [base.CollectionTask.from_dict(value) for value in manifest.get("tasks", [])] + jobs = [base.CollectionJob.from_dict(value) for value in manifest.get("jobs", [])] + if not tasks or not jobs: + raise ValueError(f"Base manifest contains no tasks/jobs: {manifest_path}") + if len({_task_identity(task) for task in tasks}) != len(tasks): + raise ValueError("Base manifest contains duplicate task identities.") + if len({job.job_id for job in jobs}) != len(jobs): + raise ValueError("Base manifest contains duplicate job IDs.") + return manifest, tasks, jobs + + +def _validate_base_semantics( + base_manifest: dict[str, Any], + args: argparse.Namespace, + *, + classification_path: Path, + input_files: dict[str, Any], +) -> None: + expected = { + "task_source": "table13", + "initialization": args.initialization, + "prompt_mode": args.prompt_mode, + "task_config": args.task_config, + "gripper_action_format": args.gripper_action_format, + "hydra_overrides": list(args.override), + } + mismatches = { + key: (base_manifest.get(key), value) + for key, value in expected.items() + if base_manifest.get(key) != value + } + if mismatches: + details = ", ".join( + f"{key}: base={actual!r} retry={expected_value!r}" + for key, (actual, expected_value) in mismatches.items() + ) + raise ValueError(f"Retry command differs from base collection semantics: {details}") + + schema_version = int(base_manifest.get("schema_version", 1)) + if schema_version >= 2: + if base_manifest.get("input_files") != input_files: + raise ValueError( + "Retry input fingerprints differ from the base collection manifest." + ) + return + + # Legacy schema 1 did not record content hashes. Require its recorded paths + # to resolve to the exact inputs supplied now instead of silently mixing a + # different checkpoint, stats file, reference dataset, or classification. + legacy_paths = { + "checkpoint": Path(args.checkpoint).expanduser().resolve(), + "dataset_stats": Path(args.dataset_stats).expanduser().resolve(), + "reference_dataset": Path(args.reference_dataset).expanduser().resolve(), + "libero_plus_root": Path(args.libero_plus_root).expanduser().resolve(), + "classification_path": classification_path, + } + legacy_mismatches: dict[str, tuple[str, str]] = {} + for key, expected_path in legacy_paths.items(): + recorded = base_manifest.get(key) + if recorded is None: + continue + recorded_path = Path(str(recorded)).expanduser().resolve() + if recorded_path != expected_path: + legacy_mismatches[key] = (str(recorded_path), str(expected_path)) + if legacy_mismatches: + details = ", ".join( + f"{key}: base={actual!r} retry={expected!r}" + for key, (actual, expected) in legacy_mismatches.items() + ) + raise ValueError(f"Retry inputs differ from legacy base manifest: {details}") + + +def _base_success_counts( + output_dir: Path, + tasks: list[base.CollectionTask], + jobs: list[base.CollectionJob], +) -> dict[str, int]: + task_by_identity = {_task_identity(task): task for task in tasks} + successes = Counter() + job_counts = Counter() + for job in jobs: + identity = _task_identity(job.task) + if identity not in task_by_identity: + raise ValueError(f"Base job refers to an unknown task: {job.job_id}") + result_path = base._result_path(output_dir, job) + if not base._is_complete_attempt(result_path, job): + raise RuntimeError( + "Base collection is not complete; finish it before starting a retry " + f"campaign. Missing/incomplete attempt: {result_path}" + ) + result = base._read_json(result_path) + job_counts[identity] += 1 + successes[identity] += int(bool(result["success"])) + missing_jobs = sorted(set(task_by_identity) - set(job_counts)) + if missing_jobs: + raise ValueError(f"Base tasks have no jobs: {missing_jobs[:5]}") + return { + identity: int(successes[identity]) + for identity in task_by_identity + } + + +def _global_success_counts( + output_dir: Path, + tasks: list[base.CollectionTask], +) -> dict[str, int]: + """Count every validated success bundle already in this collection.""" + task_by_identity = {_task_identity(task): task for task in tasks} + successes = Counter() + seen_job_ids: set[str] = set() + success_root = output_dir / "staging" / "success" + for artifact_dir in sorted( + path for path in success_root.glob("*") if path.is_dir() + ): + result_path = output_dir / "attempts" / f"{artifact_dir.name}.json" + if not result_path.is_file(): + raise ValueError( + "Global success staging has no matching attempt: " + f"{artifact_dir}" + ) + result = base._read_json(result_path) + job = base.CollectionJob.from_dict(result["job"]) + identity = _task_identity(job.task) + if ( + identity not in task_by_identity + or job.task.to_dict() != task_by_identity[identity].to_dict() + ): + raise ValueError( + f"Global attempt refers to an unknown or changed task: {result_path}" + ) + if result_path != base._result_path(output_dir, job): + raise ValueError(f"Global attempt has an unexpected path: {result_path}") + if artifact_dir.name != job.job_id or not bool(result.get("success")): + raise ValueError(f"Invalid global success bundle: {artifact_dir}") + if job.job_id in seen_job_ids: + raise ValueError(f"Duplicate global attempt job id: {job.job_id}") + seen_job_ids.add(job.job_id) + if not base._is_complete_attempt(result_path, job): + raise ValueError(f"Global attempt is incomplete or invalid: {result_path}") + successes[identity] += int(bool(result["success"])) + return { + identity: int(successes[identity]) + for identity in task_by_identity + } + + +def _base_below_target_tasks( + tasks: list[base.CollectionTask], + base_success_counts: dict[str, int], + *, + target_successes_per_variant: int, +) -> list[base.CollectionTask]: + return [ + task + for task in tasks + if base_success_counts[_task_identity(task)] + < int(target_successes_per_variant) + ] + + +def _validate_seed_interval_against_campaigns( + output_dir: Path, + campaign_dir: Path, + *, + retry_index_start: int, + retry_max_attempts: int, +) -> None: + new_start = int(retry_index_start) + new_end = new_start + int(retry_max_attempts) - 1 + for path in sorted((output_dir / "retry_campaigns").glob("*/manifest.json")): + if path.parent == campaign_dir: + continue + manifest = base._read_json(path) + old_start = int(manifest["retry_index_start"]) + old_end = old_start + int(manifest["retry_max_attempts"]) - 1 + if max(new_start, old_start) <= min(new_end, old_end): + raise ValueError( + "Retry seed interval overlaps an existing campaign: " + f"new=[{new_start},{new_end}] " + f"existing={path.parent.name}[{old_start},{old_end}]." + ) + + +def _base_zero_success_tasks( + output_dir: Path, + tasks: list[base.CollectionTask], + jobs: list[base.CollectionJob], +) -> list[base.CollectionTask]: + """Backward-compatible helper for the original target=1 policy.""" + counts = _base_success_counts(output_dir, tasks, jobs) + return _base_below_target_tasks( + tasks, + counts, + target_successes_per_variant=1, + ) + + +def _validate_tasks_against_current_inputs( + tasks: list[base.CollectionTask], + *, + classification_path: Path, + reference_dataset: Path, +) -> None: + current = { + _task_identity(task): task + for task in base._table13_robot_tasks( + classification_path, + reference_dataset, + ) + } + for task in tasks: + identity = _task_identity(task) + if identity not in current or current[identity].to_dict() != task.to_dict(): + raise ValueError( + f"Base retry target differs from current classification: {identity}" + ) + + +def _campaign_manifest( + args: argparse.Namespace, + *, + output_dir: Path, + base_manifest_path: Path, + input_files: dict[str, Any], + tasks: list[base.CollectionTask], + base_success_counts: dict[str, int], + global_success_counts_at_start: dict[str, int] | None = None, +) -> dict[str, Any]: + global_additional = _uses_global_additional_target(args) + target_successes = _effective_target_successes(args) + manifest = { + "schema_version": CAMPAIGN_SCHEMA_VERSION, + "created_at": _utc_now(), + "policy": ( + "global_below_threshold_until_additional_success_target" + if global_additional + else ( + "base_zero_success_until_first_retry_success" + if target_successes == 1 + else "base_below_target_until_cumulative_success_target" + ) + ), + "base_collection_manifest": base._file_identity(base_manifest_path), + "input_files": input_files, + "prompt_embedding": base._prompt_embedding_identity( + args.text_embedding_cache_dir + ), + "task_config": args.task_config, + "prompt_mode": args.prompt_mode, + "initialization": args.initialization, + "init_state_index": _retry_init_state_index(args.initialization), + "gripper_action_format": args.gripper_action_format, + "hydra_overrides": list(args.override), + "retry_index_start": int(args.retry_index_start), + "retry_max_attempts": int(args.retry_max_attempts), + "max_active_runtime_seconds": float(args.max_runtime_hours) * 3600.0, + "target_count": len(tasks), + "tasks": [task.to_dict() for task in tasks], + "output_layout": { + "attempts": str((output_dir / "attempts").relative_to(output_dir)), + "success_staging": str( + (output_dir / "staging" / "success").relative_to(output_dir) + ), + "dataset": str((output_dir / "lerobot_dataset").relative_to(output_dir)), + }, + } + wall_deadline = _normalize_wall_deadline( + getattr(args, "wall_deadline", None) + ) + if wall_deadline is not None: + manifest["wall_deadline"] = wall_deadline + if args.initialization == base.TABLE13_EVAL_STATE0_INITIALIZATION: + manifest["seed_derivation"] = ( + "env_seed=42; policy_seed=retry_index_start+retry_index" + ) + manifest["fixed_env_seed"] = base.TABLE13_EVAL_ENV_SEED + else: + # Preserve the exact schema/values of legacy random-reset campaigns. + manifest["retry_seed_salt"] = str(args.retry_seed_salt) + manifest["seed_derivation"] = "sha256-task-retry-channel-mod-2^31-1" + if global_additional: + if global_success_counts_at_start is None: + raise ValueError("Missing frozen global success snapshot.") + manifest["select_global_successes_below"] = int( + args.select_global_successes_below + ) + manifest["target_additional_successes_per_variant"] = target_successes + manifest["global_successes_at_start_by_task"] = { + identity: int(count) + for identity, count in sorted(global_success_counts_at_start.items()) + } + elif target_successes != 1: + manifest["target_successes_per_variant"] = target_successes + manifest["base_successes_by_task"] = { + _task_identity(task): int( + base_success_counts[_task_identity(task)] + ) + for task in tasks + } + return manifest + + +def _write_or_validate_campaign_manifest( + path: Path, + manifest: dict[str, Any], +) -> dict[str, Any]: + if not path.exists(): + base._atomic_json_dump(manifest, path) + return manifest + existing = base._read_json(path) + comparable_existing = dict(existing) + comparable_new = dict(manifest) + comparable_existing.pop("created_at", None) + comparable_new.pop("created_at", None) + if comparable_existing != comparable_new: + raise ValueError( + f"Existing retry campaign is incompatible with this command: {path}" + ) + return existing + + +def _read_state(path: Path) -> dict[str, Any]: + if not path.exists(): + return { + "schema_version": STATE_SCHEMA_VERSION, + "active_runtime_seconds": 0.0, + "run_count": 0, + "status": "prepared", + "updated_at": _utc_now(), + } + state = base._read_json(path) + if int(state.get("schema_version", -1)) != STATE_SCHEMA_VERSION: + raise ValueError(f"Unsupported retry campaign state: {path}") + active_runtime = float(state.get("active_runtime_seconds", 0.0)) + if active_runtime < 0: + raise ValueError(f"Negative active runtime in {path}") + return state + + +def _campaign_task_scan( + output_dir: Path, + staging_root: Path, + tasks: list[base.CollectionTask], + *, + args: argparse.Namespace, + base_success_counts: dict[str, int], + recover: bool, +) -> tuple[dict[str, dict[str, Any]], int]: + checkpoint_identity = str(Path(args.checkpoint).expanduser().resolve()) + states: dict[str, dict[str, Any]] = {} + recovered = 0 + for task in tasks: + identity = _task_identity(task) + complete_indices: set[int] = set() + successful_indices: list[int] = [] + success_frames = 0 + for retry_index in range(int(args.retry_max_attempts)): + job = _retry_job( + task, + retry_index=retry_index, + retry_index_start=int(args.retry_index_start), + seed_salt=str(args.retry_seed_salt), + initialization=args.initialization, + init_state_index=_retry_init_state_index(args.initialization), + prompt_mode=args.prompt_mode, + ) + if recover: + recovered += int( + base._recover_staged_attempt( + output_dir, + staging_root, + job, + checkpoint=checkpoint_identity, + ) + ) + result_path = base._result_path(output_dir, job) + if not base._is_complete_attempt(result_path, job): + continue + result = base._read_json(result_path) + complete_indices.add(retry_index) + if bool(result["success"]): + successful_indices.append(retry_index) + success_frames += int(result.get("frames", 0)) + missing_indices = [ + index + for index in range(int(args.retry_max_attempts)) + if index not in complete_indices + ] + base_successes = int(base_success_counts[identity]) + target_successes = _effective_target_successes(args) + total_successes = base_successes + len(successful_indices) + target_reached = total_successes >= target_successes + states[identity] = { + "task": task, + "base_successes": base_successes, + "complete_indices": complete_indices, + "successful_indices": successful_indices, + "successful_index": ( + None if not successful_indices else successful_indices[0] + ), + "success_frames": success_frames, + "target_successes": target_successes, + "target_reached": target_reached, + "next_index": ( + None + if target_reached or not missing_indices + else missing_indices[0] + ), + "exhausted": ( + not target_reached + and len(complete_indices) >= int(args.retry_max_attempts) + ), + } + return states, recovered + + +def _refresh_task_state( + state: dict[str, Any], + *, + retry_max_attempts: int, +) -> None: + successful_indices = sorted(set(state["successful_indices"])) + state["successful_indices"] = successful_indices + state["successful_index"] = ( + None if not successful_indices else successful_indices[0] + ) + state["target_reached"] = ( + int(state["base_successes"]) + len(successful_indices) + >= int(state["target_successes"]) + ) + missing = [ + index + for index in range(int(retry_max_attempts)) + if index not in state["complete_indices"] + ] + state["next_index"] = ( + None + if state["target_reached"] or not missing + else missing[0] + ) + state["exhausted"] = ( + not state["target_reached"] and not missing + ) + + +def _snapshot( + states: dict[str, dict[str, Any]], + *, + active_runtime_seconds: float, + completion_reason: str | None = None, + blocked_tasks: set[str] | None = None, + infra_errors: int = 0, +) -> dict[str, Any]: + blocked_tasks = blocked_tasks or set() + by_difficulty: dict[str, dict[str, int]] = {} + retry_attempts = 0 + retry_successes = 0 + base_successes = 0 + retry_frames = 0 + exhausted = 0 + targets_reached = 0 + for identity, state in states.items(): + task: base.CollectionTask = state["task"] + retry_attempts += len(state["complete_indices"]) + task_retry_successes = len(state["successful_indices"]) + retry_successes += task_retry_successes + base_successes += int(state["base_successes"]) + target_reached = bool(state["target_reached"]) + targets_reached += int(target_reached) + retry_frames += int(state["success_frames"]) + exhausted += int(bool(state["exhausted"])) + difficulty = str(task.difficulty_level) + row = by_difficulty.setdefault( + difficulty, + { + "targets": 0, + "attempts": 0, + "base_successes": 0, + "retry_successes": 0, + "rescued": 0, + "exhausted": 0, + "blocked": 0, + }, + ) + row["targets"] += 1 + row["attempts"] += len(state["complete_indices"]) + row["base_successes"] += int(state["base_successes"]) + row["retry_successes"] += task_retry_successes + row["rescued"] += int(target_reached) + row["exhausted"] += int(bool(state["exhausted"])) + row["blocked"] += int(identity in blocked_tasks) + return { + "updated_at": _utc_now(), + "completion_reason": completion_reason, + "targets": len(states), + "retry_attempts": retry_attempts, + "base_successes": base_successes, + "retry_successes": retry_successes, + "retry_failures": retry_attempts - retry_successes, + "retry_success_frames": retry_frames, + "rescued_variants": targets_reached, + "targets_reached": targets_reached, + "exhausted_variants": exhausted, + "blocked_variants": len(blocked_tasks), + "unresolved_variants": len(states) - targets_reached - exhausted, + "infra_errors": int(infra_errors), + "active_runtime_seconds": float(active_runtime_seconds), + "by_difficulty": { + key: by_difficulty[key] + for key in sorted(by_difficulty, key=int) + }, + } + + +def _global_attempt_summary(output_dir: Path) -> dict[str, Any]: + complete = 0 + successes = 0 + frames = 0 + invalid: list[str] = [] + for path in sorted((output_dir / "attempts").glob("*.json")): + try: + result = base._read_json(path) + job = base.CollectionJob.from_dict(result["job"]) + expected_path = base._result_path(output_dir, job) + if ( + path != expected_path + or result.get("job_id") != job.job_id + or not base._is_complete_attempt(path, job) + or ( + result["success"] is False + and (output_dir / "staging" / "success" / job.job_id).exists() + ) + ): + invalid.append(str(path)) + continue + complete += 1 + if bool(result["success"]): + successes += 1 + frames += int(result.get("frames", 0)) + except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError): + invalid.append(str(path)) + if invalid: + raise ValueError(f"Invalid global attempt files: {invalid[:5]}") + return { + "jobs": complete, + "complete": complete, + "successes": successes, + "failures": complete - successes, + "frames": frames, + "success_rate": successes / complete if complete else 0.0, + } + + +def _worker_args( + args: argparse.Namespace, + *, + output_dir: Path, + staging_root: Path, + tasks: list[base.CollectionTask], + libero_config_dir: Path, +) -> dict[str, Any]: + return { + "parent_pid": os.getpid(), + "output_dir": str(output_dir), + "staging_root": str(staging_root), + "config_dir": str(Path(args.config_dir).expanduser().resolve()), + "task_config": args.task_config, + "checkpoint": str(Path(args.checkpoint).expanduser().resolve()), + "checkpoint_load_path": str( + Path(args.checkpoint_load_path or args.checkpoint).expanduser().resolve() + ), + "dataset_stats": str(Path(args.dataset_stats).expanduser().resolve()), + "gripper_action_format": args.gripper_action_format, + "hydra_overrides": list(args.override), + "text_embedding_cache_dir": ( + None + if not args.text_embedding_cache_dir + else str(Path(args.text_embedding_cache_dir).expanduser().resolve()) + ), + "model_task_descriptions": sorted( + base._model_task_descriptions( + tasks, + prompt_mode=args.prompt_mode, + ) + ), + "libero_plus_root": str( + Path(args.libero_plus_root).expanduser().resolve() + ), + "libero_config_dir": str(libero_config_dir), + "egl_lock_file": os.environ.get( + "LIBERO_EGL_INIT_LOCK_FILE", + "/tmp/fastwam_libero_mujoco_egl_init.lock", + ), + "egl_lock_scope": args.egl_lock_scope, + "egl_fallback_gpu": args.egl_fallback_gpu, + "model_base_path": str(Path(args.model_base_path).expanduser().resolve()), + } + + +def _run_workers( + args: argparse.Namespace, + *, + output_dir: Path, + campaign_dir: Path, + staging_root: Path, + tasks: list[base.CollectionTask], + states: dict[str, dict[str, Any]], + state_path: Path, + persisted_state: dict[str, Any], +) -> tuple[str, set[str], int, float, bool]: + gpu_ids = base._parse_gpu_ids(args.gpus) + worker_slots = [ + (f"gpu{gpu_id}-w{replica_id}", gpu_id) + for gpu_id in gpu_ids + for replica_id in range(int(args.workers_per_gpu)) + ] + runnable = [ + state + for state in states.values() + if not state["target_reached"] + and not state["exhausted"] + and state["next_index"] is not None + ] + runnable.sort( + key=lambda state: ( + int(state["base_successes"]) + + len(state["successful_indices"]), + -int(state["task"].difficulty_level or 0), + state["task"].suite, + state["task"].task_id, + ) + ) + if not runnable: + return ( + "all_targets_rescued_or_exhausted", + set(), + 0, + float(persisted_state.get("active_runtime_seconds", 0.0)), + False, + ) + + max_runtime_seconds = float(args.max_runtime_hours) * 3600.0 + runtime_base = float(persisted_state.get("active_runtime_seconds", 0.0)) + if runtime_base >= max_runtime_seconds: + return "runtime_budget_already_consumed", set(), 0, runtime_base, False + wall_deadline_timestamp = _wall_deadline_timestamp( + getattr(args, "wall_deadline", None) + ) + if ( + wall_deadline_timestamp is not None + and time.time() >= wall_deadline_timestamp + ): + return "wall_deadline_already_reached", set(), 0, runtime_base, False + + libero_config_dir = base._ensure_libero_config( + Path(args.libero_plus_root).expanduser().resolve(), + output_dir, + ) + worker_args = _worker_args( + args, + output_dir=output_dir, + staging_root=staging_root, + tasks=tasks, + libero_config_dir=libero_config_dir, + ) + context = mp.get_context("spawn") + task_queue = context.Queue() + task_queue.cancel_join_thread() + status_queue = context.Queue() + start_event = context.Event() + stop_event = context.Event() + + workers = [ + context.Process( + target=base._worker_main, + args=( + worker_slot, + gpu_id, + worker_args, + task_queue, + status_queue, + start_event, + stop_event, + ), + name=f"libero-retry-{worker_slot}", + ) + for worker_slot, gpu_id in worker_slots + ] + process_by_slot = { + worker_slot: worker + for (worker_slot, _), worker in zip(worker_slots, workers, strict=True) + } + outstanding: set[str] = set() + jobs_by_id: dict[str, tuple[base.CollectionJob, str, int]] = {} + running_by_worker: dict[str, str] = {} + processed_executions: set[tuple[str, str]] = set() + infra_retry_counts: Counter[str] = Counter() + blocked_tasks: set[str] = set() + infra_errors = 0 + enqueued_this_run = 0 + invocation_limit = int(args.max_new_attempts) + + def can_enqueue() -> bool: + return invocation_limit <= 0 or enqueued_this_run < invocation_limit + + def enqueue_job(task_state: dict[str, Any], retry_index: int) -> bool: + nonlocal enqueued_this_run + if not can_enqueue(): + return False + task: base.CollectionTask = task_state["task"] + identity = _task_identity(task) + job = _retry_job( + task, + retry_index=retry_index, + retry_index_start=int(args.retry_index_start), + seed_salt=str(args.retry_seed_salt), + initialization=args.initialization, + init_state_index=_retry_init_state_index(args.initialization), + prompt_mode=args.prompt_mode, + ) + if job.job_id in outstanding: + raise RuntimeError(f"Job is already outstanding: {job.job_id}") + jobs_by_id[job.job_id] = (job, identity, retry_index) + outstanding.add(job.job_id) + task_queue.put(job.to_dict()) + enqueued_this_run += 1 + return True + + for task_state in runnable: + if not enqueue_job(task_state, int(task_state["next_index"])): + break + + if not outstanding: + return ( + "invocation_attempt_limit", + blocked_tasks, + infra_errors, + runtime_base, + False, + ) + + previous_signal_handlers = { + signal.SIGINT: signal.getsignal(signal.SIGINT), + signal.SIGTERM: signal.getsignal(signal.SIGTERM), + } + stop_requested = False + interrupted = False + completion_reason: str | None = None + stop_deadline: float | None = None + sentinels_sent = False + + def request_stop(_signum, _frame) -> None: + nonlocal stop_requested, interrupted, completion_reason, stop_deadline + if not stop_requested: + print( + "[retry-interrupt] stopping after in-flight attempts; campaign " + "state remains resumable and the existing dataset is unchanged.", + flush=True, + ) + stop_deadline = time.time() + MAX_GRACEFUL_STOP_SECONDS + stop_requested = True + interrupted = True + completion_reason = "interrupted" + stop_event.set() + start_event.set() + + def send_sentinels() -> None: + nonlocal sentinels_sent + if sentinels_sent: + return + for _ in workers: + task_queue.put(None) + sentinels_sent = True + + def request_internal_stop(reason: str, message: str) -> None: + nonlocal stop_requested, completion_reason, stop_deadline + if not stop_requested: + print(message, flush=True) + completion_reason = reason + stop_deadline = time.time() + MAX_GRACEFUL_STOP_SECONDS + stop_requested = True + stop_event.set() + start_event.set() + send_sentinels() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + for worker in workers: + worker.start() + workers_started_at = time.time() + + ready: set[str] = set() + fatal_slots: set[str] = set() + run_active_started_monotonic: float | None = None + rollout_started_at: float | None = None + last_status_time = 0.0 + completed_this_run = 0 + successes_this_run = 0 + + def active_runtime_now() -> float: + if bool(args.exclude_runtime_budget) or run_active_started_monotonic is None: + return runtime_base + return runtime_base + (time.monotonic() - run_active_started_monotonic) + + def persist_running_state(status: str = "running") -> None: + payload = dict(persisted_state) + payload.update( + { + "schema_version": STATE_SCHEMA_VERSION, + "status": status, + "active_runtime_seconds": active_runtime_now(), + "run_count": int(persisted_state.get("run_count", 0)) + 1, + "updated_at": _utc_now(), + "ready_workers": len(ready), + "worker_count": len(workers), + "outstanding_jobs": len(outstanding), + "enqueued_this_run": enqueued_this_run, + "completed_this_run": completed_this_run, + "successes_this_run": successes_this_run, + "infra_errors_this_run": infra_errors, + } + ) + base._atomic_json_dump(payload, state_path) + + def complete_job(job_id: str, worker_slot: str) -> None: + nonlocal completed_this_run, successes_this_run, infra_errors + execution_key = (worker_slot, job_id) + if execution_key in processed_executions: + return + if job_id not in jobs_by_id or job_id not in outstanding: + processed_executions.add(execution_key) + return + processed_executions.add(execution_key) + job, identity, retry_index = jobs_by_id[job_id] + if running_by_worker.get(worker_slot) == job_id: + running_by_worker.pop(worker_slot, None) + + checkpoint_identity = str(Path(args.checkpoint).expanduser().resolve()) + try: + base._recover_staged_attempt( + output_dir, + staging_root, + job, + checkpoint=checkpoint_identity, + ) + except Exception as error: + infra_errors += 1 + outstanding.discard(job_id) + jobs_by_id.pop(job_id, None) + blocked_tasks.add(identity) + print( + f"[retry-recovery-error] task={identity} retry={retry_index} " + f"job={job_id} error={error}", + flush=True, + ) + request_internal_stop( + "recovery_error", + "[retry-recovery-stop] staging/attempt integrity recovery failed; " + "stopping without publishing a merged dataset.", + ) + return + + result_path = base._result_path(output_dir, job) + if base._is_complete_attempt(result_path, job): + outstanding.discard(job_id) + jobs_by_id.pop(job_id, None) + result = base._read_json(result_path) + success = bool(result["success"]) + state = states[identity] + state["complete_indices"].add(retry_index) + completed_this_run += 1 + if success: + state["successful_indices"].append(retry_index) + state["success_frames"] += int(result.get("frames", 0)) + successes_this_run += 1 + _refresh_task_state( + state, + retry_max_attempts=int(args.retry_max_attempts), + ) + if ( + not stop_requested + and not state["target_reached"] + and not state["exhausted"] + and state["next_index"] is not None + ): + runtime_exhausted = ( + not bool(args.exclude_runtime_budget) + and run_active_started_monotonic is not None + and active_runtime_now() >= max_runtime_seconds + ) + wall_deadline_reached = ( + wall_deadline_timestamp is not None + and time.time() >= wall_deadline_timestamp + ) + if wall_deadline_reached: + request_internal_stop( + "wall_deadline", + "[retry-wall-deadline] absolute wall deadline reached; " + "finishing in-flight attempts.", + ) + elif runtime_exhausted: + request_internal_stop( + "runtime_budget", + "[retry-deadline] active runtime budget reached; " + "finishing in-flight attempts.", + ) + else: + enqueue_job(state, int(state["next_index"])) + return + + outstanding.discard(job_id) + infra_errors += 1 + infra_retry_counts[job_id] += 1 + if ( + infra_retry_counts[job_id] <= int(args.max_infra_retries) + and not stop_requested + and can_enqueue() + ): + print( + f"[retry-infra-retry] task={identity} retry={retry_index} " + f"job={job_id} infra_attempt={infra_retry_counts[job_id]}/" + f"{args.max_infra_retries}", + flush=True, + ) + enqueue_job(states[identity], retry_index) + else: + jobs_by_id.pop(job_id, None) + blocked_tasks.add(identity) + print( + f"[retry-blocked] task={identity} retry={retry_index} job={job_id}", + flush=True, + ) + + try: + while any(worker.is_alive() for worker in workers): + try: + event = status_queue.get(timeout=2.0) + except queue.Empty: + event = None + if event is not None: + event_type = str(event.get("type")) + worker_slot = str(event.get("worker_slot", "")) + if event_type == "ready": + ready.add(worker_slot) + print( + f"[worker-ready] slot={worker_slot} gpu={event['gpu_id']} " + f"ready={len(ready)}/{len(workers)}", + flush=True, + ) + if len(ready) == len(workers): + run_active_started_monotonic = time.monotonic() + start_event.set() + persist_running_state() + elif event_type == "started": + started_job_id = str(event["job_id"]) + processed_executions.discard((worker_slot, started_job_id)) + running_by_worker[worker_slot] = started_job_id + if rollout_started_at is None: + rollout_started_at = time.time() + elif event_type == "done": + print( + f"[retry-attempt-done] slot={worker_slot} " + f"gpu={event['gpu_id']} task={event['task']} " + f"success={event['success']} frames={event['frames']} " + f"job={event['job_id']}", + flush=True, + ) + complete_job(str(event["job_id"]), worker_slot) + elif event_type == "failed": + print( + f"[retry-attempt-error] slot={worker_slot} " + f"gpu={event['gpu_id']} task={event['task']} " + f"error={event['error']} log={event['log']}", + flush=True, + ) + complete_job(str(event["job_id"]), worker_slot) + elif event_type == "fatal": + if worker_slot not in fatal_slots: + fatal_slots.add(worker_slot) + infra_errors += 1 + print( + f"[worker-fatal] slot={worker_slot} gpu={event['gpu_id']} " + f"log={event['log']}\n{event['error']}", + flush=True, + ) + request_internal_stop( + "worker_failure", + "[retry-worker-stop] a rollout worker failed; stopping all " + "workers so the campaign can be safely resumed.", + ) + + for worker_slot, worker in process_by_slot.items(): + if worker.exitcode is None or worker_slot in fatal_slots: + continue + job_id = running_by_worker.pop(worker_slot, None) + if job_id is not None: + print( + f"[worker-exit-recovery] slot={worker_slot} " + f"exitcode={worker.exitcode} job={job_id}", + flush=True, + ) + complete_job(job_id, worker_slot) + if not stop_requested: + fatal_slots.add(worker_slot) + infra_errors += 1 + request_internal_stop( + "worker_failure", + f"[retry-worker-exit] slot={worker_slot} exited unexpectedly " + f"with code={worker.exitcode}; stopping for a safe resume.", + ) + + if ( + not start_event.is_set() + and not stop_requested + and time.time() - workers_started_at > float(args.worker_ready_timeout) + ): + waiting = sorted(set(process_by_slot) - ready) + print( + f"[worker-ready-timeout] waiting={waiting}", + flush=True, + ) + stop_requested = True + completion_reason = "worker_ready_timeout" + stop_deadline = time.time() + MAX_GRACEFUL_STOP_SECONDS + stop_event.set() + start_event.set() + + now = time.time() + active_runtime = active_runtime_now() + if ( + wall_deadline_timestamp is not None + and time.time() >= wall_deadline_timestamp + and not stop_requested + ): + request_internal_stop( + "wall_deadline", + "[retry-wall-deadline] absolute wall deadline reached; " + "finishing in-flight attempts and preserving queued work.", + ) + if ( + not bool(args.exclude_runtime_budget) + and run_active_started_monotonic is not None + and active_runtime >= max_runtime_seconds + and not stop_requested + ): + print( + "[retry-deadline] active runtime budget reached; finishing " + "in-flight attempts and preserving queued work for resume.", + flush=True, + ) + stop_requested = True + completion_reason = "runtime_budget" + stop_deadline = now + MAX_GRACEFUL_STOP_SECONDS + stop_event.set() + send_sentinels() + + if ( + int(args.max_total_infra_errors) > 0 + and infra_errors >= int(args.max_total_infra_errors) + and not stop_requested + ): + print( + f"[retry-infra-limit] errors={infra_errors}; stopping campaign.", + flush=True, + ) + stop_requested = True + completion_reason = "infra_error_limit" + stop_deadline = now + MAX_GRACEFUL_STOP_SECONDS + stop_event.set() + send_sentinels() + + if not outstanding and not stop_requested: + completion_reason = ( + "blocked_variants" + if blocked_tasks + else ( + "invocation_attempt_limit" + if invocation_limit > 0 + else "all_targets_rescued_or_exhausted" + ) + ) + stop_requested = True + stop_deadline = now + MAX_GRACEFUL_STOP_SECONDS + send_sentinels() + + if stop_requested: + send_sentinels() + if stop_deadline is not None and now >= stop_deadline: + print( + "[retry-stop] graceful-stop deadline reached; terminating " + "remaining workers.", + flush=True, + ) + break + + if now - last_status_time >= float(args.status_every): + solved = sum( + state["target_reached"] for state in states.values() + ) + collected_successes = sum( + len(state["successful_indices"]) for state in states.values() + ) + exhausted = sum(bool(state["exhausted"]) for state in states.values()) + remaining_runtime = max(0.0, max_runtime_seconds - active_runtime) + print( + f"[retry-status] targets={len(states)} reached={solved} " + f"retry_successes={collected_successes} " + f"exhausted={exhausted} blocked={len(blocked_tasks)} " + f"completed_this_run={completed_this_run} " + f"successes_this_run={successes_this_run} " + f"outstanding={len(outstanding)} " + f"ready={len(ready)}/{len(workers)} " + f"active_runtime={base._format_duration(active_runtime)} " + f"remaining_runtime={base._format_duration(remaining_runtime)} " + f"infra_errors={infra_errors}", + flush=True, + ) + persist_running_state() + last_status_time = now + finally: + stop_event.set() + start_event.set() + send_sentinels() + shutdown_deadline = ( + stop_deadline + if stop_deadline is not None + else time.time() + 30.0 + ) + while any(worker.is_alive() for worker in workers): + if time.time() >= shutdown_deadline: + break + for worker in workers: + if worker.is_alive(): + worker.join(timeout=0.2) + for worker in workers: + if worker.is_alive(): + worker.terminate() + worker.join(timeout=5) + if worker.is_alive(): + worker.kill() + worker.join(timeout=5) + + while True: + try: + event = status_queue.get_nowait() + except queue.Empty: + break + if event.get("type") in {"done", "failed"}: + complete_job( + str(event["job_id"]), + str(event.get("worker_slot", "")), + ) + elif event.get("type") == "fatal": + infra_errors += 1 + + for signum, previous_handler in previous_signal_handlers.items(): + signal.signal(signum, previous_handler) + + active_runtime = active_runtime_now() + if completion_reason is None: + completion_reason = "workers_exited" + final_state = dict(persisted_state) + final_state.update( + { + "schema_version": STATE_SCHEMA_VERSION, + "status": completion_reason, + "active_runtime_seconds": active_runtime, + "run_count": int(persisted_state.get("run_count", 0)) + 1, + "updated_at": _utc_now(), + "ready_workers": len(ready), + "worker_count": len(workers), + "enqueued_this_run": enqueued_this_run, + "completed_this_run": completed_this_run, + "successes_this_run": successes_this_run, + "infra_errors_this_run": infra_errors, + } + ) + base._atomic_json_dump(final_state, state_path) + return completion_reason, blocked_tasks, infra_errors, active_runtime, interrupted + + +def _finalize( + args: argparse.Namespace, + *, + output_dir: Path, + campaign_dir: Path, + staging_root: Path, + campaign_summary: dict[str, Any], +) -> dict[str, Any]: + from lerobot_rollout_writer import finalize_dataset + + dataset_root = output_dir / "lerobot_dataset" + reference_dataset = Path(args.reference_dataset).expanduser().resolve() + # Validate the global attempt ledger before publishing a replacement + # dataset. This keeps a malformed attempt file from creating a new + # dataset/old-report mismatch. + global_attempts = _global_attempt_summary(output_dir) + build_report = finalize_dataset( + staging_root, + dataset_root, + reference_dataset=reference_dataset, + ) + report = { + "schema_version": 2, + "updated_at": _utc_now(), + "attempts": global_attempts, + "retry_campaign": campaign_summary, + "build": build_report, + "validation": dict(build_report), + } + base._atomic_json_dump(report, campaign_dir / "final_report.json") + base._atomic_json_dump(report, output_dir / "final_report.json") + return report + + +def _should_finalize( + *, + skip_finalize: bool, + interrupted: bool, + completion_reason: str, +) -> bool: + return ( + not skip_finalize + and not interrupted + and completion_reason + not in { + "blocked_variants", + "infra_error_limit", + "invocation_attempt_limit", + "recovery_error", + "worker_failure", + "worker_ready_timeout", + "workers_exited", + } + ) + + +def _run_locked(args: argparse.Namespace, output_dir: Path) -> int: + campaign_dir = _campaign_dir(output_dir, args.campaign_name) + campaign_dir.mkdir(parents=True, exist_ok=True) + campaign_manifest_path = campaign_dir / "manifest.json" + _validate_seed_interval_against_campaigns( + output_dir, + campaign_dir, + retry_index_start=int(args.retry_index_start), + retry_max_attempts=int(args.retry_max_attempts), + ) + existing_campaign_manifest = ( + None + if not campaign_manifest_path.is_file() + else base._read_json(campaign_manifest_path) + ) + staging_root = output_dir / "staging" + base_manifest, base_tasks, base_jobs = _read_base_collection(output_dir) + + libero_plus_root = Path(args.libero_plus_root).expanduser().resolve() + classification_path = ( + Path(args.classification).expanduser().resolve() + if args.classification + else base.default_classification_path(libero_plus_root) + ) + input_files = base._validate_inputs(args, classification_path) + _validate_base_semantics( + base_manifest, + args, + classification_path=classification_path, + input_files=input_files, + ) + base_success_counts = _base_success_counts( + output_dir, + base_tasks, + base_jobs, + ) + global_success_counts_at_start: dict[str, int] | None = None + if _uses_global_additional_target(args): + if existing_campaign_manifest is None: + global_success_counts_at_start = _global_success_counts( + output_dir, + base_tasks, + ) + retry_tasks = _base_below_target_tasks( + base_tasks, + global_success_counts_at_start, + target_successes_per_variant=int( + args.select_global_successes_below + ), + ) + else: + if existing_campaign_manifest.get("policy") != ( + "global_below_threshold_until_additional_success_target" + ): + raise ValueError( + "Existing campaign does not use global-additional selection." + ) + global_success_counts_at_start = { + str(identity): int(count) + for identity, count in existing_campaign_manifest[ + "global_successes_at_start_by_task" + ].items() + } + retry_tasks = [ + base.CollectionTask.from_dict(value) + for value in existing_campaign_manifest["tasks"] + ] + campaign_start_success_counts = { + _task_identity(task): 0 for task in retry_tasks + } + selection_counts = global_success_counts_at_start + else: + retry_tasks = _base_below_target_tasks( + base_tasks, + base_success_counts, + target_successes_per_variant=int( + args.target_successes_per_variant + ), + ) + campaign_start_success_counts = base_success_counts + selection_counts = base_success_counts + _validate_tasks_against_current_inputs( + retry_tasks, + classification_path=classification_path, + reference_dataset=Path(args.reference_dataset).expanduser().resolve(), + ) + retry_tasks.sort( + key=lambda task: ( + selection_counts[_task_identity(task)], + -int(task.difficulty_level or 0), + task.suite, + task.task_id, + ) + ) + base._prepare_libero_runtime(libero_plus_root, output_dir) + base._preflight_text_embedding_cache( + args, + retry_tasks, + output_dir=output_dir, + physical_gpu_id=base._parse_gpu_ids(args.gpus)[0], + ) + manifest = _campaign_manifest( + args, + output_dir=output_dir, + base_manifest_path=output_dir / "collection_manifest.json", + input_files=input_files, + tasks=retry_tasks, + base_success_counts=base_success_counts, + global_success_counts_at_start=global_success_counts_at_start, + ) + _write_or_validate_campaign_manifest(campaign_manifest_path, manifest) + + base_report_backup = campaign_dir / "base_final_report.json" + if not base_report_backup.exists() and (output_dir / "final_report.json").is_file(): + base._atomic_json_dump( + base._read_json(output_dir / "final_report.json"), + base_report_backup, + ) + + states, recovered = _campaign_task_scan( + output_dir, + staging_root, + retry_tasks, + args=args, + base_success_counts=campaign_start_success_counts, + recover=bool(args.resume and not args.prepare_only), + ) + difficulty_targets = Counter( + int(task.difficulty_level or 0) for task in retry_tasks + ) + planned_job_ids = { + _retry_job( + task, + retry_index=retry_index, + retry_index_start=int(args.retry_index_start), + seed_salt=str(args.retry_seed_salt), + initialization=args.initialization, + init_state_index=_retry_init_state_index(args.initialization), + prompt_mode=args.prompt_mode, + ).job_id + for task in retry_tasks + for retry_index in range(int(args.retry_max_attempts)) + } + base_job_ids = {job.job_id for job in base_jobs} + if len(planned_job_ids) != len(retry_tasks) * int(args.retry_max_attempts): + raise RuntimeError("Retry seed derivation produced duplicate job IDs.") + collision = planned_job_ids & base_job_ids + if collision: + raise RuntimeError(f"Retry jobs collide with base jobs: {sorted(collision)[:3]}") + + state_path = campaign_dir / "state.json" + persisted_state = _read_state(state_path) + snapshot = _snapshot( + states, + active_runtime_seconds=float( + persisted_state.get("active_runtime_seconds", 0.0) + ), + ) + print( + f"Prepared adaptive retry campaign: name={args.campaign_name} " + f"targets={len(retry_tasks)} " + f"target_successes_per_variant={_effective_target_successes(args)} " + f"difficulty={dict(sorted(difficulty_targets.items()))} " + f"max_attempts_per_variant={args.retry_max_attempts} " + f"planned_job_ids={len(planned_job_ids)} " + f"recovered={recovered} " + f"already_rescued={snapshot['rescued_variants']} " + f"already_exhausted={snapshot['exhausted_variants']} " + f"active_runtime={base._format_duration(snapshot['active_runtime_seconds'])} " + f"runtime_budget={base._format_duration(float(args.max_runtime_hours) * 3600.0)} " + f"workers={len(base._parse_gpu_ids(args.gpus)) * int(args.workers_per_gpu)} " + f"output={output_dir}", + flush=True, + ) + if args.prepare_only: + base._atomic_json_dump(snapshot, campaign_dir / "summary.json") + return 0 + + completion_reason, blocked_tasks, infra_errors, active_runtime, interrupted = ( + _run_workers( + args, + output_dir=output_dir, + campaign_dir=campaign_dir, + staging_root=staging_root, + tasks=retry_tasks, + states=states, + state_path=state_path, + persisted_state=persisted_state, + ) + ) + states, recovered_after = _campaign_task_scan( + output_dir, + staging_root, + retry_tasks, + args=args, + base_success_counts=campaign_start_success_counts, + recover=True, + ) + summary = _snapshot( + states, + active_runtime_seconds=active_runtime, + completion_reason=completion_reason, + blocked_tasks=blocked_tasks, + infra_errors=infra_errors, + ) + summary["recovered_attempts"] = recovered + recovered_after + base._atomic_json_dump(summary, campaign_dir / "summary.json") + print( + f"Retry campaign stopped: reason={completion_reason} " + f"attempts={summary['retry_attempts']} " + f"rescued={summary['rescued_variants']}/{summary['targets']} " + f"exhausted={summary['exhausted_variants']} " + f"blocked={summary['blocked_variants']} " + f"active_runtime={base._format_duration(active_runtime)}", + flush=True, + ) + + should_finalize = _should_finalize( + skip_finalize=bool(args.skip_finalize), + interrupted=interrupted, + completion_reason=completion_reason, + ) + if not should_finalize: + print( + "Retry finalization skipped; the previously published dataset remains " + "valid and new successful staging will be merged on a later resume.", + flush=True, + ) + if interrupted: + return 130 + if completion_reason in { + "blocked_variants", + "infra_error_limit", + "recovery_error", + "worker_failure", + "worker_ready_timeout", + "workers_exited", + }: + return 1 + return 0 + + report = _finalize( + args, + output_dir=output_dir, + campaign_dir=campaign_dir, + staging_root=staging_root, + campaign_summary=summary, + ) + print( + f"Finalized merged LeRobot dataset: " + f"episodes={report['validation']['total_episodes']} " + f"frames={report['validation']['total_frames']} " + f"rescued={summary['rescued_variants']} " + f"output={report['validation']['dataset_root']}", + flush=True, + ) + return 0 + + +def run(args: argparse.Namespace) -> int: + global_additional = _uses_global_additional_target(args) + if args.task_source != "table13": + raise ValueError("Adaptive failed-variant retry requires --task-source table13.") + if args.initialization not in { + base.RANDOM_RESET_INITIALIZATION, + base.TABLE13_EVAL_STATE0_INITIALIZATION, + }: + raise ValueError( + "Adaptive retry supports --initialization random_reset or " + f"{base.TABLE13_EVAL_STATE0_INITIALIZATION}." + ) + if args.prompt_mode not in {"canonical", "benchmark"}: + raise ValueError( + "Adaptive failed-variant retry requires --prompt-mode canonical " + "or benchmark." + ) + if not args.resume: + raise ValueError( + "Retrying inside an existing collection requires the explicit --resume flag." + ) + if int(args.retry_index_start) < 0: + raise ValueError("--retry-index-start must be non-negative.") + if int(args.retry_max_attempts) < 1: + raise ValueError("--retry-max-attempts must be positive.") + if int(args.target_successes_per_variant) < 1: + raise ValueError("--target-successes-per-variant must be positive.") + if global_additional: + if int(args.target_successes_per_variant) != 1: + raise ValueError( + "Global-additional mode cannot be combined with a non-default " + "--target-successes-per-variant." + ) + if int(args.select_global_successes_below) < 1: + raise ValueError("--select-global-successes-below must be positive.") + if int(args.target_additional_successes_per_variant) < 1: + raise ValueError( + "--target-additional-successes-per-variant must be positive." + ) + _normalize_wall_deadline(getattr(args, "wall_deadline", None)) + if float(args.max_runtime_hours) <= 0: + raise ValueError("--max-runtime-hours must be positive.") + if int(args.max_new_attempts) < 0: + raise ValueError("--max-new-attempts must be non-negative.") + if int(args.max_infra_retries) < 0: + raise ValueError("--max-infra-retries must be non-negative.") + if int(args.max_total_infra_errors) < 0: + raise ValueError("--max-total-infra-errors must be non-negative.") + + output_dir = Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + lock_path = output_dir / ".collector.lock" + lock_handle = lock_path.open("a+", encoding="utf-8") + try: + fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as error: + lock_handle.close() + raise RuntimeError( + f"Another collector already holds the output lock: {lock_path}" + ) from error + try: + lock_handle.seek(0) + lock_handle.truncate() + lock_handle.write( + f"pid={os.getpid()}\ncampaign={args.campaign_name}\n" + ) + lock_handle.flush() + return _run_locked(args, output_dir) + finally: + fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN) + lock_handle.close() + + +def build_parser() -> argparse.ArgumentParser: + parser = base.build_parser() + parser.description = __doc__ + parser.add_argument("--campaign-name", default=DEFAULT_CAMPAIGN_NAME) + parser.add_argument( + "--retry-index-start", + type=int, + default=44, + help="Logical retry number used as an input to deterministic seed derivation.", + ) + parser.add_argument( + "--retry-max-attempts", + type=int, + default=96, + help="Maximum additional rollout attempts per under-target base variant.", + ) + parser.add_argument( + "--target-successes-per-variant", + type=int, + default=1, + help=( + "Required cumulative base + retry successes for each exact variant. " + "The default 1 preserves the original zero-success retry behavior." + ), + ) + parser.add_argument( + "--select-global-successes-below", + type=int, + default=None, + help=( + "On first launch, freeze variants whose validated global success " + "count is below this threshold." + ), + ) + parser.add_argument( + "--target-additional-successes-per-variant", + type=int, + default=None, + help=( + "With --select-global-successes-below, collect this many new " + "successes per frozen variant within the current campaign." + ), + ) + parser.add_argument( + "--retry-seed-salt", + default=DEFAULT_RETRY_SEED_SALT, + help=( + "Stable salt for random_reset env/policy SHA256 seed derivation. " + "table13_eval_state0 instead uses fixed env_seed=42 and the " + "retry-number sequence as policy_seed." + ), + ) + parser.add_argument( + "--max-runtime-hours", + type=float, + default=12.0, + help="Cumulative active rollout runtime budget across resumes.", + ) + parser.add_argument( + "--max-new-attempts", + type=int, + default=0, + help="Per-invocation enqueue cap; zero is unlimited. Useful for smoke tests.", + ) + parser.add_argument( + "--max-infra-retries", + type=int, + default=3, + help="Same-seed retries after worker exceptions before blocking a variant.", + ) + parser.add_argument( + "--max-total-infra-errors", + type=int, + default=32, + help="Stop the invocation after this many infrastructure errors; zero disables.", + ) + parser.add_argument( + "--skip-finalize", + action="store_true", + help="Leave new successes in staging and keep the existing published dataset.", + ) + parser.add_argument( + "--exclude-runtime-budget", + action="store_true", + help="Do not charge this invocation to the cumulative wall-clock budget.", + ) + parser.add_argument( + "--wall-deadline", + default=None, + help=( + "Absolute timezone-aware ISO-8601 deadline. Reaching it performs " + "an internal graceful stop and still finalizes merged successes." + ), + ) + return parser + + +def main() -> None: + raise SystemExit(run(build_parser().parse_args())) + + +if __name__ == "__main__": + main() diff --git a/experiments/libero/run_libero_manager.py b/experiments/libero/run_libero_manager.py index 4868d647..f6de8171 100644 --- a/experiments/libero/run_libero_manager.py +++ b/experiments/libero/run_libero_manager.py @@ -10,20 +10,50 @@ from omegaconf import DictConfig, OmegaConf -def create_task_file(output_file: Path, task_suite_names: list[str]) -> Path: +def create_task_file( + output_file: Path, + task_suite_names: list[str], + task_specs: list[str] | None = None, +) -> Path: benchmark_dict = benchmark.get_benchmark_dict() output_file.parent.mkdir(parents=True, exist_ok=True) total_tasks = 0 + seen_tasks = set() with output_file.open("w", encoding="utf-8") as f: - for suite_name in task_suite_names: - task_suite = benchmark_dict[suite_name]() - n_tasks = int(task_suite.n_tasks) - print(f"\n{suite_name}:") - print(f"- Number of tasks: {n_tasks}") - for task_id in range(n_tasks): + if task_specs: + for raw_spec in task_specs: + spec = str(raw_spec) + try: + suite_name, raw_task_id = spec.rsplit("/", 1) + task_id = int(raw_task_id) + except ValueError as exc: + raise ValueError( + f"Invalid task spec {spec!r}; expected '/'." + ) from exc + if suite_name not in benchmark_dict: + raise ValueError(f"Unknown LIBERO task suite in {spec!r}: {suite_name}") + n_tasks = int(benchmark_dict[suite_name]().n_tasks) + if not 0 <= task_id < n_tasks: + raise ValueError( + f"Task id {task_id} is out of range for {suite_name} (0-{n_tasks - 1})." + ) + task = (suite_name, task_id) + if task in seen_tasks: + raise ValueError(f"Duplicate task spec: {spec!r}") + seen_tasks.add(task) f.write(f"{suite_name},{task_id}\n") total_tasks += 1 + print(f"Selected task: {suite_name}/{task_id}") + else: + for suite_name in task_suite_names: + task_suite = benchmark_dict[suite_name]() + n_tasks = int(task_suite.n_tasks) + print(f"\n{suite_name}:") + print(f"- Number of tasks: {n_tasks}") + for task_id in range(n_tasks): + f.write(f"{suite_name},{task_id}\n") + total_tasks += 1 print(f"\nTask list created: {output_file}") print(f"Total tasks: {total_tasks}") @@ -139,7 +169,13 @@ def main(cfg: DictConfig): task_file = Path(os.path.expanduser(os.path.expandvars(str(task_file_cfg)))) else: task_file = output_dir / "tasks.txt" - task_file = create_task_file(task_file, list(manager.task_suite_names)) + task_specs_cfg = manager.get("task_specs") + task_specs = list(task_specs_cfg) if task_specs_cfg else None + task_file = create_task_file( + task_file, + list(manager.task_suite_names), + task_specs, + ) OmegaConf.save(config=cfg, f=str(output_dir / "manager_config.yaml")) diff --git a/experiments/libero/run_libero_parallel_test.sh b/experiments/libero/run_libero_parallel_test.sh index 2164ffe4..c6ebadb3 100644 --- a/experiments/libero/run_libero_parallel_test.sh +++ b/experiments/libero/run_libero_parallel_test.sh @@ -19,6 +19,12 @@ run_libero_eval() { # Basic configuration ROOT_DIR=${ROOT_DIR:-"$(pwd)"} export ROOT_DIR + printf -v WORKER_PYTHONPATH '%q' "${PYTHONPATH:-}" + WORKER_MODEL_BASE_ENV="" + if [ -n "${DIFFSYNTH_MODEL_BASE_PATH:-}" ]; then + printf -v WORKER_MODEL_BASE_PATH '%q' "$DIFFSYNTH_MODEL_BASE_PATH" + WORKER_MODEL_BASE_ENV="DIFFSYNTH_MODEL_BASE_PATH=$WORKER_MODEL_BASE_PATH" + fi # Generate a unique run_id RUN_ID=${RUN_ID:-"eval_$(date +%Y%m%d_%H%M%S)"} export RUN_ID @@ -324,6 +330,10 @@ run_libero_eval() { local task_id=$2 local gpu_id=$3 local pane_info=$4 + local visible_devices="$gpu_id" + if [ "$gpu_id" != "0" ]; then + visible_devices="$gpu_id,0" + fi local status_file="$TASK_STATUS_DIR/${suite}_task${task_id}.status" local result_file="$OUTPUT_DIR/$suite/gpu${gpu_id}_task${task_id}_results.json" local log_file="$TASK_LOG_DIR/${suite}_task${task_id}_gpu${gpu_id}.log" @@ -337,7 +347,10 @@ run_libero_eval() { tmux send-keys -t $SESSION_NAME:$pane_info "clear" C-m 2>/dev/null tmux send-keys -t $SESSION_NAME:$pane_info "source ~/.bashrc && cd $ROOT_DIR && export EXP_NAME=$EXP_NAME && \ STATUS_FILE='$status_file' LOG_FILE='$log_file' RESULT_FILE='$result_file' && \ - CUDA_VISIBLE_DEVICES=$gpu_id python experiments/libero/eval_libero_single.py \ + env PYTHONPATH=$WORKER_PYTHONPATH $WORKER_MODEL_BASE_ENV \ + CUDA_VISIBLE_DEVICES=$visible_devices MUJOCO_GL=egl PYOPENGL_PLATFORM=egl MUJOCO_EGL_DEVICE_ID=0 \ + TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 \ + python experiments/libero/eval_libero_single.py \ task=$CONFIG ckpt=$CKPT \ EVALUATION.task_suite_name=$suite EVALUATION.task_id=$task_id gpu_id=$gpu_id \ EVALUATION.num_trials=$NUM_TRIALS EVALUATION.output_dir=$OUTPUT_DIR $EXTRA_ARGS > \"\$LOG_FILE\" 2>&1; \ diff --git a/experiments/libero/summarize_libero_plus.py b/experiments/libero/summarize_libero_plus.py new file mode 100644 index 00000000..fc3689c5 --- /dev/null +++ b/experiments/libero/summarize_libero_plus.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 +"""Summarize persistent-worker LIBERO-Plus results in Kairos Table 13 form.""" + +from __future__ import annotations + +import argparse +import csv +import json +from collections import defaultdict +from pathlib import Path + +from libero_plus_eval_utils import ( + CATEGORY_LABELS, + CATEGORY_ORDER, + EXPECTED_TOTAL_TASKS, + SUITE_ORDER, + LiberoPlusTask, + is_complete_result, + result_path, +) + + +def _rate(successes: int, trials: int) -> float | None: + return None if trials <= 0 else successes / trials * 100.0 + + +def _load_manifest(output_dir: Path) -> dict: + path = output_dir / "manifest.json" + if not path.is_file(): + raise FileNotFoundError(f"Evaluation manifest not found: {path}") + with path.open(encoding="utf-8") as handle: + manifest = json.load(handle) + if int(manifest.get("schema_version", -1)) != 1: + raise ValueError(f"Unsupported manifest schema in {path}: {manifest.get('schema_version')}") + return manifest + + +def summarize(output_dir: str | Path, *, require_complete: bool = False) -> dict: + output_dir = Path(output_dir).expanduser().resolve() + manifest = _load_manifest(output_dir) + tasks = [LiberoPlusTask.from_dict(value) for value in manifest["tasks"]] + num_trials = int(manifest["num_trials"]) + + by_category = { + category: {"tasks": 0, "trials": 0, "successes": 0} + for category in CATEGORY_ORDER + } + by_suite = { + suite: {"tasks": 0, "trials": 0, "successes": 0} + for suite in SUITE_ORDER + } + by_difficulty: dict[str, dict[str, int]] = defaultdict( + lambda: {"tasks": 0, "trials": 0, "successes": 0} + ) + task_rows: list[dict] = [] + missing: list[str] = [] + invalid: list[str] = [] + + for task in tasks: + path = result_path(output_dir, task) + if not path.is_file(): + missing.append(task.key) + continue + if not is_complete_result(path, task, num_trials=num_trials): + invalid.append(task.key) + continue + with path.open(encoding="utf-8") as handle: + result = json.load(handle) + successes = int(result["successes"]) + trials = int(result["total_episodes"]) + + for bucket in (by_category[task.category], by_suite[task.suite]): + bucket["tasks"] += 1 + bucket["trials"] += trials + bucket["successes"] += successes + difficulty_key = ( + "unknown" if task.difficulty_level is None else str(task.difficulty_level) + ) + difficulty_bucket = by_difficulty[difficulty_key] + difficulty_bucket["tasks"] += 1 + difficulty_bucket["trials"] += trials + difficulty_bucket["successes"] += successes + + task_rows.append( + { + "suite": task.suite, + "task_id": task.task_id, + "classification_id": task.classification_id, + "name": task.name, + "category": task.category, + "difficulty_level": task.difficulty_level, + "successes": successes, + "trials": trials, + "success_rate": _rate(successes, trials), + "duration_seconds": result.get("duration"), + "gpu_id": result.get("gpu_id"), + } + ) + + completed_tasks = len(task_rows) + total_successes = sum(bucket["successes"] for bucket in by_category.values()) + total_trials = sum(bucket["trials"] for bucket in by_category.values()) + for bucket in by_category.values(): + bucket["success_rate"] = _rate(bucket["successes"], bucket["trials"]) + for bucket in by_suite.values(): + bucket["success_rate"] = _rate(bucket["successes"], bucket["trials"]) + for bucket in by_difficulty.values(): + bucket["success_rate"] = _rate(bucket["successes"], bucket["trials"]) + + category_rates = [ + by_category[category]["success_rate"] + for category in CATEGORY_ORDER + if by_category[category]["success_rate"] is not None + ] + macro_average = ( + None if len(category_rates) != len(CATEGORY_ORDER) else sum(category_rates) / len(category_rates) + ) + paper_average = _rate(total_successes, total_trials) + is_full_table13 = ( + manifest.get("mode") == "full" + and len(tasks) == EXPECTED_TOTAL_TASKS + and num_trials == 1 + and completed_tasks == EXPECTED_TOTAL_TASKS + and not missing + and not invalid + ) + + report = { + "schema_version": 1, + "mode": manifest.get("mode"), + "checkpoint": manifest.get("checkpoint"), + "num_trials": num_trials, + "expected_tasks": len(tasks), + "completed_tasks": completed_tasks, + "missing_tasks": missing, + "invalid_tasks": invalid, + "is_complete": completed_tasks == len(tasks) and not missing and not invalid, + "is_full_table13": is_full_table13, + "table13": { + "category_order": [CATEGORY_LABELS[value] for value in CATEGORY_ORDER], + "per_category": { + CATEGORY_LABELS[category]: by_category[category] + for category in CATEGORY_ORDER + }, + # Kairos Table 13 "Average" is the micro average over all 10,030 variants. + "average": paper_average, + "macro_average_diagnostic": macro_average, + }, + "per_suite": by_suite, + "per_difficulty": dict(sorted(by_difficulty.items())), + "overall": { + "successes": total_successes, + "trials": total_trials, + "success_rate": paper_average, + }, + "tasks": task_rows, + } + + summary_path = output_dir / "table13_summary.json" + temporary_path = summary_path.with_suffix(".json.tmp") + with temporary_path.open("w", encoding="utf-8") as handle: + json.dump(report, handle, ensure_ascii=True, indent=2) + temporary_path.replace(summary_path) + + table_path = output_dir / "table13.csv" + with table_path.open("w", encoding="utf-8", newline="") as handle: + fieldnames = ["Checkpoint", *[CATEGORY_LABELS[value] for value in CATEGORY_ORDER], "Average"] + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + row = {"Checkpoint": manifest.get("checkpoint", "")} + for category in CATEGORY_ORDER: + value = by_category[category]["success_rate"] + row[CATEGORY_LABELS[category]] = "" if value is None else f"{value:.4f}" + row["Average"] = "" if paper_average is None else f"{paper_average:.4f}" + writer.writerow(row) + + tasks_path = output_dir / "task_results.csv" + with tasks_path.open("w", encoding="utf-8", newline="") as handle: + fieldnames = [ + "suite", + "task_id", + "classification_id", + "name", + "category", + "difficulty_level", + "successes", + "trials", + "success_rate", + "duration_seconds", + "gpu_id", + ] + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(task_rows) + + print("\n=== LIBERO-Plus / Kairos Table 13 ===") + header = [CATEGORY_LABELS[value] for value in CATEGORY_ORDER] + ["Average"] + values = [] + for category in CATEGORY_ORDER: + value = by_category[category]["success_rate"] + values.append("N/A" if value is None else f"{value:.2f}") + values.append("N/A" if paper_average is None else f"{paper_average:.2f}") + print("\t".join(header)) + print("\t".join(values)) + print( + f"Completed {completed_tasks}/{len(tasks)} tasks; " + f"missing={len(missing)} invalid={len(invalid)} full_table13={is_full_table13}" + ) + print(f"Saved {summary_path}") + print(f"Saved {table_path}") + print(f"Saved {tasks_path}") + + if require_complete and not report["is_complete"]: + raise RuntimeError( + f"LIBERO-Plus results are incomplete: completed={completed_tasks}/{len(tasks)}, " + f"missing={len(missing)}, invalid={len(invalid)}." + ) + return report + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", required=True) + parser.add_argument("--require-complete", action="store_true") + args = parser.parse_args() + summarize(args.output_dir, require_complete=args.require_complete) + + +if __name__ == "__main__": + main() diff --git a/scripts/collect_fastwam_libero_plus_robot_8gpu.sh b/scripts/collect_fastwam_libero_plus_robot_8gpu.sh new file mode 100755 index 00000000..44ce3c48 --- /dev/null +++ b/scripts/collect_fastwam_libero_plus_robot_8gpu.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python3}" +DEPS_ROOT="${DEPS_ROOT:-}" +LIBERO_PLUS_ROOT="${LIBERO_PLUS_ROOT:-}" +REFERENCE_DATASET="${REFERENCE_DATASET:-}" +MODEL_BASE_PATH="${MODEL_BASE_PATH:-}" +CHECKPOINT="${CHECKPOINT:-}" +CHECKPOINT_LOAD_PATH="${CHECKPOINT_LOAD_PATH:-${CHECKPOINT}}" +DATASET_STATS_PATH="${DATASET_STATS_PATH:-}" +GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}" +WORKERS_PER_GPU="${WORKERS_PER_GPU:-1}" +TEXT_EMBEDDING_CACHE_DIR="${TEXT_EMBEDDING_CACHE_DIR:-}" +EGL_LOCK_SCOPE="${EGL_LOCK_SCOPE:-global}" +EGL_FALLBACK_GPU="${EGL_FALLBACK_GPU:-}" +MODEL_REDIRECT_COMMON_FILES="${MODEL_REDIRECT_COMMON_FILES:-false}" +TOKENIZER_MODEL_ID="${TOKENIZER_MODEL_ID:-Wan-AI/Wan2.2-TI2V-5B}" + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 OUTPUT_DIR [collector arguments ...]" >&2 + echo "Required environment: LIBERO_PLUS_ROOT, REFERENCE_DATASET, MODEL_BASE_PATH," >&2 + echo " CHECKPOINT, and DATASET_STATS_PATH." >&2 + echo "Example: $0 ./outputs/robot_rollouts --task-source heldout --mode full" >&2 + exit 2 +fi + +OUTPUT_DIR="$1" +shift + +required_variables=( + LIBERO_PLUS_ROOT + REFERENCE_DATASET + MODEL_BASE_PATH + CHECKPOINT + DATASET_STATS_PATH +) +for variable_name in "${required_variables[@]}"; do + if [[ -z "${!variable_name}" ]]; then + echo "Required environment variable is not set: ${variable_name}" >&2 + exit 2 + fi +done + +if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then + echo "Python executable not found: ${PYTHON_BIN}" >&2 + exit 1 +fi +if [[ -n "${DEPS_ROOT}" && ! -d "${DEPS_ROOT}" ]]; then + echo "Vendored dependency directory not found: ${DEPS_ROOT}" >&2 + exit 1 +fi + +REPO_PYTHONPATH="${ROOT_DIR}:${ROOT_DIR}/src:${ROOT_DIR}/experiments/libero:${LIBERO_PLUS_ROOT}" +if [[ -n "${DEPS_ROOT}" ]]; then + REPO_PYTHONPATH="${DEPS_ROOT}:${REPO_PYTHONPATH}" +fi +export PYTHONPATH="${REPO_PYTHONPATH}${PYTHONPATH:+:${PYTHONPATH}}" +export DIFFSYNTH_MODEL_BASE_PATH="${MODEL_BASE_PATH}" +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +export PYTHONUNBUFFERED=1 + +mkdir -p "${OUTPUT_DIR}" +RUNTIME_ARGS=( + --workers-per-gpu "${WORKERS_PER_GPU}" + --egl-lock-scope "${EGL_LOCK_SCOPE}" +) +if [[ -n "${EGL_FALLBACK_GPU}" ]]; then + RUNTIME_ARGS+=(--egl-fallback-gpu "${EGL_FALLBACK_GPU}") +fi +if [[ -n "${TEXT_EMBEDDING_CACHE_DIR}" ]]; then + RUNTIME_ARGS+=(--text-embedding-cache-dir "${TEXT_EMBEDDING_CACHE_DIR}") +fi + +exec "${PYTHON_BIN}" "${ROOT_DIR}/experiments/libero/collect_libero_plus_self_rollouts.py" \ + --checkpoint "${CHECKPOINT}" \ + --checkpoint-load-path "${CHECKPOINT_LOAD_PATH}" \ + --dataset-stats "${DATASET_STATS_PATH}" \ + --output-dir "${OUTPUT_DIR}" \ + --reference-dataset "${REFERENCE_DATASET}" \ + --libero-plus-root "${LIBERO_PLUS_ROOT}" \ + --model-base-path "${MODEL_BASE_PATH}" \ + --gpus "${GPU_IDS}" \ + --gripper-action-format signed_open_negative \ + --override "model.redirect_common_files=${MODEL_REDIRECT_COMMON_FILES}" \ + --override "model.tokenizer_model_id=${TOKENIZER_MODEL_ID}" \ + --resume \ + "${RUNTIME_ARGS[@]}" \ + "$@" \ + > >(trap '' INT TERM; exec tee -a "${OUTPUT_DIR}/collect.log") \ + 2>&1 diff --git a/scripts/collect_fastwam_libero_plus_robot_eval_state0_8gpu.sh b/scripts/collect_fastwam_libero_plus_robot_eval_state0_8gpu.sh new file mode 100755 index 00000000..5599e064 --- /dev/null +++ b/scripts/collect_fastwam_libero_plus_robot_eval_state0_8gpu.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +POLICY_SEEDS="${POLICY_SEEDS:-42,43}" +COLLECTION_MODE="${COLLECTION_MODE:-full}" +PROMPT_MODE="${PROMPT_MODE:-benchmark}" + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 OUTPUT_DIR [collector arguments ...]" >&2 + echo "This preset uses every selected Table 13 Robot task's saved state[0]," >&2 + echo "fixes env_seed=42, and varies only policy_seed." >&2 + exit 2 +fi + +OUTPUT_DIR="$1" +shift + +exec "${ROOT_DIR}/scripts/collect_fastwam_libero_plus_robot_8gpu.sh" \ + "${OUTPUT_DIR}" \ + --task-source table13 \ + --mode "${COLLECTION_MODE}" \ + --initialization table13_eval_state0 \ + --init-state-indices 0 \ + --policy-seeds "${POLICY_SEEDS}" \ + --prompt-mode "${PROMPT_MODE}" \ + "$@" diff --git a/scripts/eval_fastwam_libero_plus_full_8gpu.sh b/scripts/eval_fastwam_libero_plus_full_8gpu.sh new file mode 100755 index 00000000..8fcdec51 --- /dev/null +++ b/scripts/eval_fastwam_libero_plus_full_8gpu.sh @@ -0,0 +1,119 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python3}" +DEPS_ROOT="${DEPS_ROOT:-}" +LIBERO_PLUS_ROOT="${LIBERO_PLUS_ROOT:-}" +DATASET_STATS_PATH="${DATASET_STATS_PATH:-}" +MODEL_BASE_PATH="${MODEL_BASE_PATH:-}" +NVIDIA_EGL_ROOT="${NVIDIA_EGL_ROOT:-}" +TEXT_EMBEDDING_CACHE_DIR="${TEXT_EMBEDDING_CACHE_DIR:-}" +EGL_FALLBACK_GPU="${EGL_FALLBACK_GPU:-}" +GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}" +WORKERS_PER_GPU="${WORKERS_PER_GPU:-8}" +EGL_LOCK_SCOPE="${EGL_LOCK_SCOPE:-gpu}" +WORKER_READY_TIMEOUT="${WORKER_READY_TIMEOUT:-1200}" +CHECKPOINT_LOAD_PATH="${CHECKPOINT_LOAD_PATH:-}" +TOKENIZER_MODEL_ID="${TOKENIZER_MODEL_ID:-Wan-AI/Wan2.2-TI2V-5B}" +GRIPPER_ACTION_FORMAT="${GRIPPER_ACTION_FORMAT:-signed_open_negative}" +ENV_SEED="${ENV_SEED:-}" +POLICY_SEED="${POLICY_SEED:-}" +TABLE13_CATEGORIES="${TABLE13_CATEGORIES:-}" + +if [[ $# -lt 1 || $# -gt 2 ]]; then + echo "Usage: $0 /path/to/checkpoint-or-fastwam.pt [output_dir]" >&2 + echo "Required environment: LIBERO_PLUS_ROOT, DATASET_STATS_PATH, MODEL_BASE_PATH." >&2 + exit 2 +fi + +CKPT="$1" +if [[ -d "${CKPT}" ]]; then + CKPT="${CKPT%/}/fastwam.pt" +fi +CKPT_LABEL="$(basename "$(dirname "${CKPT}")")" +STAMP="$(date -u +%Y%m%d_%H%M%S)" +OUTPUT_DIR="${2:-${ROOT_DIR}/evaluate_results/libero_plus/${CKPT_LABEL}_table13_${STAMP}}" + +for variable_name in LIBERO_PLUS_ROOT DATASET_STATS_PATH MODEL_BASE_PATH; do + if [[ -z "${!variable_name}" ]]; then + echo "Required environment variable is not set: ${variable_name}" >&2 + exit 2 + fi +done +if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then + echo "Python executable not found: ${PYTHON_BIN}" >&2 + exit 1 +fi +if [[ -n "${DEPS_ROOT}" && ! -d "${DEPS_ROOT}" ]]; then + echo "Vendored dependency directory not found: ${DEPS_ROOT}" >&2 + exit 1 +fi +if [[ ! "${WORKERS_PER_GPU}" =~ ^[1-9][0-9]*$ ]]; then + echo "WORKERS_PER_GPU must be a positive integer." >&2 + exit 2 +fi +if [[ "${EGL_LOCK_SCOPE}" != "global" && "${EGL_LOCK_SCOPE}" != "gpu" ]]; then + echo "EGL_LOCK_SCOPE must be global or gpu." >&2 + exit 2 +fi +if (( WORKERS_PER_GPU > 1 )) && [[ -z "${TEXT_EMBEDDING_CACHE_DIR}" ]]; then + echo "W${WORKERS_PER_GPU} requires a complete Table 13 TEXT_EMBEDDING_CACHE_DIR." >&2 + exit 2 +fi + +ARGS=( + --checkpoint "${CKPT}" + --dataset-stats "${DATASET_STATS_PATH}" + --output-dir "${OUTPUT_DIR}" + --libero-plus-root "${LIBERO_PLUS_ROOT}" + --model-base-path "${MODEL_BASE_PATH}" + --gpus "${GPU_IDS}" + --workers-per-gpu "${WORKERS_PER_GPU}" + --egl-lock-scope "${EGL_LOCK_SCOPE}" + --worker-ready-timeout "${WORKER_READY_TIMEOUT}" + --mode full + --num-trials 1 + --gripper-action-format "${GRIPPER_ACTION_FORMAT}" + --override model.redirect_common_files=false + --override "model.tokenizer_model_id=${TOKENIZER_MODEL_ID}" + --resume +) +if [[ -n "${CHECKPOINT_LOAD_PATH}" ]]; then + ARGS+=(--checkpoint-load-path "${CHECKPOINT_LOAD_PATH}") +fi +if [[ -n "${TEXT_EMBEDDING_CACHE_DIR}" ]]; then + ARGS+=(--text-embedding-cache-dir "${TEXT_EMBEDDING_CACHE_DIR}") +fi +if [[ -n "${EGL_FALLBACK_GPU}" ]]; then + ARGS+=(--egl-fallback-gpu "${EGL_FALLBACK_GPU}") +fi +if [[ -n "${ENV_SEED}" ]]; then + ARGS+=(--env-seed "${ENV_SEED}") +fi +if [[ -n "${POLICY_SEED}" ]]; then + ARGS+=(--policy-seed "${POLICY_SEED}") +fi +if [[ -n "${TABLE13_CATEGORIES}" ]]; then + ARGS+=(--category "${TABLE13_CATEGORIES}") +fi + +REPO_PYTHONPATH="${ROOT_DIR}:${ROOT_DIR}/src:${ROOT_DIR}/experiments/libero:${LIBERO_PLUS_ROOT}" +if [[ -n "${DEPS_ROOT}" ]]; then + REPO_PYTHONPATH="${DEPS_ROOT}:${REPO_PYTHONPATH}" +fi +export PYTHONPATH="${REPO_PYTHONPATH}${PYTHONPATH:+:${PYTHONPATH}}" +export DIFFSYNTH_MODEL_BASE_PATH="${MODEL_BASE_PATH}" +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +export PYTHONUNBUFFERED=1 +if [[ -n "${NVIDIA_EGL_ROOT}" && -d "${NVIDIA_EGL_ROOT}/usr/lib/x86_64-linux-gnu" ]]; then + export LD_LIBRARY_PATH="${NVIDIA_EGL_ROOT}/usr/lib/x86_64-linux-gnu${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + export __EGL_VENDOR_LIBRARY_FILENAMES="${NVIDIA_EGL_ROOT}/usr/share/glvnd/egl_vendor.d/10_nvidia.json" +fi + +mkdir -p "${OUTPUT_DIR}" +echo "LIBERO-Plus Table 13 output: ${OUTPUT_DIR}" +exec "${PYTHON_BIN}" "${ROOT_DIR}/experiments/libero/eval_libero_plus_persistent.py" "${ARGS[@]}" \ + > >(trap '' INT TERM; exec tee -a "${OUTPUT_DIR}/eval.log") \ + 2>&1 diff --git a/scripts/eval_fastwam_libero_plus_smoke_8gpu.sh b/scripts/eval_fastwam_libero_plus_smoke_8gpu.sh new file mode 100755 index 00000000..6f728148 --- /dev/null +++ b/scripts/eval_fastwam_libero_plus_smoke_8gpu.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python3}" +DEPS_ROOT="${DEPS_ROOT:-}" +LIBERO_PLUS_ROOT="${LIBERO_PLUS_ROOT:-}" +DATASET_STATS_PATH="${DATASET_STATS_PATH:-}" +MODEL_BASE_PATH="${MODEL_BASE_PATH:-}" +NVIDIA_EGL_ROOT="${NVIDIA_EGL_ROOT:-}" +TEXT_EMBEDDING_CACHE_DIR="${TEXT_EMBEDDING_CACHE_DIR:-}" +EGL_FALLBACK_GPU="${EGL_FALLBACK_GPU:-}" +GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}" +WORKERS_PER_GPU="${WORKERS_PER_GPU:-1}" +EGL_LOCK_SCOPE="${EGL_LOCK_SCOPE:-global}" +WORKER_READY_TIMEOUT="${WORKER_READY_TIMEOUT:-1200}" +CHECKPOINT_LOAD_PATH="${CHECKPOINT_LOAD_PATH:-}" +TOKENIZER_MODEL_ID="${TOKENIZER_MODEL_ID:-Wan-AI/Wan2.2-TI2V-5B}" +SMOKE_TASKS="${SMOKE_TASKS:-14}" +SMOKE_SEED="${SMOKE_SEED:-42}" +SMOKE_TRIALS="${SMOKE_TRIALS:-1}" +SMOKE_MAX_TASKS_PER_WORKER="${SMOKE_MAX_TASKS_PER_WORKER:-0}" +SAVE_VIDEOS="${SAVE_VIDEOS:-1}" +GRIPPER_ACTION_FORMAT="${GRIPPER_ACTION_FORMAT:-signed_open_negative}" +ENV_SEED="${ENV_SEED:-}" +POLICY_SEED="${POLICY_SEED:-}" +TABLE13_CATEGORIES="${TABLE13_CATEGORIES:-}" + +if [[ $# -lt 1 || $# -gt 2 ]]; then + echo "Usage: $0 /path/to/checkpoint-or-fastwam.pt [output_dir]" >&2 + echo "Required environment: LIBERO_PLUS_ROOT, DATASET_STATS_PATH, MODEL_BASE_PATH." >&2 + exit 2 +fi + +CKPT="$1" +if [[ -d "${CKPT}" ]]; then + CKPT="${CKPT%/}/fastwam.pt" +fi +CKPT_LABEL="$(basename "$(dirname "${CKPT}")")" +STAMP="$(date -u +%Y%m%d_%H%M%S)" +OUTPUT_DIR="${2:-${ROOT_DIR}/evaluate_results/libero_plus/${CKPT_LABEL}_smoke_${STAMP}}" + +for variable_name in LIBERO_PLUS_ROOT DATASET_STATS_PATH MODEL_BASE_PATH; do + if [[ -z "${!variable_name}" ]]; then + echo "Required environment variable is not set: ${variable_name}" >&2 + exit 2 + fi +done +if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then + echo "Python executable not found: ${PYTHON_BIN}" >&2 + exit 1 +fi +if [[ -n "${DEPS_ROOT}" && ! -d "${DEPS_ROOT}" ]]; then + echo "Vendored dependency directory not found: ${DEPS_ROOT}" >&2 + exit 1 +fi +if [[ ! "${WORKERS_PER_GPU}" =~ ^[1-9][0-9]*$ ]]; then + echo "WORKERS_PER_GPU must be a positive integer." >&2 + exit 2 +fi +if [[ "${EGL_LOCK_SCOPE}" != "global" && "${EGL_LOCK_SCOPE}" != "gpu" ]]; then + echo "EGL_LOCK_SCOPE must be global or gpu." >&2 + exit 2 +fi +if (( WORKERS_PER_GPU > 1 )) && [[ -z "${TEXT_EMBEDDING_CACHE_DIR}" ]]; then + echo "W${WORKERS_PER_GPU} requires a complete selected-task TEXT_EMBEDDING_CACHE_DIR." >&2 + exit 2 +fi + +ARGS=( + --checkpoint "${CKPT}" + --dataset-stats "${DATASET_STATS_PATH}" + --output-dir "${OUTPUT_DIR}" + --libero-plus-root "${LIBERO_PLUS_ROOT}" + --model-base-path "${MODEL_BASE_PATH}" + --gpus "${GPU_IDS}" + --workers-per-gpu "${WORKERS_PER_GPU}" + --egl-lock-scope "${EGL_LOCK_SCOPE}" + --worker-ready-timeout "${WORKER_READY_TIMEOUT}" + --mode smoke + --smoke-tasks "${SMOKE_TASKS}" + --smoke-seed "${SMOKE_SEED}" + --num-trials "${SMOKE_TRIALS}" + --gripper-action-format "${GRIPPER_ACTION_FORMAT}" + --max-tasks-per-worker "${SMOKE_MAX_TASKS_PER_WORKER}" + --override model.redirect_common_files=false + --override "model.tokenizer_model_id=${TOKENIZER_MODEL_ID}" + --resume +) +if [[ -n "${CHECKPOINT_LOAD_PATH}" ]]; then + ARGS+=(--checkpoint-load-path "${CHECKPOINT_LOAD_PATH}") +fi +if [[ "${SAVE_VIDEOS}" == "1" ]]; then + ARGS+=(--save-videos) +fi +if [[ -n "${TEXT_EMBEDDING_CACHE_DIR}" ]]; then + ARGS+=(--text-embedding-cache-dir "${TEXT_EMBEDDING_CACHE_DIR}") +fi +if [[ -n "${EGL_FALLBACK_GPU}" ]]; then + ARGS+=(--egl-fallback-gpu "${EGL_FALLBACK_GPU}") +fi +if [[ -n "${ENV_SEED}" ]]; then + ARGS+=(--env-seed "${ENV_SEED}") +fi +if [[ -n "${POLICY_SEED}" ]]; then + ARGS+=(--policy-seed "${POLICY_SEED}") +fi +if [[ -n "${TABLE13_CATEGORIES}" ]]; then + ARGS+=(--category "${TABLE13_CATEGORIES}") +fi + +REPO_PYTHONPATH="${ROOT_DIR}:${ROOT_DIR}/src:${ROOT_DIR}/experiments/libero:${LIBERO_PLUS_ROOT}" +if [[ -n "${DEPS_ROOT}" ]]; then + REPO_PYTHONPATH="${DEPS_ROOT}:${REPO_PYTHONPATH}" +fi +export PYTHONPATH="${REPO_PYTHONPATH}${PYTHONPATH:+:${PYTHONPATH}}" +export DIFFSYNTH_MODEL_BASE_PATH="${MODEL_BASE_PATH}" +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +export PYTHONUNBUFFERED=1 +if [[ -n "${NVIDIA_EGL_ROOT}" && -d "${NVIDIA_EGL_ROOT}/usr/lib/x86_64-linux-gnu" ]]; then + export LD_LIBRARY_PATH="${NVIDIA_EGL_ROOT}/usr/lib/x86_64-linux-gnu${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + export __EGL_VENDOR_LIBRARY_FILENAMES="${NVIDIA_EGL_ROOT}/usr/share/glvnd/egl_vendor.d/10_nvidia.json" +fi + +mkdir -p "${OUTPUT_DIR}" +echo "LIBERO-Plus smoke output: ${OUTPUT_DIR}" +exec "${PYTHON_BIN}" "${ROOT_DIR}/experiments/libero/eval_libero_plus_persistent.py" "${ARGS[@]}" \ + > >(trap '' INT TERM; exec tee -a "${OUTPUT_DIR}/eval.log") \ + 2>&1 diff --git a/scripts/export_libero_plus_eval_prompts.py b/scripts/export_libero_plus_eval_prompts.py new file mode 100644 index 00000000..750741de --- /dev/null +++ b/scripts/export_libero_plus_eval_prompts.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Export exact LIBERO-Plus benchmark instructions for text-cache generation.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import tempfile +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +EXPERIMENT_DIR = PROJECT_ROOT / "experiments" / "libero" +for path in (PROJECT_ROOT, PROJECT_ROOT / "src", EXPERIMENT_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from eval_libero_plus_persistent import ( # noqa: E402 + _ensure_libero_config, + _load_exact_task_descriptions, +) +from libero_plus_eval_utils import ( # noqa: E402 + default_classification_path, + load_task_classification, +) + + +def _atomic_json_dump(payload: dict, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + with temporary_path.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, ensure_ascii=True, indent=2) + os.replace(temporary_path, path) + + +def _atomic_jsonl_dump(records: list[dict], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + with temporary_path.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=True) + "\n") + os.replace(temporary_path, path) + + +def run(args: argparse.Namespace) -> int: + libero_plus_root = Path(args.libero_plus_root).expanduser().resolve() + classification_path = ( + Path(args.classification).expanduser().resolve() + if args.classification + else default_classification_path(libero_plus_root) + ) + output_dir = Path(args.output_dir).expanduser().resolve() + tasks = load_task_classification(classification_path, require_full=True) + + previous_config_path = os.environ.get("LIBERO_CONFIG_PATH") + try: + with tempfile.TemporaryDirectory( + prefix="fastwam-libero-plus-prompts-" + ) as temporary_dir: + config_dir = _ensure_libero_config( + libero_plus_root, + Path(temporary_dir), + ) + os.environ["LIBERO_CONFIG_PATH"] = str(config_dir) + descriptions = _load_exact_task_descriptions(tasks) + finally: + if previous_config_path is None: + os.environ.pop("LIBERO_CONFIG_PATH", None) + else: + os.environ["LIBERO_CONFIG_PATH"] = previous_config_path + + unique_descriptions = list(dict.fromkeys(descriptions)) + records = [ + { + "task_index": task_index, + "task": description, + } + for task_index, description in enumerate(unique_descriptions) + ] + meta_dir = output_dir / "meta" + _atomic_jsonl_dump(records, meta_dir / "tasks.jsonl") + _atomic_json_dump( + { + "schema_version": 1, + "source": "LIBERO-Plus Table 13 exact task.language", + "classification_path": str(classification_path), + "total_tasks": len(descriptions), + "unique_prompts": len(unique_descriptions), + }, + meta_dir / "prompt_inventory.json", + ) + print( + "Exported LIBERO-Plus evaluation prompts: " + f"tasks={len(descriptions)} unique={len(unique_descriptions)} " + f"output={output_dir}", + flush=True, + ) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--libero-plus-root", required=True) + parser.add_argument("--classification", default=None) + parser.add_argument("--output-dir", required=True) + return parser + + +def main() -> None: + raise SystemExit(run(build_parser().parse_args())) + + +if __name__ == "__main__": + main() diff --git a/scripts/precompute_text_embeds.py b/scripts/precompute_text_embeds.py index 0c1f8e2d..8632569a 100644 --- a/scripts/precompute_text_embeds.py +++ b/scripts/precompute_text_embeds.py @@ -181,6 +181,14 @@ def main(cfg: DictConfig): ) overwrite = _to_bool(cfg.get("overwrite", True)) + text_embedding_batch_size = int( + cfg.get("text_embedding_batch_size", DEFAULT_BATCH_SIZE) + ) + if text_embedding_batch_size < 1: + raise ValueError( + "text_embedding_batch_size must be positive, got " + f"{text_embedding_batch_size}." + ) model_cfg = cfg.model if model_cfg is None: raise ValueError("`cfg.model` is required.") @@ -215,12 +223,14 @@ def main(cfg: DictConfig): enc_id = _model_id_to_enc_id(model_id) logger.info( - "Preparing text encoder with model_id=%s tokenizer_model_id=%s device=%s dtype=%s context_len=%d overwrite=%s", + "Preparing text encoder with model_id=%s tokenizer_model_id=%s device=%s " + "dtype=%s context_len=%d batch_size=%d overwrite=%s", model_id, tokenizer_model_id, device, torch_dtype, context_len, + text_embedding_batch_size, overwrite, ) @@ -306,8 +316,8 @@ def main(cfg: DictConfig): disable=is_distributed and rank != 0, ) as pbar: with torch.no_grad(): - for start in range(0, len(prompts), DEFAULT_BATCH_SIZE): - batch_prompts = prompts[start : start + DEFAULT_BATCH_SIZE] + for start in range(0, len(prompts), text_embedding_batch_size): + batch_prompts = prompts[start : start + text_embedding_batch_size] ids, mask = tokenizer(batch_prompts, return_mask=True, add_special_tokens=True) ids = ids.to(device) mask = mask.to(device=device, dtype=torch.bool) diff --git a/scripts/retry_fastwam_libero_plus_robot_eval_state0_8gpu.sh b/scripts/retry_fastwam_libero_plus_robot_eval_state0_8gpu.sh new file mode 100755 index 00000000..ce45e79a --- /dev/null +++ b/scripts/retry_fastwam_libero_plus_robot_eval_state0_8gpu.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +export RETRY_CAMPAIGN_NAME="${RETRY_CAMPAIGN_NAME:-retry_to_2_successes_eval_state0_v1}" +export RETRY_INDEX_START="${RETRY_INDEX_START:-44}" +export INITIALIZATION="table13_eval_state0" + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 EXISTING_COLLECTION_OUTPUT [retry arguments ...]" >&2 + echo "This preset keeps env_seed=42 and advances only policy_seed until each" >&2 + echo "exact Table 13 Robot variant has two cumulative successes or hits a limit." >&2 + exit 2 +fi + +OUTPUT_DIR="$1" +shift + +if [[ -z "${PROMPT_MODE:-}" ]]; then + BASE_MANIFEST="${OUTPUT_DIR%/}/collection_manifest.json" + if [[ ! -f "${BASE_MANIFEST}" ]]; then + echo "Base collection manifest not found: ${BASE_MANIFEST}" >&2 + exit 2 + fi + PROMPT_MODE="$("${PYTHON_BIN:-python3}" - "${BASE_MANIFEST}" <<'PY' +import json +import sys + +with open(sys.argv[1], encoding="utf-8") as handle: + prompt_mode = json.load(handle).get("prompt_mode") +if prompt_mode not in {"canonical", "benchmark"}: + raise SystemExit(f"Unsupported base collection prompt_mode: {prompt_mode!r}") +print(prompt_mode) +PY +)" +fi +export PROMPT_MODE + +exec "${ROOT_DIR}/scripts/retry_fastwam_libero_plus_robot_failed_8gpu.sh" \ + "${OUTPUT_DIR}" \ + --init-state-indices 0 \ + --target-successes-per-variant 2 \ + "$@" diff --git a/scripts/retry_fastwam_libero_plus_robot_failed_8gpu.sh b/scripts/retry_fastwam_libero_plus_robot_failed_8gpu.sh new file mode 100755 index 00000000..4f5e9341 --- /dev/null +++ b/scripts/retry_fastwam_libero_plus_robot_failed_8gpu.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python3}" +DEPS_ROOT="${DEPS_ROOT:-}" +LIBERO_PLUS_ROOT="${LIBERO_PLUS_ROOT:-}" +REFERENCE_DATASET="${REFERENCE_DATASET:-}" +MODEL_BASE_PATH="${MODEL_BASE_PATH:-}" +CHECKPOINT="${CHECKPOINT:-}" +CHECKPOINT_LOAD_PATH="${CHECKPOINT_LOAD_PATH:-${CHECKPOINT}}" +DATASET_STATS_PATH="${DATASET_STATS_PATH:-}" +GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}" +WORKERS_PER_GPU="${WORKERS_PER_GPU:-1}" +TEXT_EMBEDDING_CACHE_DIR="${TEXT_EMBEDDING_CACHE_DIR:-}" +EGL_LOCK_SCOPE="${EGL_LOCK_SCOPE:-global}" +EGL_FALLBACK_GPU="${EGL_FALLBACK_GPU:-}" +MODEL_REDIRECT_COMMON_FILES="${MODEL_REDIRECT_COMMON_FILES:-false}" +TOKENIZER_MODEL_ID="${TOKENIZER_MODEL_ID:-Wan-AI/Wan2.2-TI2V-5B}" +RETRY_CAMPAIGN_NAME="${RETRY_CAMPAIGN_NAME:-retry_zero_success_robot_v1}" +RETRY_INDEX_START="${RETRY_INDEX_START:-44}" +RETRY_MAX_ATTEMPTS="${RETRY_MAX_ATTEMPTS:-96}" +RETRY_MAX_RUNTIME_HOURS="${RETRY_MAX_RUNTIME_HOURS:-12}" +INITIALIZATION="${INITIALIZATION:-random_reset}" +PROMPT_MODE="${PROMPT_MODE:-canonical}" + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 EXISTING_COLLECTION_OUTPUT [retry arguments ...]" >&2 + echo "Required environment: LIBERO_PLUS_ROOT, REFERENCE_DATASET, MODEL_BASE_PATH," >&2 + echo " CHECKPOINT, and DATASET_STATS_PATH." >&2 + exit 2 +fi + +OUTPUT_DIR="$1" +shift + +required_variables=( + LIBERO_PLUS_ROOT + REFERENCE_DATASET + MODEL_BASE_PATH + CHECKPOINT + DATASET_STATS_PATH +) +for variable_name in "${required_variables[@]}"; do + if [[ -z "${!variable_name}" ]]; then + echo "Required environment variable is not set: ${variable_name}" >&2 + exit 2 + fi +done + +if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then + echo "Python executable not found: ${PYTHON_BIN}" >&2 + exit 1 +fi +if [[ -n "${DEPS_ROOT}" && ! -d "${DEPS_ROOT}" ]]; then + echo "Vendored dependency directory not found: ${DEPS_ROOT}" >&2 + exit 1 +fi + +REPO_PYTHONPATH="${ROOT_DIR}:${ROOT_DIR}/src:${ROOT_DIR}/experiments/libero:${LIBERO_PLUS_ROOT}" +if [[ -n "${DEPS_ROOT}" ]]; then + REPO_PYTHONPATH="${DEPS_ROOT}:${REPO_PYTHONPATH}" +fi +export PYTHONPATH="${REPO_PYTHONPATH}${PYTHONPATH:+:${PYTHONPATH}}" +export DIFFSYNTH_MODEL_BASE_PATH="${MODEL_BASE_PATH}" +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +export PYTHONUNBUFFERED=1 + +if [[ + "${RETRY_CAMPAIGN_NAME}" == "." || + "${RETRY_CAMPAIGN_NAME}" == ".." || + ! "${RETRY_CAMPAIGN_NAME}" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]*$ +]]; then + echo "Invalid RETRY_CAMPAIGN_NAME: ${RETRY_CAMPAIGN_NAME}" >&2 + exit 2 +fi + +CAMPAIGN_DIR="${OUTPUT_DIR%/}/retry_campaigns/${RETRY_CAMPAIGN_NAME}" +mkdir -p "${CAMPAIGN_DIR}" +RUNTIME_ARGS=( + --workers-per-gpu "${WORKERS_PER_GPU}" + --egl-lock-scope "${EGL_LOCK_SCOPE}" +) +if [[ -n "${EGL_FALLBACK_GPU}" ]]; then + RUNTIME_ARGS+=(--egl-fallback-gpu "${EGL_FALLBACK_GPU}") +fi +if [[ -n "${TEXT_EMBEDDING_CACHE_DIR}" ]]; then + RUNTIME_ARGS+=(--text-embedding-cache-dir "${TEXT_EMBEDDING_CACHE_DIR}") +fi + +exec "${PYTHON_BIN}" "${ROOT_DIR}/experiments/libero/retry_libero_plus_failed_rollouts.py" \ + --checkpoint "${CHECKPOINT}" \ + --checkpoint-load-path "${CHECKPOINT_LOAD_PATH}" \ + --dataset-stats "${DATASET_STATS_PATH}" \ + --output-dir "${OUTPUT_DIR}" \ + --reference-dataset "${REFERENCE_DATASET}" \ + --libero-plus-root "${LIBERO_PLUS_ROOT}" \ + --model-base-path "${MODEL_BASE_PATH}" \ + --gpus "${GPU_IDS}" \ + --gripper-action-format signed_open_negative \ + --override "model.redirect_common_files=${MODEL_REDIRECT_COMMON_FILES}" \ + --override "model.tokenizer_model_id=${TOKENIZER_MODEL_ID}" \ + --task-source table13 \ + --mode full \ + --initialization "${INITIALIZATION}" \ + --prompt-mode "${PROMPT_MODE}" \ + --campaign-name "${RETRY_CAMPAIGN_NAME}" \ + --retry-index-start "${RETRY_INDEX_START}" \ + --retry-max-attempts "${RETRY_MAX_ATTEMPTS}" \ + --max-runtime-hours "${RETRY_MAX_RUNTIME_HOURS}" \ + --resume \ + "${RUNTIME_ARGS[@]}" \ + "$@" \ + > >(trap '' INT TERM; exec tee -a "${CAMPAIGN_DIR}/collect.log") \ + 2>&1 diff --git a/src/fastwam/models/wan22/fastwam.py b/src/fastwam/models/wan22/fastwam.py index 106beb3d..1eb36912 100644 --- a/src/fastwam/models/wan22/fastwam.py +++ b/src/fastwam/models/wan22/fastwam.py @@ -1098,9 +1098,12 @@ def save_checkpoint(self, path, optimizer=None, step=None): torch.save(payload, path) def load_checkpoint(self, path, optimizer=None): - payload = torch.load(path, map_location="cpu") + payload = torch.load(path, map_location="cpu", weights_only=True) if "mot" in payload: - self.mot.load_state_dict(payload["mot"], strict=False) + # Evaluation must never silently continue with a partially loaded + # MoT: a missing action/video expert key would make the reported + # LIBERO score invalid while still allowing rollouts to run. + self.mot.load_state_dict(payload["mot"], strict=True) elif "dit" in payload: logger.warning("Loading legacy `dit` checkpoint into video expert only.") self.video_expert.load_state_dict(payload["dit"], strict=False) diff --git a/tests/test_collect_libero_plus_self_rollouts.py b/tests/test_collect_libero_plus_self_rollouts.py new file mode 100644 index 00000000..54f020c9 --- /dev/null +++ b/tests/test_collect_libero_plus_self_rollouts.py @@ -0,0 +1,895 @@ +from __future__ import annotations + +import hashlib +import os +import shlex +import subprocess +import sys +import tempfile +import unittest +from argparse import Namespace +from dataclasses import replace +from pathlib import Path +from unittest import mock + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +EXPERIMENT_DIR = PROJECT_ROOT / "experiments" / "libero" +for path in (PROJECT_ROOT, PROJECT_ROOT / "src", EXPERIMENT_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +import collect_libero_plus_self_rollouts as collector # noqa: E402 +import eval_libero_plus_persistent as persistent_eval # noqa: E402 +import retry_libero_plus_failed_rollouts as retry # noqa: E402 + + +def _task(task_id: int = 7) -> collector.CollectionTask: + return collector.CollectionTask( + suite="libero_spatial", + task_id=task_id, + classification_id=task_id, + classification_name=f"example_task_initstate_{task_id}", + category=collector.ROBOT_CATEGORY, + difficulty_level=4, + base_name="example_task", + canonical_task="pick up the example object", + robot_init_id=task_id, + task_source="table13", + ) + + +class EvalState0CollectionTest(unittest.TestCase): + def test_policy_seed_sequence_fixes_environment_seed(self) -> None: + args = collector.build_parser().parse_args( + [ + "--checkpoint", + "checkpoint.pt", + "--dataset-stats", + "stats.json", + "--output-dir", + "output", + "--reference-dataset", + "reference", + "--model-base-path", + "models", + "--initialization", + collector.TABLE13_EVAL_STATE0_INITIALIZATION, + "--policy-seeds", + "42,50,123", + ] + ) + self.assertEqual( + collector._seed_pairs(args), + [(42, 42), (42, 50), (42, 123)], + ) + + jobs = collector._build_jobs( + [_task()], + seed_pairs=collector._seed_pairs(args), + initialization=args.initialization, + init_state_indices=[99], + prompt_mode="canonical", + ) + self.assertEqual({job.env_seed for job in jobs}, {42}) + self.assertEqual([job.policy_seed for job in jobs], [42, 50, 123]) + self.assertEqual({job.init_state_index for job in jobs}, {0}) + + manifest = collector._build_manifest( + args, + Path("classification.json"), + [0], + [_task()], + jobs, + {}, + ) + self.assertEqual( + manifest["initial_state_strategy"], + { + "mode": "official_saved_state", + "expression": ( + "task_suite.get_task_init_states(task_id)[0]" + ), + "index": 0, + }, + ) + self.assertEqual( + manifest["seed_strategy"]["environment"], + {"mode": "fixed", "seed": 42}, + ) + self.assertEqual( + manifest["seed_strategy"]["policy"]["seeds"], + [42, 50, 123], + ) + + def test_explicit_non_42_environment_seed_is_rejected(self) -> None: + args = collector.build_parser().parse_args( + [ + "--checkpoint", + "checkpoint.pt", + "--dataset-stats", + "stats.json", + "--output-dir", + "output", + "--reference-dataset", + "reference", + "--model-base-path", + "models", + "--initialization", + collector.TABLE13_EVAL_STATE0_INITIALIZATION, + "--seed-pair", + "43:50", + ] + ) + with self.assertRaisesRegex(ValueError, "fixes env_seed=42"): + collector._seed_pairs(args) + + def test_legacy_random_reset_seed_pair_behavior_is_unchanged(self) -> None: + args = collector.build_parser().parse_args( + [ + "--checkpoint", + "checkpoint.pt", + "--dataset-stats", + "stats.json", + "--output-dir", + "output", + "--reference-dataset", + "reference", + "--model-base-path", + "models", + "--rollout-seeds", + "42,43", + ] + ) + self.assertEqual(collector._seed_pairs(args), [(42, 42), (43, 43)]) + + +class CollectionPromptCachePreflightTest(unittest.TestCase): + @staticmethod + def _cache_path(cache_dir: Path, description: str) -> Path: + full_prompt = persistent_eval.WAN_PROMPT_TEMPLATE.format( + task=description + ) + digest = hashlib.sha256(full_prompt.encode("utf-8")).hexdigest() + return cache_dir / ( + f"{digest}.t5_len{persistent_eval.WAN_PROMPT_CONTEXT_LEN}." + f"{persistent_eval.WAN_PROMPT_ENCODER_ID}.pt" + ) + + @staticmethod + def _args( + cache_dir: Path, + output_dir: Path, + *, + prompt_mode: str = "canonical", + ): + return collector.build_parser().parse_args( + [ + "--checkpoint", + "checkpoint.pt", + "--dataset-stats", + "stats.json", + "--output-dir", + str(output_dir), + "--reference-dataset", + "reference", + "--model-base-path", + "models", + "--text-embedding-cache-dir", + str(cache_dir), + "--prompt-mode", + prompt_mode, + ] + ) + + @staticmethod + def _preflight(args, tasks, output_dir: Path): + cache_cfg = Namespace(model={"tokenizer_max_len": 128}) + with mock.patch.object( + persistent_eval, + "_compose_worker_config", + return_value=cache_cfg, + ): + return collector._preflight_text_embedding_cache( + args, + tasks, + output_dir=output_dir, + physical_gpu_id=0, + ) + + def test_preflight_accepts_all_selected_canonical_prompts(self) -> None: + tasks = [ + _task(1), + replace(_task(2), canonical_task="place the example in the bowl"), + _task(3), + ] + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + cache_dir = root / "cache" + cache_dir.mkdir() + for description in {task.canonical_task for task in tasks}: + self._cache_path(cache_dir, description).touch() + report = self._preflight( + self._args(cache_dir, root / "output"), + tasks, + root / "output", + ) + self.assertEqual(report, {"tasks": 3, "unique_prompts": 2}) + + def test_preflight_reports_missing_count_before_workers(self) -> None: + tasks = [ + _task(1), + replace(_task(2), canonical_task="place the example in the bowl"), + ] + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + cache_dir = root / "cache" + cache_dir.mkdir() + with self.assertRaisesRegex( + FileNotFoundError, + r"missing=2/2 first=.*t5_len128\.wan22ti2v5b\.pt", + ): + self._preflight( + self._args(cache_dir, root / "output"), + tasks, + root / "output", + ) + + def test_benchmark_prompt_cache_is_not_a_canonical_cache(self) -> None: + task = _task(1) + benchmark_description = ( + f"{task.canonical_task} view 0 0 100 0 0 " + f"initstate {task.robot_init_id}" + ) + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + cache_dir = root / "cache" + cache_dir.mkdir() + self._cache_path(cache_dir, benchmark_description).touch() + with self.assertRaisesRegex(FileNotFoundError, r"missing=1/1"): + self._preflight( + self._args(cache_dir, root / "output"), + [task], + root / "output", + ) + + def test_preflight_accepts_exact_official_benchmark_prompt(self) -> None: + task = _task(1) + official_description = ( + f"{task.canonical_task} view 0 0 100 0 0 " + f"initstate {task.robot_init_id}" + ) + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + cache_dir = root / "cache" + cache_dir.mkdir() + self._cache_path(cache_dir, official_description).touch() + with mock.patch.object( + persistent_eval, + "_load_exact_task_descriptions", + return_value=[official_description], + ): + report = self._preflight( + self._args( + cache_dir, + root / "output", + prompt_mode="benchmark", + ), + [task], + root / "output", + ) + self.assertEqual(report, {"tasks": 1, "unique_prompts": 1}) + + def test_benchmark_mode_rejects_canonical_only_cache(self) -> None: + task = _task(1) + official_description = ( + f"{task.canonical_task} view 0 0 100 0 0 " + f"initstate {task.robot_init_id}" + ) + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + cache_dir = root / "cache" + cache_dir.mkdir() + self._cache_path(cache_dir, task.canonical_task).touch() + with mock.patch.object( + persistent_eval, + "_load_exact_task_descriptions", + return_value=[official_description], + ): + with self.assertRaisesRegex(FileNotFoundError, r"missing=1/1"): + self._preflight( + self._args( + cache_dir, + root / "output", + prompt_mode="benchmark", + ), + [task], + root / "output", + ) + + def test_benchmark_descriptions_come_from_official_task_language(self) -> None: + task = _task(1) + official_description = "official benchmark task language initstate 1" + with mock.patch.object( + persistent_eval, + "_load_exact_task_descriptions", + return_value=[official_description], + ) as loader: + descriptions = collector._model_task_descriptions( + [task], + prompt_mode="benchmark", + ) + self.assertEqual(descriptions, [official_description]) + loaded_task = loader.call_args.args[0][0] + self.assertEqual(loaded_task.name, task.classification_name) + self.assertNotEqual(descriptions, [task.classification_name]) + + +class CollectionPromptRuntimeTest(unittest.TestCase): + def test_large_prompt_set_uses_lazy_device_loading(self) -> None: + evaluator = mock.Mock() + descriptions = [ + f"benchmark prompt {index}" + for index in range(collector.MAX_EAGER_PROMPT_CONTEXTS + 1) + ] + loaded = collector._prepare_worker_prompt_contexts( + evaluator, + descriptions, + worker_slot="gpu0-w0", + ) + self.assertEqual(loaded, 0) + evaluator.preload_prompt_contexts.assert_not_called() + + def test_canonical_sized_prompt_set_is_preloaded(self) -> None: + evaluator = mock.Mock() + evaluator.preload_prompt_contexts.return_value = 40 + descriptions = [ + f"canonical prompt {index}" + for index in range(40) + ] + loaded = collector._prepare_worker_prompt_contexts( + evaluator, + descriptions, + worker_slot="gpu0-w0", + ) + self.assertEqual(loaded, 40) + evaluator.preload_prompt_contexts.assert_called_once_with(descriptions) + + def test_runtime_config_is_installed_before_official_prompt_lookup(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + output_dir = root / "output" + with mock.patch.dict(os.environ, {}, clear=False): + config_dir = collector._prepare_libero_runtime( + root / "LIBERO-plus", + output_dir, + ) + self.assertEqual( + os.environ["LIBERO_CONFIG_PATH"], + str(config_dir), + ) + self.assertTrue((config_dir / "config.yaml").is_file()) + + +class CollectionPromptProvenanceTest(unittest.TestCase): + @staticmethod + def _job(prompt_mode: str) -> collector.CollectionJob: + return collector.CollectionJob( + task=_task(1), + env_seed=42, + policy_seed=42, + initialization=collector.TABLE13_EVAL_STATE0_INITIALIZATION, + init_state_index=0, + prompt_mode=prompt_mode, + ) + + def test_benchmark_dataset_label_keeps_exact_model_prompt(self) -> None: + job = self._job("benchmark") + benchmark_prompt = ( + f"{job.task.canonical_task} view 0 0 100 0 0 initstate 1" + ) + provenance = collector._build_success_provenance( + { + "model_prompt": benchmark_prompt, + "benchmark_prompt": benchmark_prompt, + }, + job, + ) + self.assertEqual(provenance["task"], benchmark_prompt) + self.assertEqual(provenance["dataset_task"], benchmark_prompt) + self.assertEqual( + provenance["canonical_task"], + job.task.canonical_task, + ) + self.assertEqual(provenance["benchmark_prompt"], benchmark_prompt) + + def test_canonical_dataset_label_keeps_canonical_model_prompt(self) -> None: + job = self._job("canonical") + provenance = collector._build_success_provenance( + { + "model_prompt": job.task.canonical_task, + "benchmark_prompt": "official benchmark prompt initstate 1", + }, + job, + ) + self.assertEqual(provenance["task"], job.task.canonical_task) + self.assertEqual(provenance["dataset_task"], job.task.canonical_task) + self.assertEqual( + provenance["benchmark_prompt"], + "official benchmark prompt initstate 1", + ) + + def test_benchmark_runtime_prompt_must_match_environment_language(self) -> None: + job = self._job("benchmark") + collector._validate_benchmark_prompt_alignment( + job, + model_prompt="official prompt", + benchmark_prompt="official prompt", + ) + with self.assertRaisesRegex(ValueError, "Benchmark prompt mismatch"): + collector._validate_benchmark_prompt_alignment( + job, + model_prompt="classification_name", + benchmark_prompt="official prompt", + ) + + def test_canonical_runtime_prompt_may_differ_from_benchmark_language(self) -> None: + collector._validate_benchmark_prompt_alignment( + self._job("canonical"), + model_prompt="canonical prompt", + benchmark_prompt="official prompt with suffix", + ) + + +class EvalState0RetryTest(unittest.TestCase): + @staticmethod + def _write_attempt( + output_dir: Path, + job: collector.CollectionJob, + *, + success: bool, + ) -> None: + artifact_dir = output_dir / "staging" / "success" / job.job_id + if success: + artifact_dir.mkdir(parents=True) + (artifact_dir / "manifest.json").write_text("{}\n", encoding="utf-8") + (artifact_dir / "provenance.json").write_text("{}\n", encoding="utf-8") + collector._atomic_json_dump( + { + "job_id": job.job_id, + "job": job.to_dict(), + "success": success, + "artifact_path": str(artifact_dir) if success else None, + }, + collector._result_path(output_dir, job), + ) + + def test_retry_advances_only_policy_seed(self) -> None: + job = retry._retry_job( + _task(), + retry_index=3, + retry_index_start=44, + seed_salt="unused-in-eval-state0", + initialization=collector.TABLE13_EVAL_STATE0_INITIALIZATION, + init_state_index=retry._retry_init_state_index( + collector.TABLE13_EVAL_STATE0_INITIALIZATION + ), + prompt_mode="canonical", + ) + self.assertEqual(job.env_seed, 42) + self.assertEqual(job.policy_seed, 47) + self.assertEqual(job.init_state_index, 0) + + def test_target_two_continues_after_first_retry_success(self) -> None: + state = { + "base_successes": 0, + "complete_indices": {0}, + "successful_indices": [0], + "successful_index": 0, + "target_successes": 2, + "target_reached": False, + "next_index": None, + "exhausted": False, + } + retry._refresh_task_state(state, retry_max_attempts=4) + self.assertFalse(state["target_reached"]) + self.assertEqual(state["next_index"], 1) + + state["complete_indices"].add(1) + state["successful_indices"].append(1) + retry._refresh_task_state(state, retry_max_attempts=4) + self.assertTrue(state["target_reached"]) + self.assertIsNone(state["next_index"]) + self.assertFalse(state["exhausted"]) + + def test_base_one_plus_retry_one_reaches_target_two(self) -> None: + state = { + "base_successes": 1, + "complete_indices": {0}, + "successful_indices": [0], + "successful_index": 0, + "target_successes": 2, + "target_reached": False, + "next_index": None, + "exhausted": False, + } + retry._refresh_task_state(state, retry_max_attempts=4) + self.assertTrue(state["target_reached"]) + self.assertIsNone(state["next_index"]) + + def test_under_target_selection_includes_zero_and_one(self) -> None: + tasks = [_task(1), _task(2), _task(3)] + counts = { + retry._task_identity(tasks[0]): 0, + retry._task_identity(tasks[1]): 1, + retry._task_identity(tasks[2]): 2, + } + selected = retry._base_below_target_tasks( + tasks, + counts, + target_successes_per_variant=2, + ) + self.assertEqual(selected, tasks[:2]) + + def test_new_campaign_manifest_records_target_and_seed_protocol(self) -> None: + args = Namespace( + target_successes_per_variant=2, + text_embedding_cache_dir=None, + task_config="libero_uncond_2cam224_1e-4", + prompt_mode="canonical", + initialization=collector.TABLE13_EVAL_STATE0_INITIALIZATION, + gripper_action_format="signed_open_negative", + override=[], + retry_index_start=44, + retry_max_attempts=96, + retry_seed_salt="unused", + max_runtime_hours=12, + ) + task = _task() + identity = retry._task_identity(task) + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + base_manifest = output_dir / "collection_manifest.json" + base_manifest.write_text("{}\n", encoding="utf-8") + manifest = retry._campaign_manifest( + args, + output_dir=output_dir, + base_manifest_path=base_manifest, + input_files={}, + tasks=[task], + base_success_counts={identity: 1}, + ) + self.assertEqual(manifest["target_successes_per_variant"], 2) + self.assertEqual(manifest["base_successes_by_task"], {identity: 1}) + self.assertEqual(manifest["init_state_index"], 0) + self.assertEqual(manifest["fixed_env_seed"], 42) + self.assertIn( + "policy_seed=retry_index_start+retry_index", + manifest["seed_derivation"], + ) + + def test_legacy_target_one_manifest_shape_is_preserved(self) -> None: + args = Namespace( + target_successes_per_variant=1, + text_embedding_cache_dir=None, + task_config="libero_uncond_2cam224_1e-4", + prompt_mode="canonical", + initialization=collector.RANDOM_RESET_INITIALIZATION, + gripper_action_format="signed_open_negative", + override=[], + retry_index_start=44, + retry_max_attempts=96, + retry_seed_salt="legacy-salt", + max_runtime_hours=12, + ) + task = _task() + identity = retry._task_identity(task) + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + base_manifest = output_dir / "collection_manifest.json" + base_manifest.write_text("{}\n", encoding="utf-8") + manifest = retry._campaign_manifest( + args, + output_dir=output_dir, + base_manifest_path=base_manifest, + input_files={}, + tasks=[task], + base_success_counts={identity: 0}, + ) + self.assertEqual( + manifest["policy"], + "base_zero_success_until_first_retry_success", + ) + self.assertEqual(manifest["init_state_index"], -1) + self.assertEqual(manifest["retry_seed_salt"], "legacy-salt") + self.assertEqual( + manifest["seed_derivation"], + "sha256-task-retry-channel-mod-2^31-1", + ) + self.assertNotIn("target_successes_per_variant", manifest) + self.assertNotIn("base_successes_by_task", manifest) + + def test_retry_prompt_mode_must_exactly_match_base_manifest(self) -> None: + input_files = {"checkpoint": {"size": 1, "sha256": "example"}} + args = Namespace( + initialization=collector.TABLE13_EVAL_STATE0_INITIALIZATION, + prompt_mode="benchmark", + task_config="libero_uncond_2cam224_1e-4", + gripper_action_format="signed_open_negative", + override=[], + ) + base_manifest = { + "schema_version": 2, + "task_source": "table13", + "initialization": args.initialization, + "prompt_mode": "benchmark", + "task_config": args.task_config, + "gripper_action_format": args.gripper_action_format, + "hydra_overrides": [], + "input_files": input_files, + } + retry._validate_base_semantics( + base_manifest, + args, + classification_path=Path("classification.json"), + input_files=input_files, + ) + + args.prompt_mode = "canonical" + with self.assertRaisesRegex(ValueError, "prompt_mode"): + retry._validate_base_semantics( + base_manifest, + args, + classification_path=Path("classification.json"), + input_files=input_files, + ) + + def test_global_success_snapshot_counts_all_prior_campaigns(self) -> None: + tasks = [_task(1), _task(2)] + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + for policy_seed, success in ((42, True), (44, True), (45, False)): + job = collector.CollectionJob( + task=tasks[0], + env_seed=42, + policy_seed=policy_seed, + initialization=collector.TABLE13_EVAL_STATE0_INITIALIZATION, + init_state_index=0, + prompt_mode="benchmark", + ) + self._write_attempt(output_dir, job, success=success) + counts = retry._global_success_counts(output_dir, tasks) + self.assertEqual( + counts, + { + retry._task_identity(tasks[0]): 2, + retry._task_identity(tasks[1]): 0, + }, + ) + + def test_global_additional_manifest_freezes_selection_snapshot(self) -> None: + args = Namespace( + target_successes_per_variant=1, + select_global_successes_below=1, + target_additional_successes_per_variant=1, + wall_deadline="2026-07-31T18:00:00+08:00", + text_embedding_cache_dir=None, + task_config="libero_uncond_2cam224_1e-4", + prompt_mode="benchmark", + initialization=collector.TABLE13_EVAL_STATE0_INITIALIZATION, + gripper_action_format="signed_open_negative", + override=[], + retry_index_start=140, + retry_max_attempts=96, + retry_seed_salt="unused", + max_runtime_hours=12, + ) + tasks = [_task(1), _task(2)] + snapshot = { + retry._task_identity(tasks[0]): 0, + retry._task_identity(tasks[1]): 3, + } + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + base_manifest = output_dir / "collection_manifest.json" + base_manifest.write_text("{}\n", encoding="utf-8") + manifest = retry._campaign_manifest( + args, + output_dir=output_dir, + base_manifest_path=base_manifest, + input_files={}, + tasks=[tasks[0]], + base_success_counts={ + retry._task_identity(task): 0 for task in tasks + }, + global_success_counts_at_start=snapshot, + ) + self.assertEqual( + manifest["policy"], + "global_below_threshold_until_additional_success_target", + ) + self.assertEqual(manifest["select_global_successes_below"], 1) + self.assertEqual( + manifest["target_additional_successes_per_variant"], + 1, + ) + self.assertEqual( + manifest["global_successes_at_start_by_task"], + snapshot, + ) + self.assertEqual( + manifest["wall_deadline"], + "2026-07-31T10:00:00+00:00", + ) + + def test_sequential_campaign_seed_ranges_must_not_overlap(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + prior_dir = output_dir / "retry_campaigns" / "prior" + collector._atomic_json_dump( + { + "retry_index_start": 44, + "retry_max_attempts": 96, + }, + prior_dir / "manifest.json", + ) + current_dir = output_dir / "retry_campaigns" / "current" + retry._validate_seed_interval_against_campaigns( + output_dir, + current_dir, + retry_index_start=140, + retry_max_attempts=96, + ) + with self.assertRaisesRegex(ValueError, "overlaps"): + retry._validate_seed_interval_against_campaigns( + output_dir, + current_dir, + retry_index_start=139, + retry_max_attempts=96, + ) + + def test_wall_deadline_is_timezone_aware_and_allows_finalize(self) -> None: + self.assertEqual( + retry._normalize_wall_deadline( + "2026-07-31T18:00:00+08:00" + ), + "2026-07-31T10:00:00+00:00", + ) + with self.assertRaisesRegex(ValueError, "explicit timezone"): + retry._normalize_wall_deadline("2026-07-31T18:00:00") + self.assertTrue( + retry._should_finalize( + skip_finalize=False, + interrupted=False, + completion_reason="wall_deadline", + ) + ) + self.assertTrue( + retry._should_finalize( + skip_finalize=False, + interrupted=False, + completion_reason="wall_deadline_already_reached", + ) + ) + + def test_past_wall_deadline_stops_before_worker_spawn(self) -> None: + args = Namespace( + gpus="0", + workers_per_gpu=1, + max_runtime_hours=12, + wall_deadline="2000-01-01T00:00:00+00:00", + ) + task = _task(1) + states = { + retry._task_identity(task): { + "task": task, + "base_successes": 0, + "successful_indices": [], + "target_reached": False, + "exhausted": False, + "next_index": 0, + } + } + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + result = retry._run_workers( + args, + output_dir=root, + campaign_dir=root / "campaign", + staging_root=root / "staging", + tasks=[task], + states=states, + state_path=root / "state.json", + persisted_state={}, + ) + self.assertEqual(result[0], "wall_deadline_already_reached") + + +class EvalState0LauncherTest(unittest.TestCase): + def _run_wrapper( + self, + wrapper_name: str, + extra_args: list[str] | None = None, + ) -> list[str]: + env = os.environ.copy() + env.update( + { + "PYTHON_BIN": "/bin/echo", + "LIBERO_PLUS_ROOT": "/test/LIBERO-plus", + "REFERENCE_DATASET": "/test/reference", + "MODEL_BASE_PATH": "/test/models", + "CHECKPOINT": "/test/checkpoint.pt", + "DATASET_STATS_PATH": "/test/stats.json", + } + ) + env.pop("PROMPT_MODE", None) + if wrapper_name.startswith("retry_"): + # The retry wrapper normally inherits this from the immutable base + # manifest. Supplying it explicitly keeps /bin/echo as the final + # command in this launcher-only test. + env["PROMPT_MODE"] = "benchmark" + with tempfile.TemporaryDirectory() as temporary_dir: + result = subprocess.run( + [ + "bash", + str(PROJECT_ROOT / "scripts" / wrapper_name), + temporary_dir, + "--prepare-only", + *(extra_args or []), + ], + cwd=PROJECT_ROOT, + env=env, + check=False, + capture_output=True, + text=True, + timeout=30, + ) + self.assertEqual( + result.returncode, + 0, + msg=f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}", + ) + command_line = result.stdout.strip().splitlines()[-1] + return shlex.split(command_line) + + def test_collection_preset_forwards_eval_state0_protocol(self) -> None: + argv = self._run_wrapper( + "collect_fastwam_libero_plus_robot_eval_state0_8gpu.sh" + ) + self.assertEqual(argv.count("--initialization"), 1) + self.assertEqual( + argv[argv.index("--initialization") + 1], + collector.TABLE13_EVAL_STATE0_INITIALIZATION, + ) + self.assertEqual(argv[argv.index("--init-state-indices") + 1], "0") + self.assertEqual(argv[argv.index("--policy-seeds") + 1], "42,43") + self.assertEqual(argv[argv.index("--prompt-mode") + 1], "benchmark") + + def test_retry_preset_has_one_initialization_and_target_two(self) -> None: + argv = self._run_wrapper( + "retry_fastwam_libero_plus_robot_eval_state0_8gpu.sh" + ) + self.assertEqual(argv.count("--initialization"), 1) + self.assertEqual( + argv[argv.index("--initialization") + 1], + collector.TABLE13_EVAL_STATE0_INITIALIZATION, + ) + self.assertEqual(argv[argv.index("--init-state-indices") + 1], "0") + self.assertEqual( + argv[argv.index("--target-successes-per-variant") + 1], + "2", + ) + self.assertEqual(argv[argv.index("--prompt-mode") + 1], "benchmark") + + def test_retry_preset_forwards_native_wall_deadline(self) -> None: + deadline = "2099-01-02T03:04:05+08:00" + argv = self._run_wrapper( + "retry_fastwam_libero_plus_robot_eval_state0_8gpu.sh", + ["--wall-deadline", deadline], + ) + self.assertEqual(argv.count("--wall-deadline"), 1) + self.assertEqual(argv[argv.index("--wall-deadline") + 1], deadline) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_eval_libero_plus_launchers.py b/tests/test_eval_libero_plus_launchers.py new file mode 100644 index 00000000..daf14acc --- /dev/null +++ b/tests/test_eval_libero_plus_launchers.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +import os +import subprocess +import tempfile +import unittest +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +FULL_LAUNCHER = ( + PROJECT_ROOT / "scripts" / "eval_fastwam_libero_plus_full_8gpu.sh" +) +SMOKE_LAUNCHER = ( + PROJECT_ROOT / "scripts" / "eval_fastwam_libero_plus_smoke_8gpu.sh" +) + +LAUNCHER_ENV_VARS = { + "CHECKPOINT_LOAD_PATH", + "DATASET_STATS_PATH", + "DEPS_ROOT", + "EGL_FALLBACK_GPU", + "EGL_LOCK_SCOPE", + "ENV_SEED", + "GPU_IDS", + "GRIPPER_ACTION_FORMAT", + "LIBERO_PLUS_ROOT", + "MODEL_BASE_PATH", + "NVIDIA_EGL_ROOT", + "PYTHON_BIN", + "POLICY_SEED", + "SAVE_VIDEOS", + "SMOKE_MAX_TASKS_PER_WORKER", + "SMOKE_SEED", + "SMOKE_TASKS", + "SMOKE_TRIALS", + "TABLE13_CATEGORIES", + "TEXT_EMBEDDING_CACHE_DIR", + "TOKENIZER_MODEL_ID", + "WORKERS_PER_GPU", + "WORKER_READY_TIMEOUT", +} + + +def _option_value(argv: list[str], option: str) -> str: + index = argv.index(option) + return argv[index + 1] + + +class LiberoPlusLauncherTest(unittest.TestCase): + def _run_launcher( + self, + launcher: Path, + overrides: dict[str, str] | None = None, + ) -> list[str]: + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + capture_path = root / "argv.txt" + fake_python = root / "fake-python" + fake_python.write_text( + "#!/usr/bin/env bash\n" + "set -euo pipefail\n" + ': > "${CAPTURE_PATH:?}"\n' + 'printf "%s\\n" "$@" >> "${CAPTURE_PATH}"\n', + encoding="utf-8", + ) + fake_python.chmod(0o755) + + env = os.environ.copy() + for variable_name in LAUNCHER_ENV_VARS: + env.pop(variable_name, None) + env.update( + { + "CAPTURE_PATH": str(capture_path), + "DATASET_STATS_PATH": "/test/stats.json", + "LIBERO_PLUS_ROOT": "/test/LIBERO-plus", + "MODEL_BASE_PATH": "/test/models", + "PYTHON_BIN": str(fake_python), + } + ) + if overrides: + env.update(overrides) + + output_dir = root / "eval-output" + result = subprocess.run( + [ + "bash", + str(launcher), + "/test/checkpoint/fastwam.pt", + str(output_dir), + ], + cwd=PROJECT_ROOT, + env=env, + check=False, + capture_output=True, + text=True, + timeout=30, + ) + self.assertEqual( + result.returncode, + 0, + msg=f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}", + ) + self.assertTrue(capture_path.is_file()) + return capture_path.read_text(encoding="utf-8").splitlines() + + def test_full_launcher_defaults_to_eight_workers_per_gpu(self) -> None: + argv = self._run_launcher( + FULL_LAUNCHER, + {"TEXT_EMBEDDING_CACHE_DIR": "/test/table13-cache"}, + ) + + self.assertEqual( + Path(argv[0]).resolve(), + ( + PROJECT_ROOT + / "experiments" + / "libero" + / "eval_libero_plus_persistent.py" + ).resolve(), + ) + self.assertEqual(_option_value(argv, "--workers-per-gpu"), "8") + self.assertEqual(_option_value(argv, "--egl-lock-scope"), "gpu") + self.assertEqual(_option_value(argv, "--worker-ready-timeout"), "1200") + self.assertEqual(_option_value(argv, "--gpus"), "0,1,2,3,4,5,6,7") + self.assertEqual(_option_value(argv, "--mode"), "full") + self.assertEqual(_option_value(argv, "--num-trials"), "1") + self.assertEqual( + _option_value(argv, "--text-embedding-cache-dir"), + "/test/table13-cache", + ) + self.assertIn( + "model.tokenizer_model_id=Wan-AI/Wan2.2-TI2V-5B", + argv, + ) + self.assertNotIn("--env-seed", argv) + self.assertNotIn("--policy-seed", argv) + self.assertNotIn("--category", argv) + self.assertIn("--resume", argv) + + def test_smoke_launcher_defaults_to_one_worker_per_gpu(self) -> None: + argv = self._run_launcher(SMOKE_LAUNCHER) + + self.assertEqual(_option_value(argv, "--workers-per-gpu"), "1") + self.assertEqual(_option_value(argv, "--egl-lock-scope"), "global") + self.assertEqual(_option_value(argv, "--worker-ready-timeout"), "1200") + self.assertEqual(_option_value(argv, "--mode"), "smoke") + self.assertEqual(_option_value(argv, "--smoke-tasks"), "14") + self.assertEqual(_option_value(argv, "--smoke-seed"), "42") + self.assertEqual(_option_value(argv, "--num-trials"), "1") + self.assertEqual(_option_value(argv, "--max-tasks-per-worker"), "0") + self.assertIn("--save-videos", argv) + self.assertNotIn("--text-embedding-cache-dir", argv) + self.assertNotIn("--env-seed", argv) + self.assertNotIn("--policy-seed", argv) + self.assertNotIn("--category", argv) + + def test_full_launcher_forwards_runtime_overrides(self) -> None: + argv = self._run_launcher( + FULL_LAUNCHER, + { + "CHECKPOINT_LOAD_PATH": "/local/checkpoints/fastwam.pt", + "EGL_FALLBACK_GPU": "6", + "EGL_LOCK_SCOPE": "global", + "GPU_IDS": "7,5", + "GRIPPER_ACTION_FORMAT": "zero_one_open_positive", + "ENV_SEED": "42", + "POLICY_SEED": "50", + "TABLE13_CATEGORIES": "Robot,Noise", + "TEXT_EMBEDDING_CACHE_DIR": "/local/table13-cache", + "TOKENIZER_MODEL_ID": "local/tokenizer", + "WORKERS_PER_GPU": "3", + "WORKER_READY_TIMEOUT": "77.5", + }, + ) + + self.assertEqual(_option_value(argv, "--workers-per-gpu"), "3") + self.assertEqual(_option_value(argv, "--gpus"), "7,5") + self.assertEqual(_option_value(argv, "--egl-lock-scope"), "global") + self.assertEqual(_option_value(argv, "--worker-ready-timeout"), "77.5") + self.assertEqual( + _option_value(argv, "--checkpoint-load-path"), + "/local/checkpoints/fastwam.pt", + ) + self.assertEqual( + _option_value(argv, "--text-embedding-cache-dir"), + "/local/table13-cache", + ) + self.assertEqual(_option_value(argv, "--egl-fallback-gpu"), "6") + self.assertEqual( + _option_value(argv, "--gripper-action-format"), + "zero_one_open_positive", + ) + self.assertEqual(_option_value(argv, "--env-seed"), "42") + self.assertEqual(_option_value(argv, "--policy-seed"), "50") + self.assertEqual(_option_value(argv, "--category"), "Robot,Noise") + self.assertIn("model.tokenizer_model_id=local/tokenizer", argv) + + def test_smoke_launcher_forwards_smoke_overrides(self) -> None: + argv = self._run_launcher( + SMOKE_LAUNCHER, + { + "EGL_LOCK_SCOPE": "gpu", + "ENV_SEED": "42", + "POLICY_SEED": "43", + "SAVE_VIDEOS": "0", + "SMOKE_MAX_TASKS_PER_WORKER": "1", + "SMOKE_SEED": "50", + "SMOKE_TASKS": "64", + "SMOKE_TRIALS": "2", + "TABLE13_CATEGORIES": "Robot", + "TEXT_EMBEDDING_CACHE_DIR": "/local/smoke-cache", + "WORKERS_PER_GPU": "8", + }, + ) + + self.assertEqual(_option_value(argv, "--workers-per-gpu"), "8") + self.assertEqual(_option_value(argv, "--egl-lock-scope"), "gpu") + self.assertEqual(_option_value(argv, "--smoke-tasks"), "64") + self.assertEqual(_option_value(argv, "--smoke-seed"), "50") + self.assertEqual(_option_value(argv, "--num-trials"), "2") + self.assertEqual(_option_value(argv, "--max-tasks-per-worker"), "1") + self.assertEqual(_option_value(argv, "--env-seed"), "42") + self.assertEqual(_option_value(argv, "--policy-seed"), "43") + self.assertEqual(_option_value(argv, "--category"), "Robot") + self.assertEqual( + _option_value(argv, "--text-embedding-cache-dir"), + "/local/smoke-cache", + ) + self.assertNotIn("--save-videos", argv) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_eval_libero_plus_persistent.py b/tests/test_eval_libero_plus_persistent.py new file mode 100644 index 00000000..74eb989c --- /dev/null +++ b/tests/test_eval_libero_plus_persistent.py @@ -0,0 +1,354 @@ +from __future__ import annotations + +import hashlib +import json +import sys +import tempfile +import unittest +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +EXPERIMENT_DIR = PROJECT_ROOT / "experiments" / "libero" +for path in (PROJECT_ROOT, PROJECT_ROOT / "src", EXPERIMENT_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +import eval_libero_plus_persistent as evaluator # noqa: E402 + + +class WorkerTopologyTest(unittest.TestCase): + def test_build_worker_slots(self) -> None: + self.assertEqual( + evaluator._build_worker_slots([2, 7], 3), + [ + ("gpu2-w0", 2), + ("gpu2-w1", 2), + ("gpu2-w2", 2), + ("gpu7-w0", 7), + ("gpu7-w1", 7), + ("gpu7-w2", 7), + ], + ) + + def test_build_worker_slots_rejects_non_positive_count(self) -> None: + with self.assertRaisesRegex(ValueError, "at least 1"): + evaluator._build_worker_slots([0], 0) + + def test_parser_runtime_defaults(self) -> None: + args = evaluator.build_parser().parse_args( + [ + "--checkpoint", + "checkpoint.pt", + "--dataset-stats", + "stats.json", + "--output-dir", + "output", + "--model-base-path", + "models", + ] + ) + self.assertEqual(args.workers_per_gpu, 1) + self.assertEqual(args.egl_lock_scope, "global") + self.assertEqual(args.worker_ready_timeout, 1200.0) + self.assertIsNone(args.checkpoint_load_path) + self.assertIsNone(args.env_seed) + self.assertIsNone(args.policy_seed) + self.assertEqual(args.category, []) + + +class CategoryFilterTest(unittest.TestCase): + def test_short_and_canonical_names_are_normalized_in_table_order( + self, + ) -> None: + self.assertEqual( + evaluator._normalize_category_filters( + [ + "Robot,Camera", + "Robot Initial States", + "Sensor Noise", + ] + ), + [ + "Camera Viewpoints", + "Robot Initial States", + "Sensor Noise", + ], + ) + + def test_unknown_category_is_rejected_with_valid_labels(self) -> None: + with self.assertRaisesRegex( + ValueError, + "Unknown Table 13 categories.*Robot", + ): + evaluator._normalize_category_filters(["not-a-category"]) + + def test_robot_only_smoke_sampling_uses_only_robot_tasks(self) -> None: + tasks = [ + evaluator.LiberoPlusTask( + suite="libero_spatial", + task_id=index, + classification_id=index + 1, + name=f"task-{index}", + category="Robot Initial States", + difficulty_level=(index % 5) + 1, + ) + for index in range(20) + ] + selected = evaluator.select_smoke_tasks(tasks, limit=8, seed=42) + self.assertEqual(len(selected), 8) + self.assertEqual( + {task.category for task in selected}, + {"Robot Initial States"}, + ) + + +class CheckpointLoadPathTest(unittest.TestCase): + def test_byte_identical_local_checkpoint_is_accepted(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + checkpoint = root / "canonical.pt" + local = root / "local.pt" + checkpoint.write_bytes(b"same-checkpoint") + local.write_bytes(b"same-checkpoint") + self.assertEqual( + evaluator._validate_checkpoint_load_path(checkpoint, local), + local.resolve(), + ) + + def test_same_size_different_checkpoint_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + checkpoint = root / "canonical.pt" + local = root / "local.pt" + checkpoint.write_bytes(b"checkpoint-a") + local.write_bytes(b"checkpoint-b") + with self.assertRaisesRegex(ValueError, "SHA256 differs"): + evaluator._validate_checkpoint_load_path(checkpoint, local) + + +class PromptCacheTest(unittest.TestCase): + def test_exact_prompt_cache_inventory(self) -> None: + descriptions = ["pick up the mug", "open the drawer", "pick up the mug"] + with tempfile.TemporaryDirectory() as temporary_dir: + cache_dir = Path(temporary_dir) + for description in sorted(set(descriptions)): + full_prompt = evaluator.WAN_PROMPT_TEMPLATE.format( + task=description + ) + digest = hashlib.sha256( + full_prompt.encode("utf-8") + ).hexdigest() + path = cache_dir / ( + f"{digest}.t5_len{evaluator.WAN_PROMPT_CONTEXT_LEN}." + f"{evaluator.WAN_PROMPT_ENCODER_ID}.pt" + ) + path.touch() + self.assertEqual( + evaluator._validate_text_embedding_cache( + cache_dir, + descriptions, + ), + {"tasks": 3, "unique_prompts": 2}, + ) + + def test_incomplete_prompt_cache_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + with self.assertRaisesRegex(FileNotFoundError, "missing=1/1"): + evaluator._validate_text_embedding_cache( + temporary_dir, + ["missing prompt"], + ) + + def test_prompt_cache_uses_configured_context_length(self) -> None: + description = "pick up the mug" + with tempfile.TemporaryDirectory() as temporary_dir: + cache_dir = Path(temporary_dir) + full_prompt = evaluator.WAN_PROMPT_TEMPLATE.format( + task=description + ) + digest = hashlib.sha256(full_prompt.encode("utf-8")).hexdigest() + ( + cache_dir + / ( + f"{digest}.t5_len256." + f"{evaluator.WAN_PROMPT_ENCODER_ID}.pt" + ) + ).touch() + self.assertEqual( + evaluator._validate_text_embedding_cache( + cache_dir, + [description], + context_len=256, + ), + {"tasks": 1, "unique_prompts": 1}, + ) + + +class ManifestCompatibilityTest(unittest.TestCase): + def test_runtime_only_topology_fields_do_not_break_resume(self) -> None: + compatibility_values = { + "mode": "full", + "checkpoint": "/persistent/checkpoint.pt", + "dataset_stats": "/persistent/stats.json", + "classification_path": "/persistent/tasks.json", + "task_config": "libero", + "smoke_seed": None, + "num_trials": 1, + "save_videos": False, + "gripper_action_format": "signed_open_negative", + "max_tasks_per_worker": 0, + "hydra_overrides": [], + "tasks": [{"suite": "libero_spatial", "task_id": 0}], + } + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + old_manifest = { + "schema_version": 1, + "created_at": "old", + "gpu_ids": list(range(8)), + **compatibility_values, + } + (output_dir / "manifest.json").write_text( + json.dumps(old_manifest), + encoding="utf-8", + ) + w8_manifest = { + "schema_version": 1, + "created_at": "new", + "gpu_ids": [0], + "workers_per_gpu": 8, + "checkpoint_load_path": "/local/checkpoint.pt", + "egl_lock_scope": "gpu", + "text_embedding_cache_dir": None, + **compatibility_values, + } + evaluator._write_or_validate_manifest( + output_dir, + w8_manifest, + resume=True, + ) + self.assertEqual( + json.loads( + (output_dir / "manifest.json").read_text(encoding="utf-8") + ), + old_manifest, + ) + + def test_legacy_manifest_accepts_default_seed_and_category_controls( + self, + ) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + legacy_manifest = { + "mode": "smoke", + "checkpoint": "/persistent/checkpoint.pt", + "dataset_stats": "/persistent/stats.json", + "classification_path": "/persistent/tasks.json", + "task_config": "libero", + "smoke_seed": 42, + "num_trials": 1, + "save_videos": False, + "gripper_action_format": "signed_open_negative", + "text_embedding_cache_dir": None, + "max_tasks_per_worker": 0, + "hydra_overrides": [], + "tasks": [{"suite": "libero_spatial", "task_id": 0}], + } + (output_dir / "manifest.json").write_text( + json.dumps(legacy_manifest), + encoding="utf-8", + ) + current_manifest = { + **legacy_manifest, + "env_seed": None, + "policy_seed": None, + "seed_fallback": "cfg.seed", + "category_filters": [], + } + evaluator._write_or_validate_manifest( + output_dir, + current_manifest, + resume=True, + ) + + def test_explicit_seed_cannot_resume_a_legacy_manifest(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + legacy_manifest = { + "mode": "smoke", + "checkpoint": "/persistent/checkpoint.pt", + "dataset_stats": "/persistent/stats.json", + "classification_path": "/persistent/tasks.json", + "task_config": "libero", + "smoke_seed": 42, + "num_trials": 1, + "save_videos": False, + "gripper_action_format": "signed_open_negative", + "text_embedding_cache_dir": None, + "max_tasks_per_worker": 0, + "hydra_overrides": [], + "tasks": [{"suite": "libero_spatial", "task_id": 0}], + } + (output_dir / "manifest.json").write_text( + json.dumps(legacy_manifest), + encoding="utf-8", + ) + seeded_manifest = { + **legacy_manifest, + "env_seed": 42, + "policy_seed": 50, + "seed_fallback": "cfg.seed", + "category_filters": [], + } + with self.assertRaisesRegex( + ValueError, + "env_seed, policy_seed", + ): + evaluator._write_or_validate_manifest( + output_dir, + seeded_manifest, + resume=True, + ) + + def test_switching_from_online_prompt_encoding_to_cache_is_rejected( + self, + ) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + output_dir = Path(temporary_dir) + old_manifest = { + "mode": "full", + "checkpoint": "/persistent/checkpoint.pt", + "dataset_stats": "/persistent/stats.json", + "classification_path": "/persistent/tasks.json", + "task_config": "libero", + "smoke_seed": None, + "num_trials": 1, + "save_videos": False, + "gripper_action_format": "signed_open_negative", + "text_embedding_cache_dir": None, + "max_tasks_per_worker": 0, + "hydra_overrides": [], + "tasks": [{"suite": "libero_spatial", "task_id": 0}], + } + (output_dir / "manifest.json").write_text( + json.dumps(old_manifest), + encoding="utf-8", + ) + cached_manifest = { + **old_manifest, + "text_embedding_cache_dir": "/local/table13-cache", + } + with self.assertRaisesRegex( + ValueError, + "text_embedding_cache_dir", + ): + evaluator._write_or_validate_manifest( + output_dir, + cached_manifest, + resume=True, + ) + + +if __name__ == "__main__": + unittest.main()