Zero-WAM w/o pretrain 재현 학습 완료

TL;DR

목표: 논문에 절반만 보고된 "Zero-WAM w/o pretrain"(ICL 36.36%만 공개, language-only 미공개) 설정을 공개 코드로 재현해, human demo 없이(language-only) 돌렸을 때 ICL 학습의 순 기여를 분리
init 변환 검증: Wan2.2-TI2V-5B → Zero-WAM(10.79B) 변환, patch_embedding_mlp vs Conv3d 최대 오차 0.00, 저장 후 재로딩 bit-exact
학습 완주: 32,000 step × 8 GPU(own, 무중단), 20h 52m, latent loss 0.62→0.11, action loss 0.79→0.005
평가 파이프라인 사전검증: 공개 posttrain 체크포인트로 place_empty_cup ICL 7/10(70%) — 논문 84.87±0.18%와 표준오차(±14pp) 범위 내 일치
language-only / ICL 두 모드 7 unseen task × 100 episode 평가를 제출, 결과는 아직 없음(Next 참조)

10.79B

Params

32,000

Steps

20:52:26

학습 시간

0.00

변환 오차

1 배경 / 목적

Zero-WAM 논문은 Wan2.2-TI2V-5B에서 바로 RoboTwin 43개 태스크로 ICL 학습한 "w/o pretrain" 설정을 다룬다. 하지만 공개된 수치는 ICL 모드 36.36%(main.tex:1084-1091)뿐이고, 같은 모델을 human demo 없이 language-only로 평가한 수치는 없다.

비교 기준: 같은 init에서 language로만 학습·추론한 WAN-Action 10.98%, 자체 재현한 LingBot-VA 27.86%. language-only 수치가 WAN-Action(~11%)에 가깝다면 ICL 학습의 이득은 전부 "추론 시 human video 조건"에서 오는 것이고, 그보다 크게 높다면 ICL 학습 자체가 language-conditioned policy의 표현력도 함께 개선한다는 뜻이다.

기존 한계: 공개 체크포인트(zero-wam-pretrain)는 이미 사전학습이 끝난 모델이라 이 대조 실험에 쓸 수 없다. Wan2.2 원본에서 직접 Zero-WAM 아키텍처로 초기화하는 변환이 선행되어야 한다.

2 작업 내용

2.1 Wan2.2 → Zero-WAM init 변환 (`script/convert_wan22_to_zerowam.py`)

WanICLTransformer3DModel(10.79B)은 Wan2.2 5B에 없는 파라미터를 가지고 있어 from_pretrained로 그대로 로드할 수 없다. 논문 기술을 따라 다음과 같이 채웠다.

파라미터	채운 값	근거
video branch	Wan2.2 그대로	—
action branch (action_to_q/k/v, action_ffn, action_norm2, scale_shift_table_action, condition_embedder_action)	video branch 복사	main.tex:835-836 (d_a = d_v = 3072)
IFP/MCP block × 4	마지막 video block(29) 복사	main.tex:656-657
patch_embedding_mlp	Wan2.2 Conv3d weight reshape(3072, 192)	ICL forward가 Conv3d 대신 이 Linear를 사용 (icl_model.py:789)
action_embedder, action_proj_out, mcp_mlp_hidden, mcp_projections	생성자 기본 init(seed 0)	논문 미기술

저장은 diffusers save_pretrained 대신 train.py의 _prepare_safetensors_state_dict + _save_sharded_safetensors를 사용했다 — attn2의 action_to_k/v, action_norm_k가 video 쪽 모듈과 같은 객체라(icl_model.py:293-295) save_pretrained가 shared-tensor 에러를 낸다. 출력: ~/checkpoints/zero_wam/wan22_ti2v5b_zerowam_init/ (transformer 22GB + vae/text_encoder/tokenizer는 Wan2.2 symlink).

2.2 코드 수정

wan_va/train.py: --resume-from, --start-step 추가 (체크포인트가 weight만 저장하므로 step 이름 이어붙임, optimizer 상태는 미복원) evaluation/robotwin/eval_policy_client_openpi.py: --use_icl {1,0} 추가(원래 True 고정), task/camera config 경로 CONFIGS_PATH 기준으로 변경 evaluation/robotwin/deploy_zerowam.yml: fork에 없는 policy/ACT/deploy_policy.yml 대체

2.3 평가 인프라 — 공유 checkout을 오염시키지 않는 overlay

script/make_robotwin_eval_root.sh → tmp/robotwin_eval_root: 전부 symlink로 두고 envs/move_stapler_pad.py, stamp_seal.py만 Zero-WAM 수정 성공기준으로 교체. RoboTwin이 abspath(envs/_GLOBAL_CONFIGS.py)로 루트를 계산해 symlink를 안 따라가므로 overlay 루트로 정상 인식된다. fork ↔ upstream(2eeec322) 비교: unseen 7개 중 5개는 동일, 나머지는 success 판정과 무관한 차이(arm_tag 추가, B200 OIDN 환경변수)만 확인.

2.4 학습 스크립트 (`script/train_zerowam_wopretrain.sh`)

config: robotwin_train, DATASETS=robotwin:1.0 (unseen 7태스크 제외, 43개 아니면 에러) drop_icl 0.1, droptext_target 0.4, IFP on, lr 1e-4, batch 1/GPU 32K step × 8 GPU = 256K samples (자체 LingBot-VA와 동일 예산) 4000 step마다 저장, 최신 체크포인트에서 자동 resume 제출: own QOS 8 GPU (job 127257, n7)

2.5 운영 — 체크포인트 자동 정리 + 학습 중 데이터셋 이동

script/prune_checkpoints.sh <ckpt_dir> 2 600 32000: 10분마다 완료된(config.json 존재) 최신 2개만 남기고 삭제, 32000 완료 시 종료 — 실행 중인 train.py가 코드 변경을 못 읽으므로 별도 CPU job(1 CPU, 2GB)으로 분리했다. 학습 중 데이터셋 경로를 /home/nas_main/datasets/Zero-WAM/HumanGen → /home/nas_main/datasets/robot/Zero-WAM/HumanGen로 이동 — 학습이 읽는 중이라 os.rename 직후 옛 경로에 호환 symlink를 만들고, repo data/HumanGen은 os.replace로 원자적 교체. job 시간 제한도 48h → 31h로 축소(클러스터 규칙 §19: 실제 소요 + 20~30%).

3 결과 (수치)

3.1 Init 변환 검증

copied 1,694 tensors, 새 초기화 16개, video와 tied state key 170개, 총 파라미터 10.79B (README 10.8B와 일치)
patch_embedding_mlp vs Conv3d 최대 오차 0.00
저장 후 load_transformer(train.py와 동일 경로) 재로딩: 6개 텐서 bf16 왕복 후 bit-exact

3.2 평가 파이프라인 사전검증 (job 127202, 공개 zero-wam-posttrain-robotwin, place_empty_cup, seed 0)

모드	설정	결과
icl	use_icl=1, ICL CFG 5, target text CFG -1, human latent prefix cache	7/10 (70%)
lang	use_icl=0, text CFG 5, ICL cache 없음	2/2 (경로 확인용)

파이프라인 정상 동작: 논문 place_empty_cup 84.87±0.18%. n=10 표준오차 약 ±14pp라 icl 70%는 같은 범위 — fork + overlay(수정 성공기준) + client 패치 + 서버 조합이 정상 확인됨. 서버 로그의 Traceback은 client 종료 시 cleanup SIGTERM(SignalException: got signal 15)으로 오류가 아니다.

3.3 학습 진행 — dataset index 2,149 samples / 43 datasets

Step	s/it	latent_loss	action_loss	mcp(IFP)_loss	grad_norm
448	2.61	0.62 → 0.26	0.79 → 0.02	2.0 → 0.27	—
4,087	2.35	0.198	0.011	0.223	0.07
12,108	2.33	0.1601	0.0064	0.2227	0.05
31,999 (최종)	2.34 (평균)	0.1059	0.0052	0.1814	0.05

완주: 2026-09-16 00:27 UTC 완료, 총 20:52:26, 32,000 step × 8 GPU own QOS, preemption 없이 무중단. lr 1e-4 도달(warmup 200 step).

3.4 운영 검증

checkpoint pruning 가짜 폴더 테스트(4000/8000/12000 완료 + 16000 작성 중, keep=2, final=12000) → 4000만 삭제, 16000 유지 후 정상 종료. 실제 운영에서도 2026-09-15 11:29:21 UTC checkpoint_step_4000(22GB) 삭제, 8000/12000 유지 확인.
체크포인트 config.json을 extract_init_dict로 풀면 init 인자 24개가 공개 posttrain config와 차이 0 → 평가 서버 바로 로드 가능.
데이터셋 경로 이동 직후 학습 step 447, 에러 없음 확인.

4 Takeaway

의미

논문에 절반만 보고된 "w/o pretrain" 설정의 재현 학습이 수치 검증을 동반해 완료됐다. init 변환은 0.00 오차·bit-exact 재로딩으로 검증됐고, 평가 파이프라인도 공개 posttrain 체크포인트로 미리 확인(icl 70% vs 논문 84.87%, 표준오차 범위 내)했기 때문에, 앞으로 나올 w/o pretrain 수치가 논문 ICL 36.36%와 차이가 나더라도 "평가 버그"와 "학습 재현 차이"를 분리해서 판단할 수 있는 기반이 마련됐다.

자동 checkpoint pruning으로 20시간 넘는 학습 동안 디스크 사용량을 약 44GB(22GB×2)+작성 중 1개로 제한 — 별도 개입 없이 무중단 완주에 기여했다.

5 Next Steps

미해결 한계

language-only(own, 7 GPU)와 ICL(extra, 7 GPU) 평가를 checkpoint_step_32000, 7 unseen task × 100 episode(seed 0)로 이제 막 제출한 상태라 아직 결과가 없다. 또한 논문 w/o pretrain의 정확한 학습 step 수·packing 여부가 불명(공개 코드는 batch 1, packing 없음) — 재현 수치가 논문 36.36%와 차이 나면 이 요인부터 의심해야 한다. 새로 추가된 head(action_embedder/action_proj_out)의 초기화 방식도 논문에 없어 재현 오차의 잠재 원인이다.

다음 실험

1. 평가 완료 대기 → icl/lang 결과 비교. 가설: ICL ≈ 36%(논문 근접), lang은 WAN-Action 10.98%와 LingBot-VA 17.45% 사이 — w/o pretrain은 로봇 사전학습이 없고 droptext 0.4로 language 의존이 약하게 학습되기 때문.

2. lang 수치가 WAN-Action(~11%)에 가까우면 "ICL 학습 이득은 전부 human video 조건에서 온다"는 가설 채택, 유의하게 높으면 "ICL 학습이 language 표현력도 개선" 가설로 전환해 추가 분석.

Zero-WAM w/o pretrain 재현 학습 완료 — RoboTwin 32K step