Vision Components #

TorchWM provides a family of reusable vision modules — encoders, decoders, tokenizers, quantization layers, and distribution transforms — that serve as building blocks for world models and representation learning.

Overview #

All components live under world_models.vision and are accessible from the top-level package:

import torchwm

# Imports
from torchwm import ConvEncoder, ConvDecoder, VideoTokenizer
from torchwm import VectorQuantizer, VectorQuantizerEMA

Category	Component	Used by
Encoders	`ConvEncoder`	Dreamer (image → embedding)
	`CNNEncoder`	PlaNet (image → embedding)
	`IRISEncoder`	IRIS (image → discrete tokens)
Decoders	`ConvDecoder`	Dreamer (latent → image distribution)
	`CNNDecoder`	PlaNet (latent → image)
	`DenseDecoder`	Dreamer (latent → reward/value/discount)
	`ActionDecoder`	Dreamer (latent → action distribution)
	`IRISDecoder`	IRIS (tokens → image)
Video tokenization	`VideoTokenizer`	Genie (video → discrete tokens)
Quantization	`VectorQuantizer`	IRIS, Genie (embedding → codebook index)
	`VectorQuantizerEMA`	IRIS, Genie (EMA codebook updates)
Distributions	`TanhBijector`	Dreamer (action squashing)
	`SampleDist`	Dreamer (MC-sampled distribution stats)
	`_TwoHotDistribution`	DreamerV2 (symlog two-hot encoding)
VAE	`ConvVAE`	Standalone convolutional VAE

Encoders #

`ConvEncoder` — Dreamer convolutional encoder #

from torchwm import ConvEncoder

encoder = ConvEncoder(
    input_shape=(3, 64, 64),  # (C, H, W)
    embed_size=256,           # output dimension
    activation="elu",         # see below
    depth=32,                 # base channel count
)

obs = torch.randn(4, 3, 64, 64)
embedding = encoder(obs)     # (4, 256)

Architecture: 4 convolutional layers (kernel 4, stride 2) with channel doubling 32 → 64 → 128 → 256, then a linear projection to embed_size. Input values in [-0.5, 0.5].

`CNNEncoder` — PlaNet encoder #

from torchwm import CNNEncoder

encoder = CNNEncoder(embedding_size=256, activation_function="relu")

Same depth-doubling pattern as ConvEncoder but hardcoded to 3 input channels and no configurable depth parameter.

`IRISEncoder` — IRIS discrete encoder #

from torchwm import IRISEncoder

encoder = IRISEncoder(
    vocab_size=512,           # codebook size
    tokens_per_frame=16,      # 4×4 grid of tokens
    embedding_dim=512,
    in_channels=3,
    frame_shape=(3, 64, 64),
)

Architecture: 4 conv layers (stride 2, 64×64 → 4×4), self-attention at 16×16 and 8×8 resolutions, residual blocks, then a VectorQuantizerEMA produces discrete token indices. Input should be 64×64 images.

Decoders #

`ConvDecoder` — Dreamer convolutional decoder #

from torchwm import ConvDecoder

decoder = ConvDecoder(
    stoch_size=30,            # stochastic latent dimension
    deter_size=200,           # deterministic latent dimension
    output_shape=(3, 64, 64), # (C, H, W)
    activation="elu",
    depth=32,
)

features = torch.randn(4, 230)  # stoch + deter concatenated
dist = decoder(features)        # Independent(Normal(mean, 1), 3)
reconstruction = dist.mean      # (4, 3, 64, 64)
loss = -dist.log_prob(target)   # reconstruction loss

Architecture: linear projection from (stoch+deter) to 32×depth, then 4 transposed convolutions (stride 2, kernels [5, 5, 6, 6]). Returns a torch.distributions.Independent(Normal(mean, 1)) distribution so you can compute log_prob directly.

`DenseDecoder` — reward/value/discount head #

from torchwm import DenseDecoder

# Regression (reward, value)
decoder = DenseDecoder(
    stoch_size=30, deter_size=200,
    output_shape=(1,), n_layers=2, units=400,
    activation="elu", dist="normal",
)

# Binary classification (discount)
decoder = DenseDecoder(
    stoch_size=30, deter_size=200,
    output_shape=(1,), n_layers=2, units=400,
    activation="elu", dist="binary",
)

# Symlog two-hot (DreamerV2)
decoder = DenseDecoder(
    stoch_size=30, deter_size=200,
    output_shape=(1,), n_layers=2, units=400,
    activation="elu", dist="symlog_twohot",
    num_buckets=255, symlog_range=10.0,
)

`dist`	Return type	Use case
`"normal"`	`Independent(Normal)`	Reward prediction, value function
`"binary"`	`Independent(Bernoulli)`	Discount / termination prediction
`"symlog_twohot"`	`_TwoHotDistribution`	DreamerV2 reward/value
`"none"`	Raw tensor	Custom downstream processing

`ActionDecoder` — Dreamer policy head #

from torchwm import ActionDecoder

actor = ActionDecoder(
    action_size=6,
    stoch_size=30, deter_size=200,
    n_layers=2, units=400,
    activation="elu",
    min_std=1e-4, init_std=5, mean_scale=5,
)

features = torch.randn(4, 230)
action = actor(features)          # stochastic sample
action = actor(features, deter=True)  # deterministic mode

Outputs a Gaussian distribution squashed through TanhBijector to [-1, 1]. The deterministic mode (deter=True) returns the distribution mode for deployment; the stochastic mode is used during training.

Quantization #

`VectorQuantizer` and `VectorQuantizerEMA`#

from torchwm import VectorQuantizer, VectorQuantizerEMA

# Standard VQ (gradient-based codebook updates)
vq = VectorQuantizer(vocab_size=512, embedding_dim=64, commitment_weight=0.25)

# EMA VQ (more stable codebook learning)
vq = VectorQuantizerEMA(
    vocab_size=512, embedding_dim=64,
    commitment_weight=0.25, ema_decay=0.99,
)

z = torch.randn(4, 64, 8, 8)  # (B, C, H, W)
z_q, indices, loss_dict = vq(z)
#   z_q:      (4, 64, 8, 8)  quantized embeddings
#   indices:  (4, 8, 8)      codebook indices per spatial location
#   loss_dict: {"vq_loss": tensor, "perplexity": tensor}

Both layers implement the same interface. The EMA variant updates codebook vectors via exponential moving average of encoder outputs rather than gradient descent, which typically produces higher codebook utilization.

Video tokenization #

`VideoTokenizer` — Genie-style VQ-VAE #

from torchwm import VideoTokenizer

tokenizer = VideoTokenizer(
    num_frames=16,
    image_size=64,
    in_channels=3,
    encoder_dim=512, decoder_dim=1024,
    encoder_depth=12, decoder_depth=20,
    num_heads=16, patch_size=4,
    vocab_size=1024, embedding_dim=32,
    use_ema=True, ema_decay=0.99,
)

video = torch.randn(2, 3, 16, 64, 64)  # (B, C, T, H, W)
reconstructed, indices, loss = tokenizer(video)
#   reconstructed:  (2, 3, 16, 64, 64)
#   indices:        (2, 16, 16, 16)  (T × H' × W')
#   loss: {"recon_loss": ..., "vq_loss": ..., "perplexity": ...}

Architecture:

Patch embedding (B, C, T, H, W) → (B, T×P, encoder_dim)
Encoder ST-Transformer
Per-frame vector quantization
Decoder ST-Transformer
Patch unembedding → (B, C, T, H, W)

Key features: causal processing (each frame only uses previous frames), spatiotemporal transformer instead of full 3D ViT, per-frame VQ with shared codebook.

# Encode to discrete tokens
z_q, indices, vq_loss = tokenizer.encode(video)

# Decode from indices (for training downstream models)
embeddings = tokenizer.decode_indices(indices)  # (B, T, 16, 16, 32)
recon = tokenizer.decode(z_q)

# Factory shortcut
from world_models.vision.video_tokenizer import create_video_tokenizer
tokenizer = create_video_tokenizer(num_frames=16, image_size=64)

Distribution utilities #

`TanhBijector`#

Bijective tanh transform for squashing Gaussian actions to [-1, 1]. Used internally by ActionDecoder.

from torch.distributions import TransformedDistribution, Normal
from torchwm import TanhBijector

dist = TransformedDistribution(Normal(mean, std), TanhBijector())
action = dist.sample()  # bounded to [-1, 1]

`SampleDist`#

Wraps a distribution and approximates mean, mode, entropy via Monte Carlo sampling (100 samples by default). Used internally when the analytic form is unavailable (e.g., after tanh squashing).

`_TwoHotDistribution`#

DreamerV2’s symlog two-hot encoding for scalar prediction. Internally encodes targets into a categorical distribution over num_buckets evenly spaced bins within [-symlog_range, symlog_range], then decodes via symexp.

from world_models.vision.dreamer_decoder import _TwoHotDistribution

dist = _TwoHotDistribution(logits, num_buckets=255, symlog_range=10.0)
dist.log_prob(target)  # categorical cross-entropy in symlog space
dist.mean()            # expectation decoded via symexp

ConvVAE #

from world_models.vision.VAE.ConvVAE import ConvVAE

vae = ConvVAE(
    latent_dim=32,
    input_shape=(3, 64, 64),
)
recon, mu, logvar = vae(images)
loss = vae.loss_function(recon, images, mu, logvar)

A standalone convolutional VAE for representation learning and generative modeling. Can be used as a building block or baseline.

Which component to use #

Task	Encoder	Decoder
Dreamer world model	`ConvEncoder`	`ConvDecoder` + `DenseDecoder`
Dreamer policy	—	`ActionDecoder`
PlaNet world model	`CNNEncoder`	`CNNDecoder`
IRIS discrete AE	`IRISEncoder`	`IRISDecoder`
Genie video tokenizer	`VideoTokenizer`	(built-in)
Custom VQ-VAE	custom + `VectorQuantizer`	custom
Reward/value head	—	`DenseDecoder(dist="normal")`
Discount head	—	`DenseDecoder(dist="binary")`