code/world_model.py

"""Load the frozen LeWM world model for PushT.

The published `quentinll/lewm-pusht` checkpoint was saved with transformers 4.x
ViT parameter names. transformers >= 5 renamed them, so the state dict needs a
1:1 key remap before it will load (shapes are unchanged).
"""

import re

import torch
from hydra.utils import instantiate

from stable_worldmodel.data import get_cache_dir
from stable_worldmodel.wm.utils import _resolve

# transformers 4.x ViT name -> transformers 5.x name. Shapes are identical.
_VIT_RENAMES = (
    (r'^encoder\.encoder\.layer\.', 'encoder.layers.'),
    (r'\.attention\.attention\.query\.', '.attention.q_proj.'),
    (r'\.attention\.attention\.key\.', '.attention.k_proj.'),
    (r'\.attention\.attention\.value\.', '.attention.v_proj.'),
    (r'\.attention\.output\.dense\.', '.attention.o_proj.'),
    (r'\.intermediate\.dense\.', '.mlp.fc1.'),
    (r'(\.layers\.\d+)\.output\.dense\.', r'\1.mlp.fc2.'),
)


def _remap_vit_keys(state_dict: dict) -> dict:
    out = {}
    for key, value in state_dict.items():
        for pattern, repl in _VIT_RENAMES:
            key = re.sub(pattern, repl, key)
        out[key] = value
    return out


def load_lewm(
    name: str = 'quentinll/lewm-pusht',
    device: str = 'cuda',
    cache_dir: str | None = None,
):
    """Instantiate LeWM and load the pretrained weights, frozen and in eval."""
    cache_dir = get_cache_dir(cache_dir, sub_folder='checkpoints')
    ckpt_path, config = _resolve(name, cache_dir)

    model = instantiate(config)
    state_dict = torch.load(ckpt_path, map_location='cpu')
    missing, unexpected = model.load_state_dict(
        _remap_vit_keys(state_dict), strict=False
    )
    if missing or unexpected:
        raise RuntimeError(
            f'LeWM checkpoint mismatch after remap.\n'
            f'  missing: {sorted(missing)[:8]}\n'
            f'  unexpected: {sorted(unexpected)[:8]}'
        )

    model = model.to(device).eval()
    model.requires_grad_(False)
    return model