Dev/flash training - #78
Closed
CharioteerStark wants to merge 19 commits into
Closed
Conversation
- docs/STAGE_A_TO_B_PLAN.md: end-to-end plan for embodiment (Stage A) → unseen-task (Stage B) finetuning on a custom bimanual robot, grounded in the paper's YAM and AgiBot post-training recipes. - docs/DATA_COLLECTION_STAGES_zh.md: data-team-facing spec describing what to capture for each stage, including the 11 YAM reference tasks (verbatim instructions from dreamzero0.github.io/yam_gallery) and paragraph-style instruction examples for multi-stage Stage B tasks (verbatim from paper Tables 5 & 6). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Paper-aligned recipe (§4.2): all params trainable except text/image encoders and VAE. Mirrors the existing scripts/train/droid_training_ full_finetune_wan21.sh pattern. Recommended for setups with ≥8× H200 NVLink — see docs/STAGE_A_TO_B_PLAN.md §3 note. Changes: - train_architecture: lora → full - save_lora_only: true → false - defer_lora_injection flag removed - deepspeed: zero2 → zero2_offload (configurable via DEEPSPEED_CFG) - dataloader_pin_memory: false → true; num_workers 1 → 4 - NUM_GPUS default 3 → 8 - GPU_MEM_FRACTION default 0.70 → 0.90 - SAVE_TOTAL_LIMIT 5 → 3 (full checkpoints are ~30 GB each) - output_dir adam_stage_a_lora → adam_stage_a_full - updated header / comments to reflect Full-FT
- Add Adam (14-DoF bimanual) to embodiment enum and conversion script - Add adam_relative.yaml + Adam entries in base_48_wan_fine_aug_relative - Add scripts/train/adam_stage_b.sh for unseen-task specialization - Add scripts/utils/merge_lora.py for Stage A → Stage B handoff - Pin uv.lock for reproducible environment Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Diagnosed via scripts/open_loop_adam.py that Adam's 5k-step LoRA at batch 1 cannot relearn the absolute->relative gripper semantic flip introduced by adam_relative.yaml's relative_action_keys. The gripper heads collapse to a per-head constant (~67 for left, ~83 for right), while joint predictions are accurate (MSE 0.001-0.004 rad^2). YAM uses the same config but trains 100k steps at batch 4 with save_lora_only, which is enough signal for the gripper head to track conditioning. Adam Stage A now mirrors yam_training.sh: train_architecture=lora, per_device_train_batch_size=4, max_steps=100000, save_total_limit=10, deepspeed=zero2, dataloader_pin_memory=false, dataloader_num_workers=1, save_lora_only=true, +defer_lora_injection=true. save_steps stays at 1000 for frequent early checkpoints (vs YAM's 10000). Also adds scripts/open_loop_adam.py (structural clone of open_loop_yam for the Adam embodiment), the WAM inference server (serve_wam.py + scripts/inference/serve_wam.sh), VM setup docs, and AdamRoboarenaPolicy in socket_test_optimized_AR.py. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- STAGE_A_TO_B_PLAN.md §2.4: training block now reflects YAM-aligned recipe (max_steps=100000, bs=4, deepspeed=zero2, save_lora_only=true, +defer_lora_injection=true), with save_steps=1000 kept low for early debugging. Add inline rationale referencing the gripper-head convergence issue (see §8). - §2.5: add quantitative open-loop check via scripts/open_loop_adam.py with per-key MSE thresholds. - §3: note that Stage A LoRA-100k matches YAM's validated recipe. - §6: update LoRA-vs-FullFT decision row with the >=50k effective sample-updates floor. - §7: add scripts/train/adam_stage_a.sh and scripts/open_loop_adam.py to the reference file list. - §8: new risk entry documenting gripper-head absolute->relative semantic flip introduced by relative_action_keys, observed failure mode (predicts training mean), and the cheap mitigation if shorter training is required. - VM_SETUP.md: key-paths table reflects LoRA recipe and adds scripts/open_loop_adam.py. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Align save_steps (1000 -> 10000) so the Adam embodiment-adaptation recipe is identical to scripts/train/yam_training.sh (LoRA, zero2, bs=4, 100k steps, LR=1e-5, save_lora_only=true, from DreamZero-AgiBot, single dataset). Make the conda-env launcher VM-portable: fall back to plain `torchrun` when the local env path is absent. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Make all English docs/code consistent with the settled training recipe (LoRA, zero2, bs=4, 100k steps, from DreamZero-AgiBot — mirrors yam_training.sh) and the real-robot inference stack (serve_wam.py + deploy_adam.py). - Replace STAGE_A_TO_B_PLAN.md with STAGE_A_PLAN.md (Stage-A-only): save_steps=10000, measured gripper MSE, serve_wam/deploy_adam inference, paper-vs-repo LoRA note. - Remove scripts/train/adam_stage_b.sh (task-specialization deferred). - Fix VM_SETUP smoke test (was serve_wam + test_client_AR — protocol mismatch) to the hardware-free open_loop_adam.py check; refresh key-paths table. - Update embodiment_tags.py and merge_lora.py docstrings; fix zh-doc plan link. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ata fix Everything needed to run Adam Stage A embodiment adaptation on the UW Tillicum H200 SLURM cluster, plus the dataset fix that unblocks bs>1. Setup & docs - docs/TILLICUM_SETUP.md: full terminal/SSH runbook (access, storage, conda env, build gotchas, weight/data staging, smoke, train, monitor) - README.md / docs/STAGE_A_PLAN.md: measured 8x H200 configs + bs ceilings SLURM infra (scripts/slurm/) - build_env_tillicum.slurm, smoke_adam_tillicum.slurm - train_adam_full_tillicum.slurm (full-FT, zero2_offload, bs=4) - train_adam_tillicum.slurm (LoRA, zero2, bs=3) - submit_chain.sh: chained 24h afterany slices, auto-resume - timing_probe.slurm / bs_probe.slurm: batch-size / throughput sweeps - wandb_sync_loop.sh: login-node offline->dashboard sync (fallback) - scripts/train/adam_stage_a_full.sh: full-parameter recipe Data - scripts/data/convert_wam_raw_to_lerobot.py: raw WAM (ZED svo2 + teleop json) -> LeRobot v2; prepare_wam_tillicum.sh wrapper bs>1 fix - groot/vla/data/dataset/lerobot_sharded.py: REQUIRE_UNIFORM_CHUNKS guard. Edge-of-episode windows form fewer than max_chunk_size chunks (state (3,64) vs (4,64), video 25 vs 33 frames); individually model-valid (bs=1 trains) but won't np.stack -> bs>=2 collate crash. The guard drops those non-uniform windows so every sample is uniform. Gated; bs=1 unaffected. - scripts/diag_*.py: dataset/collate diagnostics used to pin it down Measured on 8x H200: full-FT bs=4 (zero2_offload), LoRA bs=3 (zero2); higher batch OOMs. save_steps must fit the 24h walltime. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- deploy_adam.py: gripper binary default (thresh 30); reanchor default OFF + --reanchor-skip; --chunk-tail-skip; joint-torque force guard (--force-stop/--estop-torque, stop+regenerate, per-arm/joint trip log); --tcp-payload-kg gravity comp; --inference-freq default 30. - action_chunk_broker.py: reanchor-skip/chunk-tail-skip; force_replan + generation guard (discard stale in-flight chunk); infer lock (no concurrent ws recv); closed-loop freq logging. - serve_wam.py: disable world-model video saving by default. - scripts/inference/deploy_adam.sh: client launcher. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Profiling + ranked non-Flash speedup plan: steps 16->4, DiT-cache tuning, torch.compile/CUDA-graphs, TensorRT, NVFP4, KV-cache streaming. Targets ~1-3 Hz on RTX PRO 6000 (GB200/7Hz + Flash out of scope). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- .sh referenced missing build_trt_engine_droid.py -> build_trt_engine.py - add conda-env activation (was running base python3.13 -> ModuleNotFoundError tianshou) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…xe_droid KeyError 'oxe_droid' when building for the Adam checkpoint (metadata keyed 'adam'). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…Adam - adam_stage_a_full.sh: EXTRA_OVERRIDES passthrough so wrappers can flip variant flags without forking the validated script. - train_adam_flash_tillicum.slurm: full-FT Flash slice — inits (weights-only) from adam_stage_a_full/checkpoint-7500, flips decouple_video_action_noise=true + video Beta(7,1) per the paper, fresh OUTPUT_DIR/wandb id, otherwise the validated Stage-A recipe (zero2_offload, bs4, LR 2e-5, 7500 steps, 24h chain). Verified: training builds the model from the launch hydra config (checkpoint config.json is not read in training), so the overrides take effect. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.