Nano World Model: Diffusion Forcing 기반의 미니멀한 비디오 월드 모델 학습 저장소

9bow · 5월 8, 2026, 3:30오전

Nano World Model 소개

월드 모델(world model)은 환경의 동역학을 학습하여, 어떤 행동을 취했을 때 미래 상태가 어떻게 변할지를 예측하는 모델을 가리킵니다. 최근에는 단순한 잠재 공간(latent space) 동역학 예측을 넘어 비디오 자체를 직접 예측하는 비디오 월드 모델(video world model)이 활발히 연구되고 있으며, 이는 로봇 조작, 자율주행 시뮬레이션, MPC(model-predictive control) 기반 계획, 영상 자료를 활용한 3D 재구성 등 다양한 응용으로 이어지고 있습니다. Nano World Model 은 이런 비디오 월드 모델 연구를 누구나 빠르게 시작할 수 있도록 하는 것을 목표로 한 미니멀한 PyTorch 기반 학습 저장소입니다.

이 프로젝트는 Max Simchowitz 가 이끄는 Simchowitz 연구실(simchowitzlabpublic) 에서 공개한 오픈소스 코드베이스이며, NanoGPT, NanoChat, Boyuan Chen 의 연구 템플릿에서 영감을 받아 "최소한의 의존성, 최대한의 투명성, 즉시 재현 가능한 결과"를 추구합니다. 클론 직후 몇 분 안에 첫 롤아웃(rollout)이 가능하도록 데이터 로딩과 환경 설정이 단순화되어 있고, 학습/검증/평가 파이프라인은 Hydra 기반 설정 시스템으로 통합되어 있습니다. 또한 prediction target, 행동 주입 방식(action injection), 모델 스케일에 대한 ablation 결과까지 모두 공개되어 있어 연구자가 디자인 선택의 영향을 직접 비교하기 쉽습니다.

기술적 토대는 디퓨전 포싱(diffusion forcing) 으로, 이는 시퀀스 모델링과 디퓨전 모델을 결합하여 프레임별 노이즈 스케줄을 다르게 적용함으로써 자기회귀(autoregressive) 비디오 생성과 일관된 장기 롤아웃을 동시에 달성하기 위한 기법입니다. Latte, Vid2World, DFoT, DINO-WM 같은 기존 코드베이스의 설계를 흡수하면서도, 단일한 학습 루프와 일관된 평가 절차를 제공하는 것이 Nano World Model 의 차별점입니다.

Nano World Model의 핵심 기능

이 저장소는 단순한 학습 스크립트 모음이 아니라, 다양한 도메인과 적용 사례를 한 묶음으로 다루는 통합 파이프라인입니다.

즉시 시작 가능한 학습 환경: conda 환경 파일과 Hydra 설정만으로 DINO-WM, RT-1(fractal), CSGO 같은 대표 도메인의 학습을 즉시 시작할 수 있도록 구성되어 있습니다. 데이터 경로는 환경 변수로 분리되어 있어 동일한 설정을 여러 머신에 배포하기 쉽습니다.

통합 파이프라인: 학습(training), 검증(validation), 평가(evaluation) 가 모두 Hydra 기반 설정 시스템으로 관리됩니다. 같은 설정 파일을 통해 같은 실험을 여러 도메인에서 동일한 절차로 재현할 수 있습니다.

과학적 투명성: 코드베이스가 깔끔하게 정리되어 있을 뿐 아니라, prediction target(예: pred-v 등), action injection 방식(additive 등), 모델 스케일(B/2, L/2 등)에 대한 head-to-head ablation 결과가 함께 제공됩니다. 모든 모델 체크포인트는 Hugging Face 컬렉션을 통해 공개되어 있어 학습 비용 없이 재현 실험이 가능합니다.

다양한 응용: 학습된 월드 모델로부터 long-horizon rollout(장기 자기회귀 생성), Depth Anything 3 기반의 video-to-3D 포인트 클라우드 재구성, MPC 스타일 계획(CEM 기반 옵티마이저)을 곧바로 수행할 수 있는 예제가 함께 제공됩니다.

Nano World Model 의 3x4 롤아웃 그리드 데모

Nano World Model의 디퓨전 포싱 학습 의사코드

이 저장소의 핵심은 디퓨전 포싱 기반의 비디오 예측입니다. 일반적인 비디오 디퓨전 모델이 모든 프레임에 동일한 노이즈 레벨을 적용하는 것과 달리, 디퓨전 포싱은 프레임 인덱스 t 마다 서로 다른 노이즈 레벨 \sigma_t 를 부여합니다. 이렇게 하면 학습 과정에서 자연스럽게 자기회귀 생성과 전체 시퀀스 디노이징의 경계를 매끄럽게 잇는 모델을 학습할 수 있고, 추론 시에는 frame-by-frame 자기회귀 디노이징(sequential scheduling)으로 일관된 장기 롤아웃이 가능합니다.

다음은 README 와 학습 스크립트를 바탕으로 정리한 PyTorch 의사코드 형태의 학습 루프 예시입니다. 실제 코드는 src/main.py 와 Hydra 설정으로 구성되어 있으며, 모델은 NanoWM-B/2, NanoWM-L/2 같은 변형을 제공합니다.

import torch
from torch.optim import AdamW

# 1) 비디오 월드 모델 (디퓨전 포싱 기반)
model = NanoWM(
    backbone="dit",        # diffusion transformer
    pred_target="v",       # pred-v 파라미터화 (cosine + ZTSNR)
    action_injection="additive",
    scale="B/2",           # 또는 "L/2" 등
).cuda()

optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.0)

for batch in train_loader:
    x = batch["frames"].cuda()      # (B, T, C, H, W) 비디오 시퀀스
    a = batch["actions"].cuda()     # (B, T, A) 액션 시퀀스 (도메인에 따라 None 가능)

    # 2) 프레임마다 다른 노이즈 레벨 σ_t 를 샘플링 (Diffusion Forcing 핵심)
    sigmas = sample_per_frame_sigmas(shape=(x.size(0), x.size(1)),
                                     schedule="cosine_ztsnr",
                                     device=x.device)
    noise = torch.randn_like(x)
    x_noisy = x + sigmas[..., None, None, None] * noise

    # 3) v-prediction 타겟 (cosine + ZTSNR)
    v_target = compute_v_target(x, noise, sigmas)
    v_pred = model(x_noisy, sigmas, action=a)

    loss = torch.nn.functional.mse_loss(v_pred, v_target)
    optimizer.zero_grad(); loss.backward(); optimizer.step()

추론 시에는 250 DDIM 스텝, sequential scheduling(frame-by-frame autoregressive denoising)을 통해 256 개의 고정 샘플(seed=42) 위에서 PSNR, SSIM, LPIPS, FID 를 측정합니다. README 의 정량 평가표에 따르면 DINO-WM Wall 에서 PSNR 34.05, SSIM 0.994, LPIPS 0.010 수준의 성능이 보고되어 있고, 더 어려운 RT-1(fractal) 도메인에서도 PSNR 24.36 수준의 결과를 얻었습니다.

Nano World Model의 응용 시나리오

학습된 월드 모델은 단순히 비디오를 생성하는 데 그치지 않고, 후속 응용에 그대로 활용될 수 있도록 도구가 함께 정리되어 있습니다.

장기 자기회귀 롤아웃: 학습된 체크포인트로부터 50 프레임 이상의 장기 시퀀스를 자기회귀적으로 생성합니다. README 에는 NanoWM-L/2 + CSGO 도메인에서 100k 스텝 학습된 모델의 50-frame 롤아웃 데모가 포함되어 있습니다.

비디오에서 3D 포인트 클라우드 재구성: 생성된 롤아웃 비디오를 입력으로 받아 Depth Anything 3 모델로 깊이를 예측하고, 이를 통해 시점 변화 가능한 3D 포인트 클라우드를 만들어 냅니다. 비디오 월드 모델이 단순한 픽셀 생성을 넘어 환경 표현으로 연결되는 방식을 보여줍니다.

MPC 스타일 계획: CEM(Cross-Entropy Method) 기반의 모델 예측 제어를 월드 모델 위에 얹어, 행동 후보를 샘플링하고 비용 함수를 평가하면서 계획을 수행합니다. 로봇 조작 도메인(DINO-WM PushT, RT-1)에서 곧바로 시도해볼 수 있도록 예제가 마련되어 있습니다.

Nano World Model 설치 및 사용법

# 1) 저장소 클론 및 conda 환경 생성
git clone https://github.com/simchowitzlabpublic/nano-world-model.git
cd nano-world-model
conda env create -f environment.yml && conda activate nanowm

# 2) 데이터 / 결과 경로 환경 변수 설정
export DATASET_DIR=/path/to/dino_wm_data
export CSGO_DATA_DIR=/path/to/csgo
export RT1_DATA_ROOT=/path/to/rt1_fractal
export RESULTS_DIR=/path/to/results

# 3) FID/FVD 평가에 사용되는 i3d torchscript 다운로드
mkdir -p pretrained_models/i3d && curl -L \
  "https://www.dropbox.com/scl/fi/c5nfs6c422nlpj880jbmh/i3d_torchscript.pt?rlkey=x5xcjsrz0818i4qxyoglp5bb8&dl=1" \
  -o pretrained_models/i3d/i3d_torchscript.pt

# 4) 첫 모델 학습 (DINO-WM PushT, NanoWM-B/2)
python src/main.py experiment=dino_wm_pusht dataset=dino_wm/pusht model=nanowm_b2

# 5) CSGO 도메인 학습 (NanoWM-L/2)
python src/main.py experiment=csgo dataset=game/csgo model=nanowm_l2_csgo

DINO-WM, RT-1, CSGO 데이터셋 다운로드 방법과 분할 방식은 docs/datasets/README.md 에 정리되어 있고, 각 모델 변형의 사전학습 체크포인트는 Hugging Face 의 knightnemo/nano-world-model 컬렉션에서 받을 수 있습니다.

라이선스

이 프로젝트는 MIT 라이선스로 공개되어 있어 학술 연구는 물론, 상업적 활용 시에도 라이선스 호환성 측면에서 부담이 적습니다. 다만 README 의 인용(citation) 섹션에 안내된 BibTeX 항목으로 저자(Siqiao Huang, Partha Kaushik, Michael Chen, Hengkai Pan, Kaiwen Geng, Omar Chehab, Fernando Moreno-Pino, Max Simchowitz)에 대한 인용을 포함하는 것이 권장됩니다.

Nano World Model 프로젝트 홈페이지

Nano World Model의 Hugging Face 모델 컬렉션

Nano World Model 프로젝트 GitHub 저장소

더 읽어보기

이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다.

파이토치 한국 사용자 모임이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

아래쪽에 좋아요를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~