EpiWM

Beating leWorldModel with the EpiJEPA objective in place of SIGReg.

Code · Checkpoints

Introduction

In EpiJEPA, I explored an epiplexity-inspired objective as an alternative to SIGReg for preventing embedding collapse. EpiJEPA does not force the embeddings to follow a prescribed distribution, which in theory gives the model more freedom to find the representation best suited to its task. EpiWM takes that idea into world modelling.

The goal here is simple:

• Replace SIGReg with the EpiJEPA objective.
• Show that representational expressivity can improve performance.

I compare EpiWM with LeWM's released checkpoints and my own LeWM reproductions (labelled SIGReg controls) across the paper's four control environments: PushT, TwoRoom, OGBench-Cube and Reacher.

Theory

World models predict the consequences of actions, letting an agent plan before acting. The argument for JEPA-based world models is that predicting in embedding space lets them discard nuisance information and focus on what matters for control. This is shown to linearise the task, turning a traditionally generative problem into a single point regression.

Personally, I think a generative objective is still essential, and that JEPA and generative methods can be combined. unlike Yann

LeWM

LeWM jointly trains an encoder and an action-conditioned predictor: frames \(x_n\) become embeddings \(z_n\), and the predictor estimates the next embedding from the recent history and actions.

Training minimises the error between \(\hat z_n\) and the target frame's embedding \(z_n\). However, this task has a trivial solution: \(z_n = C\) for all \(n\). SIGReg prevents trivial collapse by encouraging an isotropic Gaussian embedding distribution.

But an isotropic Gaussian may not be the optimal representation distribution. EpiWM replaces that regulariser with the EpiJEPA objective, which does not force a particular distribution.

EpiJEPA

Epiplexity measures the information that a computationally limited observer can learn from data.

EpiJEPA uses an approximation of this idea to reward learnable information in the embeddings. A small linear predictor learns to predict the embeddings from features extracted from the same frames by a frozen, randomly initialised CNN (this small predictor is only for loss calculation). I fit it using ridge regression, then use a log-determinant score to measure the strength and spread of its weights (how much it learnt).

If the embeddings collapse to a constant, the task is trivial and the score is zero, indicating nothing had to be learnt. Rewarding this score gives the encoder an incentive to preserve information connected to the input.

Method

Model

EpiWM starts from LeWM's released architecture, with the following changes:

† The epiplexity score can grow without bound as the embeddings grow in scale. The extra BatchNorm fixes this scale, with no learnable scale or shift. This is the same as in EpiJEPA; SIGReg fixes scale through its Gaussian target.

For PushT, TwoRoom and Cube, I use the same training datasets as LeWM. My Reacher* models were trained on the released dmc/reacher_random dataset, collected using a random policy. The paper instead describes training on data collected using a Soft Actor-Critic (SAC) policy. I evaluate all Reacher models on the released random-policy dataset, so these results are not directly comparable to the paper’s reported 86%.

Evaluation

I compare three models: the released LeWM checkpoints, EpiWM, and my own SIGReg controls. Each task asks the model to plan from a starting state to a target state. Each environment has three evaluation sets, all selected with seed 42: the original 50 tasks, a non-overlapping set of 200 tasks, and a further 500-task evaluation. Start states come from the training datasets, following LeWM’s protocol.

Planning uses CEM, which samples action sequences and refines those predicted to reach the goal. Unless noted, all models use the same released CEM settings, including 30 iterations in every environment.

Then, to compare the underlying learnt representations, I visualise them using PCA. I pass the same 3,000 randomly sampled dataset frames through each model and take the projector outputs, the embeddings used for planning. I fit PCA separately for each model and plot the first two principal components. Frames from a recorded episode are projected using each model’s fitted PCA and traced in time with the video. The percentages show how much of the embedding variance these two dimensions capture.

Results

Planning

EpiWM beats the released checkpoints across all four environments on the larger evaluation sets. Against my controls, PushT is effectively tied. The original 50 PushT tasks favour released LeWM; both larger sets reverse that ranking.

Representations

The EpiWM representations are visibly more structured in these PCA plots, especially in Cube and TwoRoom, where the first two principal components capture recognisable x and y structure from the video frames. This makes the physical motion easy to follow in latent space.

These analyses compare one selected EpiWM checkpoint per environment (the best on the 500-task set) with its released LeWM counterpart.

Conclusion

Replacing SIGReg with the EpiJEPA objective matches or improves planning performance across all four environments on the larger evaluation sets. The PCA plots also show more recognisable spatial structure, particularly in Cube and TwoRoom.

These results support the idea behind EpiWM: preventing collapse need not mean forcing the embeddings towards a prescribed distribution. Rewarding learnable information gives the model room to develop its own representation, and here that freedom produces a useful world model.

Code · Checkpoints

Comments