| --- |
| library_name: pytorch |
| tags: |
| - world-model |
| - tokenizer |
| - autoencoder |
| - dinov3 |
| - counter-strike |
| - opencs2 |
| --- |
| |
| # CS2 RAE — DINOv3 representation autoencoder tokenizers (128px) |
|
|
| Frozen latent tokenizers for the [OpenCS2 world model](https://github.com/julien-blanchon/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 |
| |
| ```python |
| 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. |
| |
