---
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.