Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Predictive Embeddings, JEPA, and Energy-Based Models › Worked Example: Latent Prediction for a Partially Observed Linear System

Michael BrenndoerferJune 29, 202617 min read

Part 5 of 12World-Model Architectures

Chapter 29Predictive Embeddings, JEPA, and Energy-Based Models › Worked Example: Latent Prediction for a Partially Observed Linear System

Reading section 6 of 10 · Worked Example: Latent Prediction for a Partially Observed Linear System · 17 minSection 6 of 10 · 17 min

Section 6 of 10 · 17 min read

Predictive Embeddings, JEPA, and Energy-Based Models › Worked Example: Latent Prediction for a Partially Observed Linear SystemOne section · Sections in this chapter6 / 10 ↓

Sections in this chapter10

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

We will use a small system where the true state is recorded exactly. The environment is a two-dimensional linear dynamical system with a rotating, contracting transition matrix. The agent observes a 32-dimensional nonlinear function of the state plus Gaussian noise. The observation is much bigger than the state, and the encoder must decide what to keep. Because the probe targets are ground-truth simulated states, held-out $R^2$ measures their linear decodability under a specified ridge-regression protocol. It does not measure all retained information or control usefulness. The experiment and evaluation cells compute the arrays used by the later presentation-only figure cells.

In[3]:

Code

import numpy as np
import torch
import torch.nn as nn

SEED = 0
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(1)

D_OBS, D_STATE, D_LATENT = 32, 2, 8

# A stable, rotating 2-D linear system
A = np.array([[0.9, 0.15], [-0.15, 0.9]])
B = np.array([0.3, 0.1])

# Fixed random observation map from state to a 32-dimensional measurement
W_OBS = np.random.default_rng(SEED).normal(size=(D_OBS, D_STATE)) * 0.6

The dynamics, observation map, and state labels are fixed and known, which lets us measure held-out linear decodability from each representation. Now let us generate trajectories.

In[4]:

Code

def observe(s, rng):
    """Nonlinear, noisy, over-complete measurement of a 2-D state.

    The additive observation noise is large enough that the informative
    directions in the measurement are not a perfect read-out of the state,
    so the raw-observation probe is a realistic reference point rather than
    a saturated R^2 = 1 baseline."""
    return np.tanh(W_OBS @ s) + rng.normal(size=D_OBS) * 0.35


def rollout(rng, T):
    """Roll out the controlled linear system and record (state, obs, action)."""
    s = rng.normal(size=D_STATE) * 0.5
    S, X, U = [], [], []
    for _ in range(T):
        u = float(rng.normal() * 0.5)
        S.append(s.copy())
        X.append(observe(s, rng))
        U.append(u)
        s = A @ s + B * u + rng.normal(size=D_STATE) * 0.02
    return np.array(S), np.array(X), np.array(U)

Each trajectory yields transitions $(o_t, a_t, o_{t+1}, s_{t+1})$ . The code pools transitions within each trajectory and then splits complete trajectories into train and test sets, so all transitions from a held-out trajectory remain together. This avoids leakage across time between adjacent transitions in the same trajectory.

In[5]:

Code

rng_data = np.random.default_rng(123)
T_LEN, N_TRAJ = 60, 40
trajs = [rollout(rng_data, T_LEN) for _ in range(N_TRAJ)]

# Each trajectory is a tuple (S[T_LEN, D_STATE], X[T_LEN, D_OBS], U[T_LEN]).
# Split complete trajectories, not individual transitions, to avoid leakage across time.
n_tr_traj = int(0.8 * N_TRAJ)
train_trajs = trajs[:n_tr_traj]
test_trajs = trajs[n_tr_traj:]


def stack_trajs(trajs):
    ctx = np.concatenate([X[:-1] for _, X, _ in trajs])
    nxt = np.concatenate([X[1:] for _, X, _ in trajs])
    act = np.concatenate([U[:-1, None] for _, _, U in trajs])
    state_n = np.concatenate([S[1:] for S, _, _ in trajs])  # state at t+1
    return ctx, nxt, act, state_n


ctx_tr, nxt_tr, act_tr, state_n_tr = stack_trajs(train_trajs)
ctx_te, nxt_te, act_te, state_n_te = stack_trajs(test_trajs)


def t(x):
    return torch.tensor(x, dtype=torch.float32)


Xc_tr, U_tr, Xn_tr, Sn_tr = t(ctx_tr), t(act_tr), t(nxt_tr), t(state_n_tr)
Xc_te, U_te, Xn_te, Sn_te = t(ctx_te), t(act_te), t(nxt_te), t(state_n_te)

Out[6]:

Console

train transitions: 1888, test transitions: 472

The printed counts report the exact number of training and held-out transitions. The two-dimensional target makes the linear probe simple and directly interpretable, although the result still depends on the split, probe family, regularization, and training seed.

Next, the two network pieces. The encoder is an MLP with a linear output layer, and the predictor takes the latent plus the action and outputs a predicted next latent.

In[7]:

Code

class Encoder(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, d_out),
        )

    def forward(self, x):
        return self.net(x)


class Predictor(nn.Module):
    def __init__(self, d_latent, d_cond):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_latent + d_cond, 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, d_latent),
        )

    def forward(self, z, c):
        return self.net(torch.cat([z, c], dim=-1))

The training function below makes the target encoder, predictor, and variance hinge optional. Setting use_target=False, var_weight=0.0, and use_predictor=False gives the fully coupled symmetric objective. Enabling the EMA target, predictor, and a positive variance weight gives the protected JEPA-style variant. The function keeps the data, encoder, optimizer, and training loop shared while changing several mechanisms, so the comparison is descriptive rather than a one-factor ablation.

In[8]:

Code

class IdentityPredictor(nn.Module):
    """No-op predictor: returns the context embedding unchanged.

    Using this for the coupled variant removes the predictor asymmetry as well as
    the stop-gradient. Constant encodings are zero-loss stationary solutions, and
    the seeded run below converges toward one."""

    def forward(self, z, c):
        return z


def train_latent_model(
    Xc,
    U,
    Xn,
    use_target=True,
    var_weight=0.0,
    ema=0.99,
    epochs=1500,
    lr=1e-3,
    seed=0,
    use_predictor=True,
):
    torch.manual_seed(seed)
    online = Encoder(D_OBS, D_LATENT)
    pred = Predictor(D_LATENT, 1) if use_predictor else IdentityPredictor()
    target = None
    if use_target:
        target = Encoder(D_OBS, D_LATENT)
        target.load_state_dict(online.state_dict())
        for p in target.parameters():
            p.requires_grad_(False)

    opt = torch.optim.Adam(
        list(online.parameters()) + list(pred.parameters()), lr=lr
    )

    for _ in range(epochs):
        opt.zero_grad()
        z_hat = pred(online(Xc), U)
        if use_target:
            with torch.no_grad():
                y = target(Xn)
        else:
            y = online(Xn)  # gradients flow into the encoder, so no asymmetry
        loss = ((z_hat - y) ** 2).mean()

        if var_weight > 0:  # VICReg-style hinge on per-dimension std
            std = torch.sqrt(online(Xc).var(dim=0) + 1e-6)
            loss = loss + var_weight * torch.relu(1.0 - std).mean()

        loss.backward()
        opt.step()

        if use_target:  # exponential moving average of target weights
            with torch.no_grad():
                for pt, po in zip(target.parameters(), online.parameters()):
                    pt.mul_(ema).add_(po, alpha=1 - ema)

    return online, pred, target

The y = online(Xn) branch is the critical line. When use_target is false, the target embedding is produced by the same optimized weights, so a constant encoding is a zero-loss solution. When use_target is true, the target is detached and lags behind, preventing instantaneous gradient cooperation between branches. That change alone does not make collapse impossible; the full variant also uses a predictor and variance hinge.

Let us train four variants: a pixel-space predictor, a fully coupled latent model with no protections, a no-target latent model with a predictor and variance hinge, and the full JEPA-style recipe.

In[9]:

Code

# Pixel-space baseline: predict o_{t+1} directly from (o_t, a_t)
torch.manual_seed(0)
pixel_model = nn.Sequential(
    nn.Linear(D_OBS + 1, 64),
    nn.ReLU(),
    nn.Linear(64, 64),
    nn.ReLU(),
    nn.Linear(64, D_OBS),
)
pixel_opt = torch.optim.Adam(pixel_model.parameters(), lr=1e-3)
for _ in range(1500):
    pixel_opt.zero_grad()
    out = pixel_model(torch.cat([Xc_tr, U_tr], dim=-1))
    loss = ((out - Xn_tr) ** 2).mean()
    loss.backward()
    pixel_opt.step()

# Latent variants
# The coupled variant is fully symmetric: no stop-gradient and no predictor.
# Constant encodings are zero-loss solutions; this seeded run converges near one.
enc_collapse, pred_collapse, _ = train_latent_model(
    Xc_tr, U_tr, Xn_tr, use_target=False, var_weight=0.0, use_predictor=False
)

enc_var, pred_var, _ = train_latent_model(
    Xc_tr, U_tr, Xn_tr, use_target=False, var_weight=0.1
)

enc_jepa, pred_jepa, tgt_jepa = train_latent_model(
    Xc_tr, U_tr, Xn_tr, use_target=True, var_weight=1.0, ema=0.99
)

Out[10]:

Console

training complete: 4 variants trained

Now we evaluate. Fit ridge regression from each representation to the true next state $s_{t+1}$ , then report mean $R^2$ on held-out transitions. A high value means that this chosen state is accessible to a regularized linear readout under the trajectory split. A poor value does not rule out nonlinear decodability, and a high value does not establish control sufficiency.

In[11]:

Code

def probe_r2(feat_tr, target_tr, feat_te, target_te, lam=0.1):
    """Ridge regression probe: how linearly decodable is the true state?"""
    Ftr = feat_tr.detach().numpy().astype(np.float64)
    Fte = feat_te.detach().numpy().astype(np.float64)
    Ttr = target_tr.numpy().astype(np.float64)
    Tte = target_te.numpy().astype(np.float64)

    mu, sd = Ftr.mean(0), Ftr.std(0) + 1e-8
    Ftr, Fte = (Ftr - mu) / sd, (Fte - mu) / sd
    Ftr = np.c_[Ftr, np.ones(len(Ftr))]
    Fte = np.c_[Fte, np.ones(len(Fte))]

    W = np.linalg.solve(Ftr.T @ Ftr + lam * np.eye(Ftr.shape[1]), Ftr.T @ Ttr)
    P = Fte @ W
    ss_res = ((Tte - P) ** 2).sum(0)
    ss_tot = ((Tte - Tte.mean(0)) ** 2).sum(0)
    return float(np.mean(1 - ss_res / np.maximum(ss_tot, 1e-12)))

In[12]:

Code

with torch.no_grad():
    pixel_next = pixel_model(torch.cat([Xc_tr, U_tr], dim=-1))
    pixel_next_te = pixel_model(torch.cat([Xc_te, U_te], dim=-1))

    rep_tr = {
        "Raw observation": Xn_tr,
        "Pixel predictor": pixel_next,
        "Coupled latent": pred_collapse(enc_collapse(Xc_tr), U_tr),
        "No target + hinge": pred_var(enc_var(Xc_tr), U_tr),
        "JEPA latent": pred_jepa(enc_jepa(Xc_tr), U_tr),
    }
    rep_te = {
        "Raw observation": Xn_te,
        "Pixel predictor": pixel_next_te,
        "Coupled latent": pred_collapse(enc_collapse(Xc_te), U_te),
        "No target + hinge": pred_var(enc_var(Xc_te), U_te),
        "JEPA latent": pred_jepa(enc_jepa(Xc_te), U_te),
    }

r2_results = {k: probe_r2(rep_tr[k], Sn_tr, rep_te[k], Sn_te) for k in rep_tr}


def effective_rank(feat):
    """Entropy effective rank of a centered embedding matrix."""
    centered = feat - feat.mean(dim=0, keepdim=True)
    singular_values = torch.linalg.svdvals(centered)
    eps = torch.finfo(singular_values.dtype).eps
    total = singular_values.sum()
    if total <= eps:
        return 0.0
    weights = singular_values / total
    entropy = -(weights * weights.clamp_min(eps).log()).sum()
    return float(entropy.exp())


with torch.no_grad():
    latent_te = {
        "Coupled": enc_collapse(Xc_te),
        "No target + hinge": enc_var(Xc_te),
        "JEPA": enc_jepa(Xc_te),
    }
    jepa_std = latent_te["JEPA"].std(dim=0).numpy()
    coll_std = latent_te["Coupled"].std(dim=0).numpy()
    effective_rank_results = {
        name: effective_rank(features) for name, features in latent_te.items()
    }
    coupled_training_mse = float(
        (pred_collapse(enc_collapse(Xc_tr), U_tr) - enc_collapse(Xn_tr))
        .square()
        .mean()
    )

Out[13]:

Console

Raw observation         probe R^2 =  0.795
Pixel predictor         probe R^2 =  0.627
Coupled latent          probe R^2 = -4044.184
No target + hinge       probe R^2 =  0.293
JEPA latent             probe R^2 =  0.640

mean latent std | JEPA    : 1.8208
mean latent std | coupled : 0.0004
coupled training MSE      : 2.632e-10
effective rank | Coupled          : 7.051
effective rank | No target + hinge: 1.299
effective rank | JEPA             : 2.856

The coupled latent model collapses in this seeded run: its per-dimension standard deviation drops near zero, its training MSE is near zero, and the held-out linear probe fails. Its effective rank remains near the number of latent dimensions because effective rank normalizes away overall scale and measures the shape of the tiny residual numerical variation; it must be read alongside coordinate variance. The EMA-target-plus-hinge model retains coordinate variation and substantially linearly decodable state. On held-out probe $R^2$ , the no-target-plus-hinge variant lies between the collapsed and full variants. It avoids per-coordinate variance collapse but has an effective rank well below the number of dimensions. Because it also adds a learnable, action-conditioned predictor relative to the coupled baseline, this comparison does not isolate the hinge's causal effect; it shows a confounded variant with dimensional collapse.

Now consider a compatibility ranking derived from the same model. We use $E(\hat{z},z')=\|\hat{z}-z'\|_2$ as an energy over candidate successors, where $\hat{z}$ is the predicted embedding, $z'$ is a candidate successor embedding, and low energy means "this candidate is compatible with the prediction." A useful discriminative latent model should rank true successors above decoys often enough to outperform chance retrieval; stochastic dynamics, partial observability, and model error mean the true item need not rank first in every candidate set. This finite-candidate ranking does not define or calibrate a probability distribution over all successors.

In[14]:

Code

def retrieval_accuracy(pred_z, cand_z):
    """Expected top-k accuracy with random tie-breaking."""
    d = torch.cdist(pred_z, cand_z)  # lower distance = lower energy
    true_d = d.diag().unsqueeze(1)
    close = torch.isclose(d, true_d, rtol=1e-5, atol=1e-8)
    strictly_better = ((d < true_d) & ~close).sum(dim=1).float()
    tied = close.sum(dim=1).float()

    def expected_topk(k):
        # If a tie straddles the top-k boundary, give the true item its
        # probability of selection under a uniformly random tie break.
        available = torch.clamp(k - strictly_better, min=0.0)
        return torch.clamp(available / tied, max=1.0).mean()

    return float(expected_topk(1)), float(expected_topk(5))


K = 256
with torch.no_grad():
    jepa_top1, jepa_top5 = retrieval_accuracy(
        pred_jepa(enc_jepa(Xc_te[:K]), U_te[:K]), tgt_jepa(Xn_te[:K])
    )
    coll_top1, coll_top5 = retrieval_accuracy(
        pred_collapse(enc_collapse(Xc_te[:K]), U_te[:K]),
        enc_collapse(Xn_te[:K]),
    )

Out[15]:

Console

JEPA        top-1 0.016   top-5 0.066
Coupled     top-1 0.003   top-5 0.017

The tie-aware calculation is important. An exactly collapsed representation gives every candidate the same distance, so its expected top-1 score is $1/K$ , not 1.0. This is the compatibility-learning view in action. A planner can roll the latent model forward under candidate action sequences and rank the resulting latents with a goal-compatibility score, alongside any other task costs and constraints. The retrieval calculation above tests one-step successor discrimination, not planning performance.

Those top-k numbers summarize a single batch; the energy distributions behind them are more informative.

In[16]:

Code

# Deterministic compatibility-energy summary over the same held-out batch (K = 256).
# Energy is E(z_hat, z') = ||z_hat - z'||_2, so low energy means compatible.
# Selection rule: candidates are the first K held-out successors in index order,
# with no sampling and no shuffling. Diagonal entries pair each prediction with its
# own successor (compatible pairs); off-diagonal entries pair it with every other
# candidate in the batch (decoys). Tie-breaking is irrelevant for these summaries
# because we report the full distributions and their medians, not ranks.
with torch.no_grad():
    d_jepa = torch.cdist(
        pred_jepa(enc_jepa(Xc_te[:K]), U_te[:K]), tgt_jepa(Xn_te[:K])
    )
    d_coll = torch.cdist(
        pred_collapse(enc_collapse(Xc_te[:K]), U_te[:K]),
        enc_collapse(Xn_te[:K]),
    )

mask = ~torch.eye(K, dtype=torch.bool)
energy_true_jepa = d_jepa.diag().numpy()
energy_decoy_jepa = d_jepa[mask].numpy()

Out[17]:

Visualization

Overlaid density histograms of true and decoy successor energies. True-successor energies are lower on average, but the distributions overlap heavily. — Compatibility energy E = ||z_hat - z'||_2 on held-out transitions for the JEPA model. The true successors sit lower than decoys on average, but the distributions overlap heavily and the nearest decoy is usually closer to the prediction than the true successor, so point retrieval at this scale is only slightly above chance. A collapsed model produces the same energy for every candidate and no separation at all.

The predictor is recursive, so we can also test open-loop error instead of stopping at one step. Starting from each held-out trajectory's first observation, we feed the predicted latent back into the predictor for 20 actions. The target encoder supplies a reference latent at each horizon. This metric measures alignment with target-encoder coordinates as rollout length increases; it is not observation error and does not establish closed-loop control quality. Downstream planning must still be evaluated directly.

In[18]:

Code

ROLLOUT_HORIZON = 20
rollout_sq_errors = []

with torch.no_grad():
    for S, X, U in test_trajs:
        z_roll = enc_jepa(t(X[0:1]))
        trajectory_errors = []
        for h in range(ROLLOUT_HORIZON):
            u_h = t(U[h : h + 1, None])
            z_roll = pred_jepa(z_roll, u_h)
            z_target = tgt_jepa(t(X[h + 1 : h + 2]))
            trajectory_errors.append(float(((z_roll - z_target) ** 2).mean()))
        rollout_sq_errors.append(trajectory_errors)

rollout_mse_by_horizon = np.mean(np.asarray(rollout_sq_errors), axis=0)

Out[19]:

Console

latent rollout MSE at horizon 1 : 2.3980
latent rollout MSE at horizon 20: 1.7569

The final MSE is below the horizon-one value, and the intervening curve fluctuates. The curve is not monotone: on this dataset the error fluctuates with horizon rather than accumulating steadily. A useful one-step representation can still drift in open loop, which is why rollout error must be reported separately from one-step retrieval and probe scores.

The shape of that error curve is worth reading directly rather than through its first and last values.

Out[20]:

Visualization

Line plot of latent rollout mean squared error over 20 prediction steps. Error fluctuates nonmonotonically and finishes below its horizon-one value. — Open-loop latent rollout error for the JEPA model on the held-out trajectories. Error is measured against the target encoder at each horizon and fluctuates as the predictor consumes its own outputs; on this dataset it does not grow monotonically, so one-step retrieval accuracy does not describe long-horizon usability.

Out[21]:

Visualization

Symlog bar chart of held-out linear-probe R-squared. The coupled latent is a far-negative outlier; the other four representations are positive. — Linear-probe $R^2$ for recovering the true two-dimensional state from each representation. The coupled latent model contains almost no linearly recoverable state signal under this probe despite a training MSE near zero.

Log-scale grouped bars with different hatch patterns. Coupled coordinate standard deviations are far smaller than JEPA values. — Per-dimension standard deviation of the learned latents on a logarithmic scale. Coupled values are about four orders of magnitude below the EMA-target JEPA values.

Two patterns stand out. First, the raw-observation reference has the strongest linear probe: the state is substantially linearly decodable from informative directions in the noisy observation, although the noise prevents $R^2=1$ . Second, the pixel predictor is competitive on this toy problem. The setup lacks the structured, high-dimensional nuisance content of real images even though all 32 observation coordinates contain independent noise. Latent prediction shows no probe advantage here, and the experiment does not measure compute efficiency or scaling. Predictive objectives can become more useful as unpredictable nuisance variation grows while task-relevant structure remains predictable, but that is a conditional design argument rather than a monotonic guarantee.