Latent-Variable Representations › Code Implementation: A VAE Discovers the Pendulum's Angle
World Models HandbookAll chapters Latent-Variable Representations › Code Implementation: A VAE Discovers the Pendulum's AngleOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us make all of this concrete. We will build a static image dataset where the ground-truth factor for each snapshot is a single number, the pendulum angle , but the agent only sees 16 by 16 images of a glowing bob. This is a minimal laboratory for factor retention, geometry, topology, bottleneck strength, and collapse because the data has a known, one-dimensional, circular factor. It does not test temporal prediction or control sufficiency: there are no transitions, actions, rewards, or control objective, and a dynamical pendulum would generally also need angular velocity in its state. Then we will:
- Train a small VAE on the images and inspect its 2-dimensional latent space.
- Sweep the bottleneck strength and watch reconstruction, KL, and latent geometry respond.
- Train a tiny VQ-VAE on the same data, and check codebook usage.
The example uses CPU-compatible PyTorch operations. First, the hidden setup cell that the book's plotting conventions require.
The synthetic pendulum dataset
Each observation is a Gaussian blob placed at angle on a ring of radius 5 pixels around the image center. The only thing that varies across the dataset is , so the data manifold is intrinsically one-dimensional and topologically a circle. That makes it a sharp test of latent geometry: a topology-faithful 2-dimensional latent space should trace a ring. It is also a deliberately friendly test because the single factor of variation is visually dominant. Reduced data complexity makes failures easier to diagnose here, though optimization, architecture, sampling, and loss-scaling confounds remain.
import numpy as np
import torch
import torch.nn as nn
rng = np.random.default_rng(11)
def pendulum_image(theta, grid=16, radius=5.0, sigma=1.1):
cx = cy = (grid - 1) / 2
bx = cx + radius * np.sin(theta)
by = cy - radius * np.cos(theta)
ii, jj = np.mgrid[0:grid, 0:grid]
return np.exp(-((jj - bx) ** 2 + (ii - by) ** 2) / (2 * sigma**2))
n_train, n_test = 1200, 300
theta_train = rng.uniform(0, 2 * np.pi, n_train)
theta_test = rng.uniform(0, 2 * np.pi, n_test)
def make_set(thetas):
imgs = np.stack([pendulum_image(t) for t in thetas]).astype(np.float32)
return torch.tensor(imgs.reshape(len(thetas), -1))
x_train = make_set(theta_train)
x_test = make_set(theta_test)x_train shape: torch.Size([1200, 256]), x_test shape: torch.Size([300, 256])
The dataset is 1200 training and 300 test images, each flattened to a 256-dimensional vector. The angles are sampled uniformly, so no part of the ring is privileged; a dataset concentrated on one arc would leave the rest of the circle as extrapolation territory, which is exactly the hole problem from earlier in an intentional, controlled form. We keep the true angles around so we can grade the latents afterward, which in a real system corresponds to the (rare) luxury of a simulator that exposes the ground-truth factor. Closedness and gaps in the posterior-mean locus can be inspected without labels. A single traversal of hue is a supervised probe: it is available here only because the simulator supplies , and in a real system it requires ground-truth factor labels or a defensible measured proxy.
Here is what the agent sees:
probe_thetas = [0.0, np.pi / 2, np.pi, 3 * np.pi / 2]
samples = np.stack([pendulum_image(t).reshape(16, 16) for t in probe_thetas])
A glowing dot on a ring. No explicit angle label is supplied to the VAE; angle is only implicit in the blob's pixel location. The VAE must exploit the one-dimensional nonlinear correlation structure generated by motion around a closed curve in pixel space.
A small VAE in PyTorch
The model is deliberately tiny: an encoder and decoder with two hidden layers of 128 units and a 2-dimensional Gaussian latent. The physical snapshot has one intrinsic degree of freedom, while two Euclidean coordinates are the minimum for a seam-free embedding of a circle. The decoder outputs unconstrained reconstruction means. Under a unit-variance Gaussian likelihood, the negative log likelihood is one-half the summed squared error plus a parameter-independent constant, so the training loop uses that exact scale.
class VAE(nn.Module):
def __init__(self, dim_in=256, dim_h=128, dim_z=2):
super().__init__()
self.enc = nn.Sequential(
nn.Linear(dim_in, dim_h),
nn.SiLU(),
nn.Linear(dim_h, dim_h),
nn.SiLU(),
)
self.mu = nn.Linear(dim_h, dim_z)
self.logvar = nn.Linear(dim_h, dim_z)
self.dec = nn.Sequential(
nn.Linear(dim_z, dim_h),
nn.SiLU(),
nn.Linear(dim_h, dim_h),
nn.SiLU(),
nn.Linear(dim_h, dim_in),
)
def forward(self, x):
h = self.enc(x)
mu = self.mu(h)
logvar = self.logvar(h).clamp(-8.0, 8.0)
eps = torch.randn_like(mu)
z = mu + torch.exp(0.5 * logvar) * eps
return z, mu, logvar, self.dec(z)Two implementation details in the forward pass are worth connecting to the theory. The network predicts logvar rather than the variance directly, which keeps the variance positive by construction after exponentiation and is numerically better behaved; the clamp guards against the extreme values that can arise early in training. The sampling line z = mu + torch.exp(0.5 * logvar) * eps is the reparameterization trick written literally: the random draw eps is parameter-free, so gradients reach the encoder through mu and through logvar while the clamp is inactive. When the raw log-variance lies outside , clamp saturation intentionally blocks its variance-path gradient. Also note what the model returns: both the sampled and the posterior mean , because for visualization and evaluation we usually want the deterministic posterior mean, while the training objective needs the sampled latent.
The training loop minimizes a negated -ELBO: unit-variance Gaussian reconstruction loss plus times the closed-form Gaussian KL, one Adam step per minibatch.
def train_vae(beta, epochs=60, lr=1e-3, batch_size=128, seed=0):
torch.manual_seed(seed)
model = VAE()
opt = torch.optim.Adam(model.parameters(), lr=lr)
history = []
for epoch in range(epochs):
perm = torch.randperm(len(x_train))
for i in range(0, len(x_train), batch_size):
xb = x_train[perm[i : i + batch_size]]
z, mu, logvar, xrec = model(xb)
recon = 0.5 * ((xrec - xb) ** 2).sum(dim=1).mean()
kl = (-0.5 * (1 + logvar - mu**2 - logvar.exp())).sum(dim=1).mean()
loss = recon + beta * kl
opt.zero_grad()
loss.backward()
opt.step()
history.append((recon.item(), kl.item()))
return model, historyThe factor on the summed squared error matches a unit-variance Gaussian negative log likelihood up to an additive constant. The KL is summed over latent dimensions and averaged over the batch. This scaling matters: changing the reconstruction scale changes the effective strength of a fixed .
Now the experiment: train one model per and, on the held-out set, record the reconstruction error, the per-dimension KL, and the posterior means that we will visualize.
## The chosen range crosses from informative codes into near-complete collapse for this seed.
betas = [0.01, 0.025, 0.1, 0.5]
results = {}
for beta in betas:
model, history = train_vae(beta)
model.eval()
with torch.no_grad():
_, mu, logvar, xrec = model(x_test)
kl_dim = (-0.5 * (1 + logvar - mu**2 - logvar.exp())).mean(dim=0)
results[beta] = {
"model": model,
"history": history,
"mu": mu.numpy(),
"kl_dim": kl_dim.numpy(),
"recon_mse": ((xrec - x_test) ** 2).mean().item(),
}
recon_by_beta = np.array([results[b]["recon_mse"] for b in betas])
kl_by_beta = np.stack([results[b]["kl_dim"] for b in betas])Reading the latent geometry
The moment of truth. Because this controlled dataset retains the true angle as evaluation metadata, we can use hue ordering as a supervised probe of whether the VAE retains that factor. A topology-faithful two-dimensional encoding would form a closed curve without a seam, but angle retention alone does not imply that geometry. We compare an informative run at with the near-collapsed endpoint at .
At , hue follows an ordered one-dimensional curve, so the means retain angle. The curve is spiral-like rather than a seam-free circle, directly illustrating that reconstruction and a Gaussian prior do not guarantee topological fidelity. At , the means retain a tiny residual trace, but their mean radius is about 42 times smaller and total KL is about 0.003 nats. Samples from that near-prior posterior carry essentially no usable angle information, while reconstruction error rises sharply. Use measured KL, posterior scale, and latent sensitivity together; posterior-mean radius alone is not an information measure.
Sweeping the bottleneck
The per- summary quantities tell the quantitative side of the same story.
These are achieved operating points, not an empirical estimate of the optimal rate-distortion curve. At low , both coordinates carry substantial KL. That establishes posterior deviation from the prior, but by itself establishes neither observation-dependent usage nor topology because expected KL also contains aggregated-posterior mismatch. A seam-free Euclidean embedding of a circular factor needs at least two coordinates; this learned curve is not seam-free. At , both KL values approach zero while reconstruction error jumps: this seed crosses into near-complete posterior collapse rather than pruning an unused coordinate. Per-dimension KL is a primary warning signal, not a complete discriminator; pair it with reconstruction, latent-sensitivity or supervised probes, and downstream evaluation to learn whether and what each dimension encodes.
A tiny VQ-VAE on the same data
Now the discrete variant. The encoder is deterministic and its output is snapped to the nearest of 8 learned codes, with straight-through gradients. We initialize the codebook from encoder outputs of real data, placing every code initially in a region visited by the untrained encoder. Whether those codes remain active still has to be measured.
class VQVAE(nn.Module):
def __init__(
self, dim_in=256, dim_h=128, dim_z=2, n_codes=8, lambda_commit=0.25
):
super().__init__()
self.enc = nn.Sequential(
nn.Linear(dim_in, dim_h),
nn.SiLU(),
nn.Linear(dim_h, dim_h),
nn.SiLU(),
)
self.mu = nn.Linear(dim_h, dim_z)
self.codebook = nn.Embedding(n_codes, dim_z)
self.dec = nn.Sequential(
nn.Linear(dim_z, dim_h),
nn.SiLU(),
nn.Linear(dim_h, dim_h),
nn.SiLU(),
nn.Linear(dim_h, dim_in),
)
self.lambda_commit = lambda_commit
def forward(self, x):
ze = self.mu(self.enc(x))
dists = (
(ze.unsqueeze(1) - self.codebook.weight.unsqueeze(0)) ** 2
).sum(-1)
idx = dists.argmin(dim=1)
zq = self.codebook(idx)
zq_st = ze + (zq - ze).detach() # straight-through estimator
return self.dec(zq_st), idx, ze, zqThe structure mirrors the VAE closely, with the stochastic layer replaced by quantization. The distance computation broadcasts every encoder output against every codebook entry and takes the argmin, which is the quantization rule from the discrete-latents section in executable form. The straight-through line ze + (zq - ze).detach() is the expression verbatim: the forward value is , the backward gradient flows as if through alone. Training adds the codebook and commitment losses to reconstruction. Note the .detach() calls, which implement exactly the gradient routing described in the discrete-latents section: the codebook moves toward encoder outputs, the encoder is held to its committed codes, and the decoder's gradient flows straight through the quantization. Each .detach() selects which party a given loss term is allowed to train, and getting this routing right is the entire craft of implementing VQ-VAEs.
def train_vqvae(epochs=60, lr=1e-3, batch_size=128, seed=1):
torch.manual_seed(seed)
model = VQVAE()
with torch.no_grad():
init_z = model.mu(model.enc(x_train[:512]))
pick = torch.randperm(len(init_z))[: model.codebook.num_embeddings]
model.codebook.weight.copy_(init_z[pick])
opt = torch.optim.Adam(model.parameters(), lr=lr)
for epoch in range(epochs):
perm = torch.randperm(len(x_train))
for i in range(0, len(x_train), batch_size):
xb = x_train[perm[i : i + batch_size]]
xrec, idx, ze, zq = model(xb)
recon = ((xrec - xb) ** 2).sum(dim=1).mean()
codebook_loss = ((zq - ze.detach()) ** 2).sum(dim=1).mean()
commit_loss = (
model.lambda_commit
* ((zq.detach() - ze) ** 2).sum(dim=1).mean()
)
loss = recon + codebook_loss + commit_loss
opt.zero_grad()
loss.backward()
opt.step()
return model
vq_model = train_vqvae()
vq_model.eval()
with torch.no_grad():
vq_rec_test, idx_test, _, _ = vq_model(x_test)
vae_rec3 = results[0.025]["model"](x_test[:3])[3]
vq_rec3 = vq_rec_test[:3]
code_counts = np.bincount(
idx_test.numpy(), minlength=vq_model.codebook.num_embeddings
)The initialization block encodes real training data with the untrained encoder and copies a random subset of those outputs into the codebook. In this seed, all eight codes remain active on the held-out set. That is evidence for this run, not a multi-seed guarantee; repeat the usage check across seeds and training regimes.
The VQ-VAE emits one token from an eight-code vocabulary, at most three bits before accounting for the learned usage distribution, so it can preserve the angle only coarsely. The seeded counts are [80, 49, 18, 44, 20, 35, 24, 30]: uneven, but with no dead code. Unequal partition sizes and decoder preferences can produce this pattern without collapse. Prolonged zero use is a strong dead-code warning, not a mathematical point of no return; track perplexity and stability over training and reinitialize codes that remain unused.
The continuous VAE exposes posterior parameters that can be incorporated into a temporal belief model after calibration and consistency checks. The VQ-VAE exposes a discrete vocabulary suited to token dynamics. Neither static representation is automatically a belief state or decision-sufficient. The experiment shows both the benefit of a useful bottleneck and the failure boundary: the low- VAE and VQ-VAE retain the angle at different resolutions, while the high- VAE nearly abandons it.