scripts/WIDOWX.md

WidowX / full-Bridge adaptation (updated 2026-09-24)

This trains the released Zeva Stage-2 Cosmos model, selecting its action heads, behavior injection modules and proprio encoder for fine-tuning. The Qwen/MoT backbone, VAE, CTE and static retrieval head remain frozen. This is an initial cross-robot adaptation, not a reproduction of a WidowX-trained release.

The smoke root is /opt/zhangchenyu/datasets/bridge_orig_lerobot_smoke. The old sampler used only 397 of its 654 metadata episodes (1,255 windows). The new full sampler includes short episodes and all stride-4 start positions. Camera image_0, 5 FPS, 256-pixel resolution, batch size 1. Per-window task IDs are read from parquet. Six motion dimensions use the documented Bridge quantile constants; gripper is -1 closed / +1 open. Proprio is Bridge's eight raw state fields plus the repeated gripper value, matching the inference adapter's nine fields.

prepare_widowx_features.py caches frozen release CTE features for every window. The phase/effect features use only observations and actions up to that window's start, including late-trajectory windows. Global context is retrieved from the episode's first observation and instruction; this bank is still RoboCasa-trained.

Use the full-IPEC command below after reserving GPUs; do not invoke the generic smoke launcher with its historical GPU/cache defaults. Both launchers use the existing Cosmos venv. The practical step limit is 2,147,483,647; training continues across data epochs and does not stop after an evaluation. The installed Cosmos source/venv are reused, with compatibility fixes scoped to the launcher. No new model download is required.

Run python scripts/watch_widowx_eval.py to watch atomic delta checkpoints. For saved multiples of 2,000 steps, the selected GPU serves the release plus that trained delta and runs ten complete widowx_carrot_on_plate episodes, seeds 0–9, in the existing SimplerEnv environment. --exec-horizon controls how many predictions are executed before replanning (1–32, default 4). It is independent of the number of denoising steps. The causal history records every executed control and makes a CTE boundary every four controls, including when execution horizons differ from four. XYZ commands are denormalized once; Euler rotation increments are converted to axis-angle. Only complete evaluations get success-rate JSON files; errors remain errors.

Checkpoint output is $WIDOWX_RUN_DIR/checkpoints/. Checkpoints contain selected parameters, optimizer, scheduler, scaler, RNG, a data contract and per-rank CONSUMED data cursors. They require the original Stage-2 weights. All checkpoints are retained; none are automatically deleted, even after evaluation. Writes are atomic, existing checkpoint files cannot be overwritten, and saving fails loudly below 2 GiB free space. This does not reserve disk space against other jobs.

Full IPEC run

bash scripts/run_ipec_full_training.sh uses the full local IPEC Bridge export and .artifacts/ipec_full_v2 for new features and training state. The old .artifacts/ipec_full run is preserved as a historical baseline. As of 2026-09-24, the WidowX recipe and TOML both save every 10,000 optimizer steps. Existing checkpoints are retained; the launcher resumes the latest complete checkpoint when explicitly started. It does not start an evaluator automatically.

The default recipe records W&B metrics in offline mode. Records are stored under $WIDOWX_RUN_DIR/train/<project>/widowx/<name>/wandb/. Logging begins at the resumed checkpoint's step; previously disabled W&B history is not reconstructed automatically. Training state remains in checkpoints/; W&B logging does not replace model checkpoints.

To send a future full-IPEC run to your own W&B account, set WANDB_API_KEY and WANDB_ENTITY in the launching shell. Do not put the API key in this repository or replace the shared machine's W&B login. Pass your project and a new run name as Hydra overrides; the full-IPEC launcher forwards these to the trainer:

# Example after WANDB_API_KEY and WANDB_ENTITY are exported in this shell:
CUDA_VISIBLE_DEVICES=2,4 NPROC_PER_NODE=2 WANDB_MODE=online \
  bash scripts/run_ipec_full_training.sh -- \
  job.project=YOUR_PROJECT job.name=ipec_full_v2 job.wandb_mode=online

The new run name gives W&B a fresh ID and local directory. The model still resumes the latest delta under $WIDOWX_RUN_DIR/checkpoints/ because that directory is controlled by WIDOWX_RUN_DIR, independently of the W&B run name. The offline history from earlier training is separate; after authenticating, it can be previewed with wandb sync --dry-run and uploaded with wandb sync using --entity and --project on its offline-run-* directory.

The launcher defaults WANDB_BASE_URL to https://api.wandb.ai, overriding unrelated global W&B settings on this shared account. Set WANDB_BASE_URL explicitly before launching if using a self-hosted W&B server. If online W&B initialization raises a W&B error, the local training entrypoint retries once in offline mode with the same project, name and persisted run ID. An offline initialization failure or an unrelated training error is not hidden.

The full-IPEC launcher defaults to GPU 1 and one process. To use a different reserved GPU set, supply both CUDA_VISIBLE_DEVICES and NPROC_PER_NODE in the environment. The TOML uses data_parallel_shard_degree = -1, so FSDP infers the shard count from the launched world size. Eight GPUs with one sample per rank yield an effective batch size of eight (two GPUs: two). Check GPU processes and availability before launching; a momentary zero utilization reading does not mean a card is unoccupied.

To evaluate one checkpoint on 20 carrot-on-plate episodes (seeds 0–19), with videos and per-control action/state traces:

WIDOWX_RUN_DIR=/home/gpu4/tianyi/Zeva/.artifacts/ipec_full WIDOWX_EVAL_GPU=1 \
  /home/gpu4/tianyi/cosmos-framework/.venv/bin/python scripts/watch_widowx_eval.py \
  --once --checkpoint /home/gpu4/tianyi/Zeva/.artifacts/ipec_full/checkpoints/step_000290000.pt \
  --episodes 20 --port 18785 --record-traces \
  --num-steps 30 --max-steps 150 --exec-horizon 8

Choose a free GPU and port before launching. This mode exits after the requested evaluation and never resumes training or backfills other checkpoints. Results are written under eval/, MP4s under eval/videos/step_000290000/, and traces under eval/traces/step_000290000/. It refuses to overwrite an existing completed result for the checkpoint; use a separate WIDOWX_RUN_DIR for repeat evaluations. The watcher resolves the GPU to its UUID and checks the renderer's actual GPU before loading the policy. Rendering requests cuda:0 within that visible set. This installed SAPIEN version can confuse identical PCI bus numbers in different domains: GPU 4 incorrectly renders on GPU 0 even with cuda:0. The preflight rejects that mismatch instead of starting a rollout on someone else's card. GPU 2 passed the physical-device probe; check its workload before evaluation. WIDOWX_RENDER_GPU can select a separate, verified renderer GPU (defaults to WIDOWX_EVAL_GPU). Both physical UUIDs are recorded in the evaluation report. GPU 6 also collides with GPU 2 in this SAPIEN version, so this experiment uses policy GPU 6 and renderer GPU 2 explicitly, not an unverified renderer on GPU 6.

For a controlled sampling comparison, pass --num-steps 30 --max-steps 150 to the watcher. These set the policy's UniPC denoising count and the simulator's control-step limit, respectively. --seed-start defaults to 0; disjoint ranges can run in separate artifact directories and use the same per-step inference seeds as a single sequential run. Completed reports record inference_num_steps.

Full sampling and cache provenance

The local IPEC export has 53,192 episodes and 1,893,026 recorded steps. The old sampler used 30,770 episodes and 99,075 windows. The new sampler uses 53,192 episodes / 492,160 windows, starts at 0, 4, 8, ... through the last recorded frame, and keeps the model's 32-action / 33-image shape. Tail images repeat the last observation and actions are padded, but separate validity masks exclude padded action, vision and Gaussian-prior targets from their losses. The terminal recorded action remains supervised even if its next image was not exported. This changes loss normalization; old and new scalar losses are not directly comparable. Full sampling does not synthesize missing language annotations or balance task frequencies.

The current v2 feature cache is stored beside the raw dataset at /home/gpu4/tianyi/TTTVLA/datasets/bridge_orig_1.0.0_lerobot_zeva_features_v2. Its manifest.json binds metadata SHA256, model-file identities and versioned preprocessing contracts to the cache. Small model files are hashed; large shards use path, size and modification time, so this is provenance validation, NOT a full dataset-content integrity scan. Merely adding a fingerprint does not alter features. This migration prepares a separate cache because old caches lack provenance and omit late prefixes/short episodes. The old cache is not deleted. The full launcher runs one preparation worker per selected GPU, validates every episode's feature tensors, and starts training only after all are present. Interrupting preparation is safe; rerun to validate and skip completed episodes.

Do not set WIDOWX_MAX_EPISODES for production; it is an optional smoke-test limit included in the checkpoint's data contract. WIDOWX_WINDOW_MODE=legacy exists only for explicit historical diagnostics, not the full-IPEC launcher.

Resume semantics

WIDOWX_RESUME=auto (default) selects the highest saved optimizer step in the current checkpoint directory. This is unrelated to trainer.max_iter, the total optimizer-step stopping limit, and to --max-steps, the evaluation episode limit. WIDOWX_RESUME=none requires a new/empty checkpoint directory. An explicit checkpoint path is also accepted; use a new output directory when branching.

New v3 checkpoints restore the next actually consumed data window, discarding stale prefetched batches. Exact sample-order continuity currently requires the same data/cache contract, GPU count, seed, one worker/rank and one packed sample/rank. It is not a promise of bitwise-identical GPU arithmetic. Old v2 checkpoints do not contain this cursor and cannot retroactively recover it.

Changing GPU count is rejected by default. To deliberately re-shard data, set WIDOWX_ALLOW_DATA_RESHARD=1; this begins the next shuffled data epoch instead of pretending to preserve exact sample order. WIDOWX_ALLOW_LEGACY_RESUME=1 is an explicit compatibility escape hatch for old/data-changed checkpoints; it is NOT the recommended migration for the corrected full-data recipe. Restart this corrected experiment in a new directory from released Stage-2.

Manual checkpoint deletion

Nothing is pruned automatically. Check sizes/free space first:

df -h /home/gpu4/tianyi/Zeva
ls -lh /home/gpu4/tianyi/Zeva/.artifacts/ipec_full/checkpoints/
ls -lh /home/gpu4/tianyi/Zeva/.artifacts/ipec_full_v2/checkpoints/

Example: explicitly select and interactively confirm deletion of ONLY an old 2,000-step checkpoint (do not run this unless you want to delete that file):

rm -i -- /home/gpu4/tianyi/Zeva/.artifacts/ipec_full/checkpoints/step_000002000.pt

Keep the latest checkpoint, any checkpoint currently being loaded/evaluated, and your comparison baselines. Copy to another disk before deletion if recovery may be needed; rm does not use a trash bin. These deltas also depend on the released base model: do not delete /home/gpu4/tianyi/zeva-release/weights/stage2. Stopping a trainer does not force a final save; unsaved progress may be lost.

Frozen behavior memory

This change does not claim to adapt the release's five-task RoboCasa memory bank to Bridge. CTE/retrieval are frozen and precomputed; changing requires_grad alone cannot train them. A proper adaptation trains the Stage-1 CTE on Bridge, rebuilds a Bridge effect bank, adapts Stage-3 retrieval, then rebuilds features and trains Stage-2. An ablation can measure whether the current bank helps or harms. Making horizon 32 smaller does not fix memory-domain mismatch, and changes to that horizon must also update behavior-PBD heads and checkpoint shapes. Keep horizon 32 and use masked tail windows for now.

External-bank ablation

--memory-mode bank_off does not load the external task-context bank or retrieval head, and supplies zero global context. It keeps the CTE, within-episode phase and effect history, and all policy weights (including learned injection biases). The default on is unchanged. The client checks the server's memory-mode metadata before running any episode. This is an inference-only ablation, not training without memory or disabling every Zeva component.

Run a matched bank-off/on comparison, each with 20 episodes, 30 denoising steps, 150 control steps and guidance 3, after checking GPU availability:

WIDOWX_EVAL_GPU=6 WIDOWX_RENDER_GPU=2 OMP_NUM_THREADS=2 \
  /home/gpu4/tianyi/cosmos-framework/.venv/bin/python scripts/run_widowx_memory_ablation.py \
  --checkpoint /home/gpu4/tianyi/Zeva/.artifacts/ipec_full/checkpoints/step_000290000.pt \
  --output /home/gpu4/tianyi/Zeva/.artifacts/CHOOSE_NEW_ABLATION_DIRECTORY \
  --episodes 20 --exec-horizon 16 --port 18825

The execution horizon is independent of the 32-action prediction horizon and the 30-step denoising count. At 5 Hz, executing 16 actions means replanning every 3.2 seconds (four actions: 0.8 seconds). History still records all executed controls at four-control CTE boundaries. Each mode saves videos and JSON traces under its own eval/; comparison.json appears only after BOTH modes finish. The controller validates checkpoint, seeds and protocol before accepting reports. --after /absolute/path/to/previous/comparison.json queues a comparison without loading a model until the preceding comparison has completed successfully.

The 2026-09-24 runs are .artifacts/ipec_memory_ablation_290000_20260924 (four actions) and .artifacts/ipec_memory_ablation_290000_h16_20260924 (16 actions, queued after the first). The latter runs in tmux Zeva:memory_eval16, with launcher.log in its output directory. These evaluate the old 290k baseline; the corrected full-data pipeline is still preparing its feature cache.

See Bridge memory audit and implementation plan for the actual released bank contents, missing training/export components, data-label requirements, and disk-space estimates. No Bridge bank has been built or substituted as part of this diagnostic evaluation.

CPU regressions:

PYTHONPATH=$PWD:/home/gpu4/tianyi/cosmos-framework CUDA_VISIBLE_DEVICES='' \
  /home/gpu4/tianyi/cosmos-framework/.venv/bin/python -m unittest \
  scripts.test_widowx_full scripts.test_fp32_master_adamw scripts.test_wandb_fallback \
  scripts.test_widowx_memory -v

The actual rollout loop also has CPU-only fake-environment tests for horizons 3/4/16/32, all 150 control steps, partial final chunks, history and mode checks. Run these with the installed SimplerEnv dependencies (no GPU context is created):

PYTHONPATH=$PWD:/opt/benchmarks/SimplerEnv CUDA_VISIBLE_DEVICES='' \
  /home/gpu4/miniconda3/envs/simpler_env/bin/python -m unittest \
  scripts.test_widowx_execution -v