fuisloy/r-lewm-uav

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-22): lewm, m0, m1, m2 and m2s are here.
  • m2s is the current best model and the first with a clear memory effect.
  • m2 is kept for reference: its frame encoder collapsed (finding 7).
  • M3 and M4 were trained with the same flaw, so their weights are not published. They are being re-run with the fix (M3s, M4s).

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.

The collapse and the fix, at a glance. Left: effective rank of each model's frame embedding. Middle: what an obstacle probe reads out. Right: the memory effect against time out of view, as a share of re-entry error:

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 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.
M3, M4	jepa.BeliefJEPA	M2 on look-away data with blind spans; M3 with an ego-centric latent map	not published: same encoder collapse	to be re-run with the fix
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 and M2s about 26M), one seed each (3072), and frameskip 5 at 20 Hz (0.25 s per frame). 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 M2s effects are 9–10 s.e. over 60 scenarios, but have not been replicated across seeds.
• M2s's encoder is not fully recovered (rank 47 against 68 for M1; visible AP 0.164 against 0.208).
• Zeroing M2s's belief at the turn-back lowers its re-entry error by ≈ 5 %, though the wall memory is in the belief.
  • On ordinary flight, a reset costs one frame and then gives a ≈ 2 % advantage for about 1 s.
  • The likely cause is that training always starts the belief from zeros.
  • One-step prediction on this data uses little history. M3s / M4s, trained on look-away data, are in training.
• 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.