Summary

In this post, I train LeWM without SIGReg and compare its embeddings with normal runs.

Q: What’s the purpose of this experiment?

I want to see what SIGReg does. To do this, I compare the embeddings of models trained with and without SIGReg, and look for representation collapse in the model without it.

The non-SIGReg models is trained with the same setup as my baseline (image size 112, 3 epochs), except that the SIGReg weight ($\lambda) is set to 0.

In addition, I also look at the embeddings of another model trained with image size 224 and 3 epochs.

Q: What types of collapse am I looking for?

• Complete collapse. All images map to almost the same point and the spread becomes very small.
• Dimensional collapse. The points spread out, but only along a few directions. The other directions are not used.

Q: What is the input data?

I created a script that encodes a fixed set of frames with each checkpoint and saves the embeddings (cls, emb) and the true state and proprio of each frame.

I use all 3 checkpoints from each of these runs:

• pusht/2026-09-28_114752_img112_s3072 (image size 112, 3 epochs)
• pusht/2026-10-14_115337_img112_s3072_nosigreg (same as above but without SIGReg)

and all 10 checkpoints from:

• pusht/2026-10-10_120411_img224_s3072 (image size 224, 10 epochs)

The frames are in two groups: 20,000 random frames from the whole dataset, which I use for the statistics, and 1,094 frames of 10 random episodes, which I use to draw trajectories in the PCA plots. That is 21,094 frames in total.

Q: For each test, what metrics are compared?

• Scale: how big is the point cloud.
  • length of mean vector (|mean|)
  • spread
  • spread / |mean|
• Direction: are all the points in the same directions?
  • histogram of 1 - cosine similarity
  • calculated between random frames
• Eigenvalue spectrum: how many directions are used?
  • eigenvalues of covariance matrix
  • effective rank
  • participation ratio
• Is the distribution close to N(0,I)?
  • SIGReg statistic
  • histogram along random directions
• Linear probe
  • train a ridge regression from latent to states (agent x, agent y, block x, block y, cos(block angle), sin(block angle))
• PCA
  • does any principal component match the position or angle of the T block?

Q: What does the result say? What’s the restriction or drawback of SIGReg?

Without SIGReg, the embeddings collapse completely: the spread is about 10⁻⁵, compared with about 14 with SIGReg. The planner cannot use this model: its success rate on PushT is only 2.8% (mean of 5 seeds).

With SIGReg, there is no complete collapse, but there is some dimensional collapse: the effective rank is only about 69 of 192 at epoch 3.

Longer training moves the distribution closer to N(0, I), but it does not get very close. In the 10-epoch run, the effective rank stops at about 98. (This run also uses a larger image size, so the comparison is not clean.)

Q: Future experiments I want to run?

The PushT scene has only about 5 degrees of freedom, but the latent has 192 dimensions. This can limit how many dimensions the encoder can use. A more complex scene may show if this is the cause.

Experiment Details

All results are from the notebook playground/sigreg_collapse_notebook.py in the le-wm repo. Unless I say otherwise, the latent is emb (the projector output, which is the input of the predictor and the latent that SIGReg acts on), and the epoch is 3, the last epoch that all three runs have.

Summary table

emb, epoch 3:

1. Scale: how big is the point cloud?

• |mean|: the length of the mean latent vector, i.e. the position of the center of the cloud.
• spread: the square root of the total variance (the sum of the variances of all 192 dimensions). It is the typical distance from a point to the center. For N(0, I) in 192 dimensions, the spread is $\sqrt{192} \approx 13.9$.
• spread / |mean|: the size of the cloud relative to its distance from 0. A very small value means that all points are at almost the same place, far from 0.

The plot below shows the std of each of the 192 dimensions, sorted from the largest to the smallest. The y-axis is logarithmic.

Result. Without SIGReg, the cloud is complete collapse: the spread is 1.2 × 10⁻⁵, and it becomes smaller at each epoch (2.0 × 10⁻⁴ → 3.6 × 10⁻⁵ → 1.2 × 10⁻⁵). With SIGReg, every dimension has a std between 0.8 and 1.2, and the spread is very close to $\sqrt{192}$.

2. Direction: are all points in the same direction?

I take 5,000 random pairs of frames and calculate the cosine similarity of their latents. The plot shows 1 − cosine similarity on a log x-axis. For N(0, I), two random vectors are almost perpendicular, so 1 − cos is near 1. If the latents collapsed, all pairs point in the same direction, and 1 − cos is near 0.

Result. The median of 1 − cos is 1.01 for SIGReg, 1.02 for the img224 run, and 4.7 × 10⁻⁹ for no SIGReg. Without SIGReg, all latents point in the same direction (the cloud is a very small ball at distance 0.16 from 0).

3. Eigenvalue spectrum: how many directions are used?

The eigenvalues of the covariance matrix are the variances along the principal components. The left plot shows the raw eigenvalues (for N(0, I), all are 1). The right plot divides them by their sum, so it shows only the shape of the cloud.

Two numbers summarize the shape:

• Effective rank: the exponential of the entropy of the normalized eigenvalues:

• Participation ratio: $(\sum \lambda_i)^2 / \sum \lambda_i^2$. It gives more weight to the large eigenvalues.

Result: SIGReg runs have a dimensional collapse.

• There is a “cliff”. In the right plot, the SIGReg run drops near component 60, and the img224 run drops near component 75. At the cliff, the values fall by about 100 times. After the cliff, each direction has only 10⁻³ to 10⁻⁵ of the total variance, which is 5 to 500 times less than an equal share (1/192 ≈ 5 × 10⁻³).
• The used directions are larger than in N(0, I). In the left plot, the first 60–75 eigenvalues are about 2–4. SIGReg keeps the std of each coordinate near 1, so the total variance stays near 192. But only about 65 directions share it, so each gets about 192 / 65 ≈ 3. The cloud is a flat “pancake”: wide in about 65 directions and very thin in the others.

The no-SIGReg line in the right plot has no meaning. It looks “normal”: it decreases slowly, with no cliff. But the left plot shows that its eigenvalues are only 10⁻¹¹ to 10⁻¹⁶, so the cloud is almost one point. The right plot removes the size, so the shape of this tiny cloud tells nothing. For the same reason, its “directions for 99%” (86) is higher than for the SIGReg run, but this is not a good sign.

4. Is the distribution close to N(0, I)?

4.1 The SIGReg statistic

I calculate the SIGReg loss as in training: random batches of 128 frames, 1024 random directions, averaged over the batches. All models use the same batches and directions (fixed seed).

The batches must be random. The frames in the files are sorted by dataset row, and nearby episodes look alike. A first version without the shuffle gave about 14 for the SIGReg run at epoch 3, compared with 1.5 with random batches.

The value is about $c + B \times D$, where $c \approx 1$ is a noise floor (a finite batch is never exactly Gaussian), $B$ is the batch size, and $D$ is the real distance from N(0, I). Real samples from N(0, I) give 1.05. Only the part above this floor is a real distance.

• Raw: on the latents as they are. This is what the training loss sees.
• Standardized: after each dimension is shifted and scaled to mean 0 and std 1. This removes the scale problem and shows only the shape problem.

4.2 Histograms along random directions

I take 4 random unit directions (the same for all runs) and project every latent on them. The gray curve is the N(0, 1) density. All plots use the same x-axis (−4 to 4), so that a collapsed cloud shows as one thin spike and is not stretched to look like a bell curve.

Result. The SIGReg runs follow the gray curve well. For no SIGReg, all values fall into one bar (its height is 10, because the bar is 0.1 wide and holds all the frames).

5. Linear probe: is the information about the scene still there?

I fit a ridge regression from the standardized latent to the true state: agent x, agent y, block x, block y, cos(block angle), sin(block angle). Train and test frames come from different episodes (80% / 20%). Because the latent is standardized first, the probe does not see the scale, so even a cloud that is 10⁻⁵ wide can give a good result.

Result. The collapsed latent still has some information about the positions, but much less about the block angle. The probe R² of no SIGReg does not change from epoch 1 to 3 (0.54). For the planner, this information is not useful: the planner compares raw latents with a squared distance, and when all distances are about 10⁻¹⁰, the cost differences between action plans are almost 0.

The block angle is the hardest target for the 3-epoch SIGReg run (R² ≈ 0.65). The img224 run learns it much better (0.92–0.95 at epoch 3).

6. PCA

For each run, I fit a PCA on the 20,000 random frames and plot PC1 against PC2. The colors are the true state (the model never sees it): block angle in the top row, agent x in the bottom row. The black lines are 3 full episodes, with frames connected in time order.

• Size. The no-SIGReg cloud is about 10⁶ times smaller than the other two. The three plots look the same size only because each has its own axes.
• Share of variance. In the SIGReg runs, about 65 directions share the variance, so each PC has only about 2%, and the plot shows only about 4% of the cloud.
• Orange outer frames. Many frames have a block angle near 0.8 rad. This is possibly the goal angle of PushT, π/4 ≈ 0.785. I did not verify this.
• No PC matches one state. A PC is a mix of several states, so the patterns are rough. Agent x shows the clearest pattern. The block angle shows no clear pattern in PC1–PC2.

7. Change over epochs

• No SIGReg collapses already at epoch 1 (total variance 4 × 10⁻⁸), and it becomes smaller at each epoch. A checkpoint is saved only at the end of each epoch, so I cannot see when the collapse starts inside epoch 1.
• SIGReg: the effective rank goes up from 55 to 69 in 3 epochs.
• SIGReg 10 epochs (img224): the total variance is constant at about 192. The effective rank goes up from 51 (epoch 1) to 98 (epoch 10), and it stops at about epoch 8.

The SIGReg statistic goes down while the rank goes up (2.31 → 1.33), so the SIGReg loss “sees” the low rank and pushes against it. At epoch 10, the part above the N(0, I) floor is only about 0.3.

Why the rank does not go higher. A separate test on the SIGReg run at epoch 3 whitened the latents (it made all eigenvalues equal to 1 and kept the rest of the shape). This removed about 90% of the distance above the floor. So most of the remaining SIGReg loss comes from this anisotropy. But with a batch size of 128, it adds only about 0.45 to the loss, which is less than the noise floor (about 1.05). So the signal against the remaining anisotropy is weak. A larger batch or a larger SIGReg weight makes this signal stronger.

Be careful with these comparisons:

• Learning-rate schedule. The scheduler is a linear warmup and a cosine decay to 0 over all training steps. In the 3-epoch run, the learning rate is almost 0 at the end of epoch 3. In the 10-epoch run, it is still about 80% of the peak at the end of epoch 3. This is one reason why the 3-epoch run stops at about 69. It also means that the plateau at epochs 8–10 can come from the learning rate that goes to 0, and not from a real limit.
• Image size. The 10-epoch run uses 224-pixel images, and the baseline uses 112. The image size and the training length change together.
• One seed. Each run is one training run.

8. The encoder output (cls)

The projector is Linear → BatchNorm → GELU → Linear, so the BatchNorm is in the hidden layer and does not prevent a collapse of emb. To see if the encoder itself collapsed, I also look at cls (the encoder CLS token, before the projector).

• Without SIGReg, the encoder output also collapses completely (spread 1.6 × 10⁻⁴ at a distance of 13.6 from 0). So the collapse is not only in the projector.
• With SIGReg, cls is not collapsed, but its effective rank is much lower than emb (31 vs 69 at epoch 3; 36 vs 98 at epoch 10 for the img224 run). SIGReg acts on emb, not on cls, so the projector spreads a lower-rank encoder output over more directions.

9. Planning success rate

I evaluated the no-SIGReg model (epoch 3) with the CEM planner on PushT, with 5 different seeds. Each seed gives a different set of 50 start states from the dataset. The goal is the expert’s state 25 steps later, and the planner has 50 steps to reach it.

python eval.py --config-name=pusht.yaml policy=pusht/2026-10-14_115337_img112_s3072_nosigreg/weights_epoch_003.pt eval.img_size=112 seed=<0..4>

Result. The planner reaches the goal in only 0–2 of 50 episodes for each seed. This agrees with section 5: when all latents are at almost the same point, the planning cost is almost the same for all action plans, so the planner cannot select a good plan.

I did not evaluate the SIGReg baseline or a random policy with the same seeds yet, so I cannot say if 2.8% is better than chance.