File size: 6,275 Bytes
242cc21 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """Train a small MIRA latent world model on the CS2 codec, and test action-controllability.
train: python scripts/train_cs2_wm.py
test : python scripts/train_cs2_wm.py --test (seed real context, feed fixed actions, generate)
Frozen codec is loaded from the consolidated checkpoint; a tiny flow-matching diffusion transformer
learns to predict the next codec latent conditioned on CS2 keyboard+mouse actions.
"""
import os, sys, time, argparse
os.environ.setdefault("RS_DINO_HF", "facebook/dinov3-vitl16-pretrain-lvd1689m")
sys.path.insert(0, "src")
import torch
from mira.world_model.config import LatentWorldModelConfig
from mira.world_model.latent_world_model import LatentWorldModel
from mira.world_model.actions_config import ActionConfig
from mira.ml.image_config import ImageConfig
from mira.data.cs2_stream import create_cs2_loader, CS2_KEYS
dev = "cuda" if torch.cuda.is_available() else "cpu"
T = int(os.environ.get("CS2_T", "16"))
CODEC = os.environ.get("CS2_CODEC", "runs/cs2_codec_consolidated/codec.pth")
OUT = os.environ.get("CS2_WM_OUT", "runs/cs2_wm")
def build_model():
actions = ActionConfig(valid_keys=list(CS2_KEYS), source_fps=24, target_fps=24)
video = ImageConfig(height=288, width=512, channels=3, timesteps=T, fps=24)
cfg = LatentWorldModelConfig(
actions=actions, video=video, codec_checkpoint=CODEC,
latent_mean_std=[0.0, 1.0], # identity norm (we didn't estimate latent stats)
causal=True, use_clean_past=True, learned_temporal_pool=True,
use_codec_posterior_mean=True, attention_gating=True, ada_attn_ln=True,
patch_size=1, n_register_tokens=0, n_context_frames=max(2, T // 2 - 2),
dropout_action_prob=0.1,
# size from env (defaults = tiny/16GB-safe; scale up on big GPUs, e.g. 1B: 2048/16/4/16)
hidden_dim=int(os.environ.get("WM_HIDDEN", "384")),
n_head=int(os.environ.get("WM_HEADS", "6")),
n_kv_head=int(os.environ.get("WM_KV", "2")),
n_layers=int(os.environ.get("WM_LAYERS", "4")),
time_attention_every=int(os.environ.get("WM_TATT", "2")),
activation_checkpointing=os.environ.get("WM_ACT_CKPT", "0") == "1",
)
return LatentWorldModel(cfg).to(dev)
def train():
from pathlib import Path
model = build_model().train()
n_tr = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6
print(f"[wm] trainable params: {n_tr:.1f}M (+ frozen codec)", flush=True)
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-4)
loader = create_cs2_loader(subset=os.environ.get("CS2_SUBSET", "sample"), n_players=1, clip_len=T,
target_fps=24, frame_size=(288, 512),
batch_size=int(os.environ.get("WM_BATCH", "1")),
num_workers=int(os.environ.get("CS2_WORKERS", "0")), infinite=True)
ckpt_dir = Path(OUT); ckpt_dir.mkdir(parents=True, exist_ok=True)
STEPS = int(os.environ.get("CS2_WM_STEPS", "5000"))
step, ema, t0 = 0, None, time.time()
for batch, meta in loader:
losses = model(batch.to(dev))
total = losses["loss_total"]
opt.zero_grad(); total.backward(); opt.step()
ema = float(total) if ema is None else 0.98 * ema + 0.02 * float(total)
if step % 20 == 0:
print(f"[wm] step {step:5d} loss={float(total):.4f} ema={ema:.4f} "
f"({(time.time()-t0)/(step+1):.2f}s/step)", flush=True)
step += 1
if step % 500 == 0:
# exclude the frozen codec (rebuilt from CS2_CODEC) -> ~50MB instead of ~1.3GB
wm_only = {k: v for k, v in model.state_dict().items() if not k.startswith("codec.")}
torch.save({"step": step, "model": wm_only, "ema": ema}, ckpt_dir / f"wm_{step:06d}.pt")
for old in sorted(ckpt_dir.glob("wm_0*.pt"))[:-2]:
old.unlink(missing_ok=True)
print(f"[wm] saved wm_{step:06d}.pt (ema={ema:.4f})", flush=True)
if step >= STEPS:
break
print("[wm] done", flush=True)
@torch.no_grad()
def test():
"""Seed a real context clip, then generate under three fixed action regimes and save frames."""
import glob, numpy as np
from PIL import Image
from mira.world_model.latent_world_model import WorldModelInferenceConfig
model = build_model().eval()
ck = os.environ.get("CS2_WM_CKPT") or (sorted(glob.glob(f"{OUT}/wm_0*.pt")) or [None])[-1]
if ck:
model.load_state_dict(torch.load(ck, map_location=dev, weights_only=False)["model"], strict=False)
print(f"[wm-test] loaded {ck}", flush=True)
loader = create_cs2_loader(subset="sample", n_players=1, clip_len=T, target_fps=24,
frame_size=(288, 512), batch_size=1, num_workers=0, infinite=True)
batch, _ = next(iter(loader))
# key indices in the CS2 vocab
kW, kMOUSE = CS2_KEYS.index("W"), None
regimes = {
"idle": lambda a: (a.key_presses.zero_(), a.mouse_movements.zero_()),
"forward": lambda a: (a.key_presses.zero_().__setitem__((slice(None), slice(None), kW), 1),
a.mouse_movements.zero_()),
"turn_right": lambda a: (a.key_presses.zero_(),
a.mouse_movements.zero_().__setitem__((slice(None), slice(None), 0), 25.0)),
}
icfg = WorldModelInferenceConfig() if False else None
outdir = os.path.join(OUT, "control"); os.makedirs(outdir, exist_ok=True)
for name, override in regimes.items():
b = batch.clone()
override(b.actions)
out = model.inference(b.to(dev), progress_bar=False)
vid = out.output_video[0] # (T,C,H,W) in [-1,1]
img = ((vid[-1] * 0.5 + 0.5).clamp(0, 1) * 255).byte().permute(1, 2, 0).cpu().numpy()
p = os.path.join(outdir, f"gen_{name}_lastframe.png")
Image.fromarray(img).save(p)
print(f"[wm-test] {name}: saved {p}", flush=True)
print("[wm-test] compare the three frames: forward/turn should differ from idle if controllable.", flush=True)
if __name__ == "__main__":
ap = argparse.ArgumentParser(); ap.add_argument("--test", action="store_true"); a = ap.parse_args()
(test if a.test else train)()
|