Upload README.md with huggingface_hub

README.md

@@ -0,0 +1,114 @@
+---
+license: mit
+tags:
+  - world-model
+  - jepa
+  - dino-wm
+  - robotics
+  - pusht
+library_name: stable-worldmodel
+---
+
+# DINO-WM (PreJEPA) — PushT · patch tokens + proprio
+
+프레임을 얼린 **DINOv2-small** 백본으로 latent 인코딩하고, causal predictor 로
+다음 latent 을 예측하는 world model (DINO-WM 계열, JEPA loss). 픽셀 재구성 없음.
+
+- backbone: `dinov2_small` (frozen), `pixel_token=patch` → 프레임당 256 패치 × 384-d
+- predictor: `CausalPredictor`, `dim=404` (= pixel 384 + proprio_emb 10 + action_emb 10)
+- `history_size=3`, `num_pred=1`, `frameskip=5`
+- 부가입력: `proprio`(in_chans=4: agent pos+vel), `action`(in_chans=10 = raw 2 × frameskip 5)
+- env: `swm/PushT-v1`
+
+## 파일
+
+| 파일 | 설명 |
+|---|---|
+| `weights.pt` | 모델 가중치 (epoch 10) |
+| `config.json` | 구조 (hydra instantiate 용) |
+| `norm_stats.json` | proprio/action ZScore mean·std (eval 정규화 복원) |
+
+## 설치
+
+```bash
+pip install stable-worldmodel   # 또는 저장소에서 editable 설치
+```
+
+## 로드 (public repo → 내장 로더)
+
+```python
+import stable_worldmodel as swm
+
+model = swm.wm.utils.load_pretrained("kotmul/dinowm_patch_prop_pusht")
+model = model.eval().requires_grad_(False)
+model.interpolate_pos_encoding = True
+```
+
+`load_pretrained` 는 `config.json` + `weights.pt` 를 `<cache>/checkpoints/` 아래로 받아
+`instantiate(config)` 후 가중치를 로드한다.
+
+## 정규화 (중요)
+
+- **pixels**: ImageNet mean/std 정규화 후 224×224
+- **proprio / action**: 아래 `norm_stats.json` 의 ZScore (학습과 반드시 동일해야 함)
+
+```python
+import json, numpy as np
+from huggingface_hub import hf_hub_download
+
+norm = json.load(open(hf_hub_download("kotmul/dinowm_patch_prop_pusht", "norm_stats.json")))
+p_mean, p_std = np.array(norm["proprio"]["mean"][0]), np.array(norm["proprio"]["std"][0])
+a_mean, a_std = np.array(norm["action"]["mean"][0]),  np.array(norm["action"]["std"][0])
+```
+
+## 추론 — 프레임 인코딩 & 다음 스텝 예측
+
+```python
+import torch, numpy as np
+import stable_pretraining as spt
+from torchvision.transforms import v2 as T
+
+tf = T.Compose([
+    T.ToImage(), T.ToDtype(torch.float32, scale=True),
+    T.Normalize(**spt.data.dataset_stats.ImageNet), T.Resize(224),
+])
+H, FS = model.history_size, 5   # 3 history steps, frameskip 5
+
+# frames_uint8: (H, 224, 224, 3) uint8  — history_size 개의 연속 프레임(frameskip 간격)
+# proprio_raw : (H, 4)   각 스텝의 [agent_x, agent_y, agent_vx, agent_vy]
+# action_raw  : (H, FS*2) 각 model-step 의 raw action FS개 묶음 ([-1,1]^2 × FS)
+
+pixels  = torch.stack([tf(im) for im in frames_uint8])[None]                 # (1,H,3,224,224)
+proprio = torch.tensor(((proprio_raw - p_mean) / p_std)[None], dtype=torch.float32)  # (1,H,4)
+action  = ((action_raw.reshape(H, FS, 2) - a_mean) / a_std).reshape(H, FS * 2)
+action  = torch.tensor(action[None], dtype=torch.float32)                    # (1,H,10)
+
+with torch.no_grad():
+    # (a) 단일 프레임 인코딩 (patch latent)
+    emb_img = model._encode_image(pixels[:, :1])   # (1, 1, 256, 384)
+
+    # (b) 다음 스텝 예측 (action/proprio 반영)
+    info = {"pixels": pixels, "proprio": proprio, "action": action}
+    info = model.encode(info, target="emb", is_video=False)
+    pred = model.predict(info["emb"][:, :H])       # (1, H, 256, 404)
+    next_latent = pred[:, -1]                       # 예측한 다음 latent (1, 256, 404)
+    # 404 = pixel(384) + proprio_emb(10) + action_emb(10).
+    # planning cost 등에는 보통 action 구간(마지막 10)을 제외한 actionless 부분 사용:
+    actionless = next_latent[..., :394]
+```
+
+## Planning / eval 에서 쓰기
+
+`stable-worldmodel` 의 planning(eval) 스크립트는 체크포인트 옆의 `norm_stats.json` 을
+자동으로 찾아 학습 때 정규화를 복원한다(option B). 따라서 세 파일을
+`<cache>/checkpoints/<run>/` 한 폴더에 두고 `policy` 를 그 `weights.pt` 로 지정하면 된다:
+
+```
+<cache>/checkpoints/dinowm-pusht-patch-prop/
+    weights.pt
+    config.json
+    norm_stats.json      # eval 이 여기서 mean/std 복원
+```
+
+MPC(CEM) planning 은 world model 을 imagination 으로 굴려 cost 를 최소화하고,
+실제 env 에서 실행한다. 자세한 진입점은 리포지토리의 planning 스크립트를 참고.