blanchon/cs2-rae

CS2 RAE — DINOv3 representation autoencoder tokenizers (128px)

Frozen latent tokenizers for the OpenCS2 world model. Each turns a 128×128 game frame into a tiny C×8×8 latent (an 8×8 = 64-token grid); the world model lives entirely in that latent space.

• Architecture: DINOv3 ViT-S/16 encoder (frozen, LoRA rank 8) → learned variational bridge → ViT pixel decoder, trained natively at 128px with the full 3-stage recipe (stage1 → stage2 → stage3, LPIPS + late GAN). Latent grid is always 8×8 (latent_size=8); the variants differ only in latent_channels.
• Canonical recipe: the d64 subpixel tokenizer (64×8×8) — the production tokenizer for the world model.

Variants

File	Latent	latent_channels	Decoder head	Source run	Notes
rae_8x8x64.pt	64×8×8	64	subpixel (conv + PixelShuffle)	full_subpixel/s3/stage3.pt	Canonical / recommended. Production tokenizer.
rae_8x8x128.pt	128×8×8	128	non-subpixel (older head)	stage3_d128/stage3.pt	Higher-capacity latent. Earlier decoder generation — see caveat below.
stage3.pt	32×8×8	32	subpixel (conv + PixelShuffle)	full_d32_128/stage3.pt	Original upload (d32).

Each checkpoint is {"model": <state_dict>, "step": int, "metadata": {...}}; model is the trainable state (LoRA adapters + bridge + decoder) loaded onto the frozen DINOv3 backbone.

d128 caveat: rae_8x8x128.pt comes from an earlier run whose decoder is the non-subpixel head. It does not load against the current subpixel DecoderConfig default — load it with the matching (pre-subpixel) decoder config from the world-model codebase at that revision. d32 and d64 use the current subpixel decoder and load with the defaults below.

Usage

from diamond_modern.rae import RAE, RAEConfig
from diamond_modern.rae.config import BridgeConfig
import torch

# d64 (canonical) — subpixel decoder, LoRA rank 8 are the defaults
cfg = RAEConfig(bridge=BridgeConfig(latent_channels=64, latent_size=8))
rae = RAE.from_config(cfg, device="cuda")
payload = torch.load("rae_8x8x64.pt", map_location="cuda", weights_only=False)
rae.load_trainable_state_dict(payload["model"])
rae.eval()

latents = rae.encode_images(frames)   # frames: (N, 3, 128, 128) in [0, 1] -> (N, 64, 8, 8)
recon   = rae.decode_latents(latents) # -> (N, 3, 128, 128)

For stage3.pt (d32) use latent_channels=32. The LoRA rank must match the checkpoint (rank 8) or the encoder adapters silently skip-load.

Internal research checkpoints. Provided as-is, no warranty.