ImageWAM 레포 구조 분석 — Image-Editing 기반 World Action Model

TL;DR

구조: image-editing DiT(FLUX.2 klein 4B/9B, OmniGen2, Ovis-U1)를 video expert로, 같은 layer·head 수의 slim Action DiT를 붙여 layer마다 attention만 공유하는 MoT
비디오를 버림: 33프레임 / 9프레임 생성 대신 num_frames=17 + endpoint_frames_only → 프레임 [0, 16] 2장. 입력 t, 편집 타깃 t+16 (action chunk 끝 지점)
추론 시 이미지 denoising 루프 = 0회 — 타깃 이미지 latent 토큰 0개. 백본은 KV cache prefill로 1회만 forward. 논문 §3.4에 명시된 핵심 기여
뷰: LIBERO 2뷰(224×448 가로 concat), RoboTwin 3뷰(288×256 L자 타일) — 멀티뷰를 한 캔버스에 타일링해 image-edit 사전학습 재활용
불일치 3건 action denoising steps (논문 3 vs 코드 10), loss 가중 (논문 동일 vs 코드 0.5/1.0), τ⋆ (논문 미명시 vs 코드 0) — 재현 시 함정

Image denoise steps

3 / 10

Action steps 논문/코드

백본 forward / replan

2·3

뷰 (LIBERO·RoboTwin)

Action horizon

1 배경 / 목적

WAM(World Action Model) 계열은 보통 video generation 백본(Wan2.2 등)으로 미래 관측을 다프레임 롤아웃하고 그 표현으로 action을 뽑는다. 문제는 두 가지다.

기존 한계: (a) 비디오 디노이징 비용이 추론 지연을 지배한다. (b) 정말 "비디오"가 필요한지 검증된 적이 없다.

ImageWAM은 이 질문에 "image editing 한 장이면 충분하다"로 답한다. 레포를 뜯어서 실제 구조 / 뷰 개수 / 추론 시 denoising step 수를 코드로 확인했다.

2 MoT (Mixture-of-Transformers) 구조

두 expert가 layer마다 attention만 공유하고 weight는 완전히 분리된다.

Expert	실체	크기
Video (= Image-edit) Expert	사전학습 image-editing DiT 그대로	FLUX.2 klein-base 4B / 9B (권장), OmniGen2, Ovis-U1 (DiT 1.1B)
Action Expert	백본과 layer·head 수를 맞춘 slim DiT	`hidden_dim = 1024`

MoT.__init__ (mot.py:48-61)가 num_layers / num_heads / num_kv_heads / attn_head_dim 동일성을 강제 검증
from_flux2_klein_pretrained (imagewam.py:502-517)가 action expert config를 실제 FLUX.2 값으로 override → yaml의 num_layers_double: 5 / num_layers_single: 20은 placeholder (omnigen2 config엔 주석으로 명시됨)
각 layer: Attention(cat[Q_v,Q_a], cat[K_v,K_a], cat[V_v,V_a], mask) → split (_forward_flux2, mot.py:562-606)

요약: Action expert는 사실상 백본의 shadow 트랜스포머. 폭(hidden)만 1024로 줄이고 attention 차원은 백본과 동일하게 강제해 joint attention을 성립시킨다.

3 토큰 배치 & Attention Mask — 설계의 핵심

_build_mot_attention_mask_flux2 (imagewam.py:2072-2104):

[ text(Qwen3) | proprio 1tok | ref image tokens | target image tokens | action tokens ] txt_len +1 cond_len target_len action_len └────────── stable prefix ──────────┘ └─── noisy ───┘ └─── noisy ───┘

Query	볼 수 있는 것
text / ref image	prefix만 (noisy 토큰 차단)
target image	prefix + 자기 자신
action	prefix + 자기 자신 — target image는 못 봄 (2094-2095행)

왜 중요한가: action이 미래 이미지 토큰을 직접 참조하지 않도록 끊어놨기 때문에, 추론 시 이미지 브랜치를 통째로 잘라내도 action 예측이 그대로 유효하다. 학습/추론 비대칭을 mask 설계로 만들어낸 것.

4 뷰(View) — 멀티뷰를 한 장으로 타일링

별도 토큰 스트림이 아니라 한 캔버스에 붙여서 image-editing 모델에 그대로 먹인다 (robot_video_dataset.py:265-302).

벤치마크	뷰 수	카메라	합성	최종 해상도	image token
LIBERO / LIBERO-plus	2	`image`(agentview) + `wrist_image`, 각 224×224	가로 concat	224×448	14×28 = 392
RoboTwin 2.0	3	cam_high + cam_left_wrist + cam_right_wrist	L자 타일	288×256 (`compact`, eval 기본) 또는 384×320 (`legacy`)	18×16 = 288
InternData-A1 (pretrain)	3	cam_high / cam_left / cam_right	L자 타일	288×256	288

compact_288x256

head 192×256
+ L/R 96×128

legacy_384x320

head 256×320
+ L/R 128×160

FLUX.2 AE

16× down, C=128

주의: RoboTwin wrist 뷰는 96×128까지 강하게 다운샘플된다. 타일링의 해상도 예산 배분이 성능에 어떤 영향을 주는지는 논문/코드 어디에도 근거가 없다.

5 시간축 — 비디오 vs 이미지 편집

	비디오 baseline (`imagewam.yaml`, Wan2.2-TI2V-5B)	ImageWAM (image-edit)
`num_frames`	33	17
`action_video_freq_ratio`	4	1
이미지 프레임	9장 생성	`endpoint_frames_only: true`, `image_obs_steps: 2` → 인덱스 [0, 16] 2장
action	32 step	16 step

입력 = t 프레임, 편집 타깃 = t+16 프레임 — 즉 action chunk가 끝나는 시점의 관측. "16프레임 비디오 롤아웃"을 "16스텝 뒤 결과 한 장 편집"으로 치환한 것.

6 학습

Flow matching (rectified flow), continuous-time, shift=5.0, num_train_timesteps=1000 (scheduler_continuous.py)
video / action이 독립적인 timestep을 샘플링 (imagewam.py:2410-2426) — 이미지 노이즈 레벨과 action 노이즈 레벨이 무관
loss = 0.5·MSE(video velocity) + 1.0·MSE(action velocity)
timestep weighting: exp(-2((t-500)/1000)²) 정규화 (중간 timestep 강조)
학습 대상: self.dit = self.mot (imagewam.py:54) → trainer가 model.dit.requires_grad_(True) (trainer.py:477-479) → FLUX.2 DiT + Action DiT 둘 다 full fine-tune. VAE / 텍스트 인코더 freeze. LoRA 경로 존재하나 기본 off
텍스트는 Qwen3 임베딩 precompute 캐시. text_dim = 7680(4B, Qwen3-4B hidden×3) / 12288(9B, Qwen3-8B hidden×3)
proprio: nn.Linear(proprio_dim, text_dim) → 1 토큰, 유효 텍스트 토큰 직후 삽입 (imagewam.py:766-826)
prompt: "A video recorded from a robot's point of view executing the following instruction: {task}"
bf16 + DeepSpeed ZeRO-1/2, 8 GPU, lr 1e-4 cosine, wd 1e-2

7 추론 — image denoising loop = 0회 핵심

infer_action_flux2 (imagewam.py:3515-3630) 근거 라인:

라인	내용
3545-3546	`empty_target = ref_tokens.new_zeros(batch_size, 0, ...)` — 타깃 이미지 토큰 길이 0
3562	`video_timestep = torch.zeros(...)` — 이미지 timestep 고정 0
3563-3571	`video_expert.pre_dit(x=empty_target, ...)` → `target_len=0`
3575 / 3591	두 attention mask 모두 `target_len=0`
3581	`prefill_flux2_video_cache(...)` — 루프 바깥, 1회
3611-3628	루프 대상은 `infer_timesteps_action`뿐. `latents_action`만 step
3630	`return {"action": ...}` — 이미지 반환 없음, VAE decode 없음

정확한 표현: 이미지 flow-matching denoising 루프는 0회, 타깃 이미지 latent 토큰 0개. 단, video expert 트랜스포머 블록 자체는 1회 forward된다 — text + ref image 토큰에 대해 layer별 KV cache를 만드는 조건 인코딩 용도 (prefill_flux2_video_cache, mot.py:612-667).
"미래 이미지를 만들지 않는다"는 맞고, "백본을 안 돌린다"는 틀리다. 비용은 백본 1회 + action expert 10회.

항목	값
image edit denoising steps (추론)	0 — 이미지 브랜치는 학습 전용 auxiliary objective
action denoising steps	10 (`eval_num_inference_steps: 10`, `configs/train.yaml`) — 함수 시그니처 default는 20이나 sim config가 덮어씀
video expert forward	replan당 1회 (denoising step당 아님)

다른 백본도 동일 패턴 (교차 확인)

infer_action_ovis_u1 (3987-4011): empty_target = torch.empty(B, 0, ...), timestep_video = zeros
infer_action_omnigen2 (3850-3858): x=None, timestep_video = zeros

eval 진입점

deploy_policy.py:356 → self.model.infer_action(...)
eval_libero_single.py:568-573 → visualize_future_video가 False면 model.infer_action(...). sim config 기본값 false (sim_libero.yaml:29, sim_libero_omnigen2.yaml:29; RoboTwin config엔 키 자체가 없음)
디스패치: imagewam.py:3088-3100 stack == "flux2" → infer_action_flux2

예외: model.infer() dispatcher(4317-4345)를 직접 부르면 infer_video_flux2도 같이 돌아 미래 이미지를 실제로 N step 디노이징한다. 시각화 / PSNR 측정용이며 실제 eval 경로가 아니다.

8 Output

{"action": Tensor[T, D]} # normalized

	T	D	정규화	replan
LIBERO	16	7 = eef delta pose(6) + gripper(1)	min/max	16개 중 12개 실행 후 재계획
RoboTwin	16	14 = dual-arm joint (7+7)	z-score	16개 중 16개 전부
InternData-A1	16	16 (14D + 2 masked dim)	z-score	—

delta_action_dim_mask: pose 6차원은 delta, gripper는 absolute
후처리 — LIBERO: denorm → gripper 부호 복원 → optional binarize / RoboTwin: take_action(action, action_type="qpos")

9 변형 / 베이스라인 (`configs/model/`)

imagewam.yaml — Wan2.2-TI2V-5B 기반 비디오 생성 WAM baseline (9프레임)
imagewam_idm / imagewam_noise_idm / imagewam_cache_idm — inverse dynamics 계열 ablation
imagewam_joint, task의 *_uncond_* — joint denoising / 비디오 loss 없이 action만
image-edit 백본 3종: imagewam_flux2_klein_{4b,9b}_base, imagewam_omnigen2, imagewam_ovis_u1

10 Takeaway

"world model이 필요한 게 아니라, world model의 표현이 필요했다"

미래 관측 예측을 λ=0.5 auxiliary loss로만 쓰고, attention mask에서 action ↛ target image 경로를 명시적으로 끊어놨기 때문에 추론 시 이미지 생성을 0 step으로 절단할 수 있다. 결과적으로 4B DiT를 10번이 아니라 1번만 돌려서, image-editing 백본을 쓰면서도 실시간 제어가 성립한다.

우리 작업(EgoX / X-WAM)에 이식 가능한 아이디어

멀티뷰 타일링 — 별도 뷰 토큰 스트림 대신 한 캔버스에 붙여 image-edit 사전학습을 그대로 재활용
auxiliary 브랜치의 추론 시 절단 — mask 설계로 학습/추론 비대칭을 만드는 패턴. 학습 땐 신호로 쓰고 추론 땐 비용 0

11 논문 대조 완료

arXiv 2606.19531v1 (2026-06-17) 본문과 대조. "추론 시 이미지 denoising을 안 한다"는 논문의 핵심 기여로 명시되어 있다.

§3.4 Efficient Inference 원문

"At inference time, ImageWAM avoids full future-video generation and also does not require decoding a complete edited image. Instead of running the full image editing denoising trajectory, we select a fixed editing denoising timestep τ⋆ and perform only one editing-branch forward step to obtain 𝒞_edit^τ⋆ = f_edit^τ⋆(o_t, l). Action expert generates the action chunk by denoising action samples conditioned on this cache: â_t:t+H ∼ p_θ(a_t:t+H | o_t, l, 𝒞_edit^τ⋆)."

주목: 논문 표기 f_edit^τ⋆(o_t, l)의 인자가 현재 관측 + 언어뿐이고 노이즈 낀 타깃 latent가 없다. 코드의 target_len = 0과 정확히 일치 — "코드가 논문보다 한 발 더 나갔나" 의심했던 부분은 논문 수식이 이미 그렇게 정의돼 있었다.

대조표

항목	논문	코드
이미지 편집 denoising trajectory	"only one editing-branch forward step"	`prefill_flux2_video_cache` 루프 바깥 1회	✅
편집 이미지 decode	"does not require decoding a complete edited image"	VAE decode 없음	✅
cache 의존 인자	`f_edit^τ⋆(o_t, l)` — 타깃 latent 없음	`target_len = 0`	✅
LIBERO 뷰	2뷰 "horizontally concatenated ... 224×448"	동일	✅
RoboTwin 뷰	3뷰 "wrist-horizontal + vertical", 288×256	`compact_288x256`	✅
action chunk / 편집 타깃	16, "future observation 16 frames ahead"	16, target = t+16	✅
τ⋆ 값	미명시	0 (`imagewam.py:3562`)	⚠️
action denoising steps	3 (Table 11)	10 (`eval_num_inference_steps`)	❌
loss 가중	ℒ = ℒ_act + ℒ_img (동일)	λ_img=0.5, λ_act=1.0	❌
action ↛ target image mask	명시적 문장 없음	`imagewam.py:2094-2095` 명시적 차단	⚠️

불일치 3건 — 재현 시 함정

1) action denoising steps: 논문 3 vs 코드 10. Table 11 캡션 "All models use three action denoising steps during inference." 릴리스 configs/train.yaml 기본값은 10. 논문 latency(ImageWAM 263ms / 9.72 TFLOPs vs FastWAM-IDM 1081ms / 63.65 TFLOPs)는 3 step 기준이므로 기본 config로는 논문 latency가 재현되지 않는다. 성공률 표가 3 step인지 10 step인지도 불명확.

2) loss 가중: 논문 동일 vs 코드 0.5/1.0. 논문은 ℒ = ℒ_act + ℒ_img로만 표기하고 λ가 없다. 코드는 image 쪽을 0.5로 눌러놨다. 논문의 단순화 표기인지 실제 세팅 차이인지 판단 불가.

3) τ⋆ 값 미명시 — 재현성 이슈. 코드는 0. 무의미한 디테일이 아니다 — τ⋆는 adaLN modulation(vec = time_in(timestep_embedding(τ⋆)))을 거쳐 context 토큰 표현 전체를 바꾼다. flow-matching 컨벤션상 τ⋆=0은 σ=0, 즉 "타깃 이미지가 완전히 denoise된 상태인 것처럼 modulate"하는 선택. §3.4만 읽고 구현하는 사람은 이 값을 알 수 없다.

attention mask — 논문 서술이 부정확, 코드가 authoritative

논문: "action tokens attend to the other tokens in a one-way manner, while noisy tokens attend only to context tokens, thereby keeping the information in the context tokens clean."

이 문장만으로는 action이 노이즈 낀 타깃 이미지 토큰을 보는지 확정 불가 — "attend to the other tokens"는 오히려 본다고 읽힌다. 코드는 명시적으로 차단한다. §3.4가 성립하려면 차단이 필수: action이 타깃 토큰을 봤다면 추론 때 그 토큰을 빼는 순간 학습/추론 분포가 깨진다.

논문의 주력 백본은 OmniGen2

논문은 OmniGen2를 primary로 쓰고 Ovis-U1, FLUX.2 4B/9B를 함께 평가한다("Klein"이라는 이름은 논문에 없음). FLUX.2 9B가 LIBERO-Plus 평균 SR 83.1% → 85.21%로 최고. 반면 레포 README는 FLUX.2를 권장 경로로 밀고, InternData-A1 pretrain 체크포인트는 논문 이후(2026-07-30) 추가됐다. λ / step 수 차이 일부는 이 시차 때문일 수 있다.

12 Next

3 step vs 10 step 실측 — 논문 재현 시 성공률/latency 어느 쪽이 맞는지 확인
λ_img=0.5 vs 1.0 ablation — 논문 표기와 코드가 다른 이유 확인
τ⋆ ablation — 0이 최적인지, 논문이 왜 값을 안 밝혔는지
EVALUATION.timing_enabled=true로 백본 prefill(1회) vs action loop 비중 실측
third_party/flux2 clone 후 Klein4BParams로 action expert 파라미터 수 정확 산출
멀티뷰 타일링의 해상도 예산 배분(RoboTwin wrist 96×128) 성능 민감도 검증
환경 셋업(uv sync / FLUX.2 clone / 데이터 다운로드) 및 학습·평가 실행