R-LeWM: does a JEPA world model remember what it can no longer see?

These are LeWorldModel (LeWM) checkpoints trained for UAV depth navigation, plus a ladder of variants (M0–M4) that try to give LeWM memory. A drone flies to a goal through boxes and walls, including dead-end traps, seeing only a 48×80 forward depth image. Planning around a trap needs memory of walls that have gone out of view.

Code (training, evaluation, figures): github.com/fuisl/single-agent-wm, branch lewm
Full write-up: docs/single_agent_memory.md
Status (2026-10-26): lewm, m0, m1, m2, m2s, m3s and m4s are here.
- m3s has the strongest and longest-lasting memory (15 % of the re-entry error still at 9.5 s out of view). m4s decodes remembered obstacles best but forgets within ≈ 2 s (findings 9–10).
- m2s is the same architecture as m3s on the original data, the first model with a clear memory effect.
- m2 is kept for reference: its frame encoder collapsed (finding 7). M3 and M4 had the same flaw and are not published; m3s / m4s are the fixed re-runs.

What LeWM does, in one picture

 depth frame o_t ──► ViT-tiny ──► CLS ──► projector ──► z_t ────────┐
 (48×80, 1 ch)        (patch 8)                       (192-d)       │   3-frame window
                                                                    ▼   (0.75 s)
 actions a_t ─────────────────────► action embedder ──► transformer predictor ──► ẑ_{t+1}
                                                                    │
 loss = ‖ẑ_{t+1} − z_{t+1}‖²  +  0.09 · SIGReg(z)                    │   no decoder, no EMA target:
        (predict the next       (keeps z Gaussian,                 │   SIGReg alone prevents collapse
         embedding)              so it cannot collapse)

The model never reconstructs pixels: it only predicts the next embedding. Planning uses CEM over action sequences in this latent space. The predictor sees only the last 3 frames, so any obstacle older than 0.75 s has to survive inside the embedding, or it is gone.

Findings so far

#	finding	evidence
1	LeWM's global embedding throws away obstacle geometry, but its own encoder keeps it.	A small MLP probe that decodes the ego-centric obstacle map scores 0.07 AP on visible obstacles from the embedding, against 0.28 from the same model's patch tokens and 0.33 from a frame-only CNN. Depth decodes from the patch tokens to 0.03 m error.
2	Adding the body state into the embedding creates a shortcut: the model memorises training layouts.	Depth decoded from the embedding is 0.09 m off on training episodes but 0.62 m on held-out layouts. Nearest neighbours in the embedding come from the same episode 62 % of the time (chance is 5 %).
3	LeWM has no memory of out-of-view obstacles.	In the twin look-away test, prediction error is exactly the same whether or not the wall was seen before looking away. A probe for remembered obstacles scores 0.058 AP, at the frame-only level (0.062); a DreamerV3 RSSM scores 0.105.
4	The next-frame objective almost never asks for memory.	Only about 1 % of obstacle cells in the next frame were last seen before the 3-frame window (2.8 % at the 1.25 s planning horizon).
5	Fixing the shortcut (M0) brings geometry back; adding spatial tokens (M1) does not help further.	Visible-obstacle AP rises from 0.07 to 0.24 with M0, then falls to 0.21 with M1. Neither has memory, as expected for 3-frame window models.
6	A recurrent belief (M2) gave only a small memory effect, because its frame encoder had collapsed.	M2's effect in the twin test was real (3–6 s.e.) but only ≈ 0.4 % of the re-entry error. Its frame embedding had an effective rank of 8.6 (M0 75, M1 68), and the wall barely changed it.
7	Cause: SIGReg's strength scales with the number of samples it is estimated on.	The SIGReg statistic carries a factor N. To fit 32-frame BPTT, M2–M4 trained at batch 16, so SIGReg ran at N = 16 instead of LeWM's 128: ≈ 8× weaker at the same λ. It looked normal at N = 16 (≈ 2), but scored 10–11 at N = 128 (M0: 1.9). The fix pools the batch × time frames and draws 128 per token position (as in LeWM-Flat, MotionJEPA), with λ unchanged.
8	With the fix (M2s), the belief remembers.	Encoder rank recovers from 8.6 to 47, and the wall signal from 0.63 to 1.71. In the twin test the memory effect is 9–10 s.e. at every gap. It is 4.7 % of the re-entry error at 1.5 s out of view and decays to 1.1 % at 9.5 s, 6–50× M2's. A removed wall is noticed in 100 % of scenarios up to 3.5 s. Memory-cell probe AP is 0.081: above the frame-only CNN (0.062), its own percept (0.053) and the state-only level (≈ 0.05) at 0.5–4 s out of view. DreamerV3 scores 0.105.
9	Training data that asks for memory makes it large and necessary (M3s).	Trained on flights with look-away glances and with 2–12 frames of the image blanked at random, the memory effect is 31 % of the re-entry error at 1.5 s and still 15 % at 9.5 s (12–15 s.e.; decay constant 14 s against 6 s for M2s). A removed wall is noticed in 100 % of scenarios at every gap. Zeroing the belief at the turn-back now raises the error by 7–10 %: the model depends on it.
10	An ego-centric map belief (M4s) decodes best but forgets fast.	Memory-cell probe AP 0.085, within 0.015 of DreamerV3 from 1 s out of view on. Memory effect 58 % at 1.5 s, but the decay constant is 2 s (0.8 % at 9.5 s), and after ≈ 2.5 s a reset helps (13–27 % at ≥ 3.5 s): the map holds a stale belief. Cause: the gated map write is not restricted to the camera's field of view, so every frame spent looking away overwrites the cells of the wall behind.

A probe on the state and action history alone (no image) scores 0.057 on remembered cells; the M2s–M4s beliefs are above it for up to ≈ 4 s out of view, so what they hold there is memory of the image, not a prior from the state.

The collapse and the fix, at a glance. From left: effective rank of each model's frame embedding; what an obstacle probe reads out; the memory effect against time out of view, as a share of re-entry error (log scale, so the slope is the decay rate); the change in re-entry error when the belief is zeroed at the turn-back:

A flight through a trap episode. The panels show the top-down world, a chase camera, the onboard depth, depth decoded from LeWM latents, and probe maps of where obstacles are (truth, DreamerV3, LeWM):

The same kind of flight with the probe maps read from the M2s, M3s and M4s beliefs instead of LeWM's patch tokens. Orange cells in the top-down view were seen and are now out of view:

The twin look-away memory test. The drone sees a wall W (or not), looks away for 2 s, and turns back. At re-entry, the seen and novel twins show pixel-identical frames, so any difference in prediction error can only come from memory. For LeWM the two error curves lie on top of each other. For M2s the novel curve (orange) rises above the seen one (blue) at re-entry. The figures below average over 60 scenarios and show the effect as a share of each model's re-entry error:

does the model expect the wall it saw? (E1)	is it surprised when the wall is gone? (E2)

what a probe can read from each latent	why: encoder vs embedding vs state shortcut

memory fades within seconds (LeWM ladder vs DreamerV3)	the training objective rarely needs memory

Models

folder	class	change from the row above	probe AP: visible / remembered	memory test
`lewm/`	`jepa.JEPA`	LeWM architecture and losses, retrained on UAV depth (ViT-tiny, 1-ch 48×80); body state added to the CLS token; 30k steps	0.070 / 0.058	0
`m0/`	`jepa.TokenJEPA`	pixels-only embedding; state is a predictor input and a separate next-state target; 15k steps	0.236 / 0.057	0
`m1/`	`jepa.TokenJEPA`	16 tokens per frame (CLS + 3×5 pooled patches)	0.208 / 0.054	0
`m2/`	`jepa.BeliefJEPA`	recurrent belief (closed-loop GRU, 1024-d) plus an 8-step open-loop action-driven predictor; 32-frame windows, batch 16. Frame encoder collapsed (rank 8.6)	percept 0.112 / 0.041, belief 0.065 / 0.059	E1 0.1–0.4 % of re-entry error
`m2s/`	`jepa.BeliefJEPA`	M2 with SIGReg on 128 frames pooled over batch × time (`loss.sigreg.samples=128`)	percept 0.164 / 0.053, belief 0.196 / 0.081	E1 4.7 → 1.1 %, 9–10 s.e.
`m3s/`	`jepa.BeliefJEPA`	M2s trained on `kin_v2` (look-away glances) with blind spans: p = 0.5 per window, 2–12 frames replaced by a learned blind vector	percept 0.186 / 0.056, belief 0.170 / 0.078	E1 31 → 15 %, 12–15 s.e.; reset +7–10 %
`m4s/`	`jepa.BeliefJEPA`	M3s with a map belief (`module.MapBelief`): 32×16×16 ego-centric latent map (0.6 m cells), warped by ego-motion, gated write	percept 0.224 / 0.058, belief 0.158 / 0.085	E1 58 → 0.8 %, decay 2 s
M3, M4	`jepa.BeliefJEPA`	M3 / M4 before the SIGReg fix	not published: same encoder collapse	—
DreamerV3 RSSM	—	reference (not in this repo)	0.342 / 0.105	—

lewm/ is not the official LeWM weights (quentinll/lewm). Those are trained on RGB control tasks and cannot read depth.

All models have about 18M parameters (M2, M2s and M3s about 26M, M4s 32M), one seed each (3072), and frameskip 5 at 20 Hz (0.25 s per frame). m3s and m4s were trained on kin_v2, all others on kin_v1; every model is evaluated on the same kin_v1 held-out episodes and twin scenarios. Probe numbers are average precision (AP) on held-out episodes.

Use

The files follow the official LeWM layout: each model folder has weights.pt (a state dict) and config.json (the resolved training config, the z-score stats and the training step). The model classes live in the code repo.

git clone -b lewm https://github.com/fuisl/single-agent-wm && cd single-agent-wm/lewm
pip install torch einops transformers hydra-core omegaconf huggingface_hub
hf download fuisloy/r-lewm-uav --include "m2s/*" --include "examples/*" --local-dir ../hf
python hf_example.py --model ../hf/m2s --data ../hf/examples     # training losses on 8 held-out episodes

from export_hf import load_pretrained, example_batch     # from lewm/
model, stats, cfg = load_pretrained("../hf/m2s")         # eval mode; device="cuda" for GPU
out = model.encode(example_batch(cfg, stats))            # out["emb"]: (B, T, N tokens, 192)
# belief models: b = model.belief.observe(emb, state, previous action) -> (B, T, 1024); see train.belief_forward

examples/ holds 8 held-out episodes as .npz files (goal, trap and explore flights). Each has:

depth (T, 48, 80), in metres, clipped to [0.6, 6.0]
state (T, 15): goal direction and distance, velocity, attitude, body rates
action (T, 5)
reward
occ / seen / vis: 64×64 ego-centric obstacle grids, used as probe labels
pos, yaw, and boxes (the world)

The model takes:

pixels: depth quantised to uint8, then scaled to [-1, 1] (utils.depth_to_pixels)
state: z-scored with stats["state"]
action: the 5 actions after each frame, concatenated (25-D) and z-scored with stats["action"]

WindowDataset in utils.py does all three.

Limitations

The simulator is a kinematic stand-in with analytic depth, not Isaac Sim.
Each model has one seed. The effects are 9–19 s.e. over 60 scenarios, but have not been replicated across seeds.
The encoders are not fully recovered: rank 47 (M2s) and 41 (M3s, M4s) against 68 for M1, below the pre-registered gate of 50.
M3s changes two things against M2s (data and blind spans); which one carries the gain is not yet separated.
The twin test's look-away manoeuvre resembles the kin_v2 glances, so M3s and M4s are tested nearer their training distribution than M2s is. The seen-vs-novel contrast is within one model, so this cannot create an effect, but it can inflate the comparison across models.
Zeroing M2s's belief at the turn-back lowers its re-entry error by ≈ 5 % (likely because training always starts the belief from zeros). With look-away data (M3s) this is gone: a reset hurts.
The probe runs on ordinary kin_v1 flight, which rarely needs memory, and does not see M3s's gain over M2s.
Closed-loop planning with these models has not been evaluated yet.
The "remembered" probe class is small, so absolute AP values are low. Compare against the frame-only and DreamerV3 rows, not against 1.0.

Downloads last month: -; Downloads are not tracked for this model. How to track

Inference Providers NEW

This model isn't deployed by any Inference Provider. 🙋 Ask for provider support