LeWM TwoRooms latent-to-image decoder
This repository contains a PyTorch decoder from the frozen 192-dimensional
projected CLS embedding of
quentinll/lewm-tworooms
to a 224 x 224 RGB TwoRooms observation.
The decoder follows the learned patch-query and repeated cross-attention
architecture in
lucas-maes/57804ad5e347368e7a219719336c31de.
It has 2,667,264 trainable parameters, 3
cross-attention blocks, model width 256, and
8 attention heads.
Important training detail
The final checkpoint was trained with uniform RGB binary cross-entropy only. It uses no position labels, agent mask, red-pixel weighting, position head, or environment-specific reconstruction term. The agent-centroid result below is a held-out diagnostic and was not used for optimization or checkpoint selection.
Evaluation
The source sample contains 8,192 observations with an episode-disjoint split:
- train: 5,530 images / 3,766 episodes
- validation: 614 images / 423 episodes
- test: 2,048 images / 1,382 episodes
Final held-out test results:
| Metric | Value |
|---|---|
| RGB MSE | 0.00002485 |
| RGB MAE | 0.00074265 |
| PSNR | 46.05 dB |
| Agent-centroid RMSE | 0.407 px |
The reconstruction_preview.png file shows fixed held-out target/decoded pairs.
Usage
import sys
from huggingface_hub import snapshot_download
import torch
decoder_dir = snapshot_download("fracapuano/lewm-tworooms-image-decoder")
sys.path.insert(0, decoder_dir)
from decoder import CLSDecoder
decoder = CLSDecoder.from_pretrained(decoder_dir, device="mps")
# raw_embedding is the unnormalized encoded["emb"][:, 0] output from LeWM,
# with shape (batch, 192). Input normalization is stored in the checkpoint.
raw_embedding = torch.randn(1, 192, device="mps")
rgb = decoder.decode(raw_embedding) # (1, 3, 224, 224), values in [0, 1]
For an encoder batch containing a single frame:
with torch.inference_mode():
encoded = lewm_model.encode({"pixels": normalized_pixels[:, None]})
raw_embedding = encoded["emb"][:, 0]
reconstructed_rgb = decoder.decode(raw_embedding)
The LeWM encoder expects the ImageNet-normalized, resized preprocessing used by
stable_pretraining. The decoder output is ordinary unnormalized RGB in [0, 1].
Files
model.safetensors: decoder weights and stored latent normalizationconfig.json: architecture and input/output contractdecoder.py: standalone architecture andfrom_pretrainedloadertraining_metrics.json: complete epoch history, split sizes, metrics, and provenancereconstruction_preview.png: target/decoded held-out examples
Provenance
- Dataset revision:
a6dc0810c55cffb74500906735fa6e9ac323810c - Encoder revision:
77adaae0bc31deab21c93740d1f8bb947cd0bdec - Embedding SHA-256:
81a2a71d0deed97a219f66f8ff237c712b957c51e7b6143f33468a754e49286d - Source-row SHA-256:
a35811f60bf9d322b6afd623fab4e6280de0bb0676b1d777d7083cbf2bd97627 - Training seed:
1729 - Training hardware: Apple M1 Pro via PyTorch MPS