fracapuano/lewm-tworooms-image-decoder

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 normalization
• config.json: architecture and input/output contract
• decoder.py: standalone architecture and from_pretrained loader
• training_metrics.json: complete epoch history, split sizes, metrics, and provenance
• reconstruction_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