Experiments on Hierarchy

8 June 2026

Planning in world models is slow, and one of the oldest fixes is hierarchy: instead of searching over every primitive action, let a high-level model propose a sparse sequence of subgoals, and let a cheap low-level controller fill in the gaps. This is a project log of three attempts to make a two-level world model (H-JEPA) plan better than its single-level counterpart on the Push-T task, plus a diagnosis of why each one fell short. The three differ only in how the abstract, high-level state is formed: by merging tokens, by a learnt encoder on the latent, and by a learnt encoder on raw pixels.

Setup: the two-level world model

Both levels are predictors over frozen DINOv2 image features, and we never fine-tune the encoder. A 224×224 Push-T frame (one camera) becomes a 16×16 = 256-token feature map (ViT-S/14, 384-d). The low-level (LL) model is the usual JEPA world model: it predicts the next 256-token feature map one short step ahead, conditioned on primitive per-frame actions. A token-merge step compresses 256 → 128 tokens into an abstract state, and the high-level (HL) model predicts the next abstract waypoint over a longer time step, so it abstracts in both space and time. The single-level planner runs CEM in the LL latent; the hierarchical planner asks the HL model for a subgoal and lets the LL planner chase it.

One question organizes everything below: does the high-level subgoal help planning at all?

Idea 1: Token Merging for the abstract space

The higher-level state space is obtained by applying Token Merging (ToMe) to the low-level patch tokens.

Token merge masks: random vs importance-based attention encoder — How tokens are selected for merging: a random mask versus an importance-based attention encoder.

Diagnosis: is the high-level subgoal helping at all?

The control experiment is brutal. We keep the same final approach but drive the first leg of the trajectory three ways: follow the predicted subgoal, take random actions, or do nothing. Following the model's predicted subgoal succeeds 10% of the time; random actions or staying idle succeed 45-50%. So the predicted subgoal steers the controller to a worse state than the start. It carries negative information.

Bar chart: model subgoal 10%, random actions 50%, stay idle 45% — Replacing the subgoal-guided first leg with random actions, or no motion at all, beats following the predicted subgoal.

Two hypotheses

Either the subgoals are bad, or the space they live in is bad:

H1, bad subgoals (prediction error). The token-merged space is fine, but the HL model's predicted next abstract state is inaccurate, so the LL planner is sent to the wrong place.
H2, bad abstract space (aliasing). Merging tokens discards detail, so many physically different states map to nearly the same abstract code. Then "minimise abstract distance to the subgoal" is satisfiable by the wrong states, and guidance is no better than random.

We separate them with a clean trick: feed a perfect subgoal, which removes prediction error and isolates the space.

Testing H2: is the abstract space faithful?

We take the true subgoal and a large pool of real Push-T states (its own trajectory plus 1500 from other trajectories) and rank each one by distance to the subgoal, once in the 128-token abstract space and once at full resolution. The rankings agree almost perfectly: Spearman ρ = 0.99 at both strides, and the abstract-nearest state is also the full-resolution-nearest. Distinct states do not collide, so H2 is refuted.

Scatter of abstract-space distance vs full-resolution distance, Spearman rho 0.99 — Distance in the 128-token abstract space ranks real states almost identically to full resolution (ρ = 0.99). No aliasing.

Testing H1: the decisive comparison

Holding the low-level model and CEM budget fixed and changing only the LL planner's target: the predicted subgoal gives 10%; planning in the abstract space toward the true goal (no HL model) gives 62%, essentially tied with single-level's 70%; and handing over the true intermediate subgoal gives 78%, above single-level. So the abstract objective supports planning, the decomposition genuinely helps, and the only thing separating 10% from 78% is the accuracy of the predicted subgoal.

Bar chart: predicted subgoal 10%, abstract-space true goal 62%, single-level 70%, true subgoal 78% — With a correct subgoal, the hierarchy beats single-level. Only the low-level planner's target changes across the four bars.

Why a 'small' error is fatal

Most of a Push-T frame is static background; only a few tokens move between waypoints. The HL model's one-step error is about 8% of the full feature but roughly 50% of the actual inter-waypoint motion. Because the abstract space is faithful, the LL controller is accurate, so it faithfully drives to the wrong predicted state, and open-loop error compounds over successive macros. Staying idle wins because it leaves the agent where the final approach can still finish.

Two panels: open-loop error compounds over macros; error is 8% of feature but 50% of motion — Open-loop error compounds across macro-steps (left); the one-step error is small against the whole feature but large against the tiny inter-waypoint motion (right).

Conclusion

Token Merging is a fine way to abstract a static image, but it throws away exactly the information needed for next-state prediction, especially over longer time scales. This "blind spot" of token-reduction methods has been seen elsewhere too (Chun et al., 2025). We need a learnt abstraction, not a hand-merged one.

Idea 2: a learnt encoder on the low-level state

The higher-level state space is obtained by applying a learnt encoder on the lower-level state.

From here the representation is a CLS token, since we move to the LeWorldModel (LeWM) backbone.

Two-level architecture with high-level and low-level predictors and macro-actions — The two-level architecture: a learnt high-level encoder and predictor sit on top of the low-level latent.

Tests

Latent space PCA. Looks good enough, and an L2 cost metric should work in HL space.

Temporal straightening should ideally sit near 1. LL gets 0.6; HL gets -0.45. LeWM showed SIGReg enforces this automatically, but it isn't kicking in for the high level.

Temporal straightening curve, high level far below low level — Temporal straightening: the high level fails to straighten (-0.45) where the low level reaches ≈0.6.

L2-norm satisfaction. HL and LL should show a similar gradient of pairwise distances along an expert trajectory, but they don't. Pearson correlation is only 0.6-0.8 across 100 episodes.

Pairwise-distance heatmaps for low level and high level along an expert trajectory — Pairwise-distance structure along an expert trajectory differs between the two levels.

Decoding rollouts. The planner can't succeed even when handed expert actions and expert subgoals.

Decoded rollout frames showing failure even with expert inputs — Decoded rollouts: the planner fails even when handed expert actions and expert subgoals.

Problem 1: states are off-manifold

Errors diverge as autoregressive rollouts proceed, even with expert actions provided.

Rollout error diverging over steps even with expert actions — Predicted states drift off-manifold; error diverges as the autoregressive rollout proceeds.

Attempt 1, SIGReg. Train the higher level with SIGReg: -20% success at horizons beyond ~50 env steps, less noticeable at short horizons.
Attempt 2, RDMReg + sparsity. -5% at 25 steps but -50% at longer horizons; this solves the prediction-error problem itself.

RDMReg prediction error vs horizon — Training the high level with RDMReg and enforcing sparsity solves the prediction-error problem.

Problem 2: actions are off-manifold

The deeper issue: macro-actions produced by HL CEM correspond to no physical state, so the subgoals are unreachable and the LL CEM can't get to them, which tanks the success rate.

Macro-action PCA showing off-manifold actions — Macro-actions from HL CEM land off-manifold, producing subgoals that correspond to no physical state.

We threw a lot at this:

Attempt 1, improve the eval strategy. Optimise LL actions directly through the MAE instead of the LL predictor, giving on-manifold macro-actions at the cost of slower evals: 36% → 40% (5 seeds, 50 random episodes, fixed budget).
Attempt 2, scale the model. SOTA is 75M params; we scaled 18.5M → 25M → 33M, giving 40% → 54% → 50%.
Attempt 3, change architecture. Bottleneck the macro-action dimension to 4 to cut off-manifold risk (on RDMReg checkpoints). Error stayed high versus simply using the LL predictor; only the first subgoal was accurate.
Attempt 4, better action prior. Regularise the MPC cost so actions look expert-like, sampled from a normal prior.
Attempt 5, modify the MPC cost. Add action smoothness and jerk penalties.

Eval strategy that optimises low-level actions through the MAE — The improved eval strategy: optimise low-level actions directly through the MAE to keep macro-actions on-manifold.

Summary of open problems

The recurring theme is an interface problem: the two planners don't agree on what a macro-action is.

Representation learning. Sparser representations, and an invariance loss to learn slower features.
Off-manifold HL actions. An IDM to learn macro-actions, or an action codebook.
HL/LL interface. An HL world model that predicts at multiple strides, plus contracts between the levels, since HL CEM macro-actions don't match LL CEM planned actions encoded into macros.
Eval strategy. A flexible horizon, since the LL planner can't reach subgoals that are non-uniformly spaced.

Idea 3: a learnt encoder on the pixel space

The higher-level state space is obtained by applying a learnt encoder on the pixel space.

Motivation. If the LL latent is lossy and collision prone, a high level built on top of it inherits its failure modes. So we train the HL as its own ViT-tiny over raw pixels (Pixel → HL); if pixels feed both LL and HL independently, the HL is no longer chained to the LL's representational choices.

Architecture. A PixelHLEncoder (ViT-tiny + linear projector → 96-d embedding), an HL predictor refactored into an LL-style contextual autoregressive model (5.07M params), a frame-skip K = 5 so each HL step covers 25 env steps (matching LeWM's macro horizon), a standard JEPA L2 + SIGReg loss, and a small distance head D(target_hl, ll_state) to bridge the HL pixel space and the LL latent for planning. Trained on PushT expert data, ~6M params total.

Diagnostics: temporal straightening

Straightening climbs from 0.088 (epoch 1) to 0.272 (epoch 4) before plateauing. That's still short of LL's ≈0.6, but a real improvement over the -0.45 of the LL-latent HL. PCA traces of expert rollouts are smoother than LL and coarser in time, which suggests the HL is capturing macro-trajectory shape rather than per-frame structure.

Pixel-HL temporal straightening improving across epochs then plateauing — Pixel-HL straightening improves (0.088 → 0.272) then plateaus. Better than the LL-latent HL, still below LL.

Decoder

A symmetric ViT-tiny decoder over LL (192-d) and HL (96-d), trained with L2 + LPIPS, reaches validation reconstruction of LL 0.048 and HL 0.046. The HL slightly beats LL despite a 2× tighter bottleneck.

Probing the frozen low-level encoder

Every downstream planner inherits the frozen LL encoder's flaws, so we probed it hard, and the news is bad.

Natural collisions exist. Across 2000 frames, many pairs have near-zero LL embedding distance while their pixel distance is large: different scenes map to nearly the same LL latent.

Scatter of pixel distance vs low-level embedding distance showing natural collisions — Pixel L2 distance vs LL embedding distance across 2000 frames. Many pairs collide in latent space despite large pixel differences.

Adversarial budget is tiny. A targeted perturbation of roughly 8 unit-radius steps in the LL latent is enough to decode to a completely different Push-T state. Any planner searching the LL latent without a manifold prior will drift into invalid states.

Long-horizon reachability fails. Sweeping CEM over horizons 20 and budgets 4096, the final distance-to-target plateaus. More horizon doesn't help, and neither does more compute. The self-pair drift floor sits around 1.3 latent units, and expert actions reach states CEM cannot. Long-horizon and cross-trajectory pairs are effectively unreachable under CEM in the LL latent.

CEM convergence curves plateauing well above the target distance — CEM convergence is flat from early iterations and plateaus above the target, so long-horizon targets are unreachable.

The decisive comparison

The one genuinely encouraging result. Ranking HL cost across four strategies: warm-starting HL from a quick LL+MAE search (0.30) < macro CEM (0.54) ≪ pure LL CEM (3.75) ≪ expert+MAE (27.77). Warm-starting the high level from a cheap low-level search beats both macro-only HL search and pure LL CEM.

Decoded planning comparison across four strategies — Decoded rollouts for the four planning strategies. Warm-starting HL from an LL+MAE search yields the lowest cost (0.30).

And the honest counterweight: when the HL macro-action is used as a soft constraint on LL CEM (regularise rather than replace), the score plateaus at 30-45% across many λ settings and HL checkpoints, while pure LL MPC at a matched budget hits 55-90%. On Push-T, HL-as-soft-constraint never beats pure LL.

Takeaways

Three abstractions, one consistent lesson. The abstract space is rarely the problem. Token-merged or learnt, distances in it rank states faithfully. The trouble is always downstream: high-level prediction error that looks small against the whole frame but huge against the tiny motion that matters, macro-actions that drift off-manifold into states no controller can reach, and a frozen low-level encoder whose collisions and unreachable regions every planner inherits. The one move that clearly pays off is warm-starting the high level from a cheap low-level search, and it points the way. Hierarchy helps when the two levels agree on what an action is, and hurts when they don't. The next round is about closing that gap: learnt macro-actions, multi-stride prediction, and a low-level space that's actually safe to search.