latent-nav

A JEPA-style latent world model and a CEM planner for last-mile sidewalk navigation, plus the int8 + TorchScript work to run the planner on a Jetson Orin at 20 Hz. Built as an ML research intern at an early-stage last-mile delivery robotics startup, shadowing a senior ML engineer, over about three and a half weeks. This repository is the work: the env, the model, the planner, the training loop, the on-device deployment, and the lab notes that record how it actually went, dead-ends included.

The question the senior engineer set: the existing kinematic-MPC path tracker handles uneven sidewalk dynamics poorly. Can a learned latent world model plan better? The honest answer is no, not the way this was built. The model is accurate open-loop (3.3 cm), but CEM planning on top of it loses to the tuned analytic baseline (0.800 m vs 0.250 m RMSE) because the planner exploits model error off-distribution, a known model-based RL failure mode. The finding is characterized across three planner variants (from-scratch, warm-start, pessimistic-ensemble); all lose. The portable engineering wins are real and stand on their own: an 18.6x CEM-vectorization speedup and int8 PTQ at latent parity.

Results

docs/03-results.md is the single source of truth, generated by scripts/make_results.py from experiments/*/metrics.jsonl. Closed-loop numbers come from full runs (configs/default.yaml, configs/world_model.yaml); deploy numbers come from deploy/benchmark.py on the x86 workstation. The CI smoke run (configs/smoke.yaml) is directional only and is called out separately. Every number is sim or on-device output from this repo. There is no external leaderboard claim.

Open-loop model accuracy

Metric	Value
Position error	0.033 m
Final-step position error	0.039 m
Latent prediction MSE	0.0034

Rolling the latent model forward from a start state, decoded position error is 0.033 m (about 3.3 cm) over the eval horizon. The model's predictions are good. The gap that follows is in planning on top of the model, not in the model itself.

Portable engineering wins

Optimization	Before	After	Result
CEM vectorization (x86, portable)	0.27 FPS	5.1 FPS	18.6x

The headline speedup is 18.6x: replacing the per-sample Python CEM loop with one batched model.rollout over all 512 candidate plans takes the planner from 0.27 FPS to 5.1 FPS. It removes Python loop overhead, not a quantization trick, so it is portable and hardware-independent. Measured on x86 CPU via deploy/benchmark.py.

int8 vs fp32 parity	Value	Bar
Latent rank-preservation delta	-0.04 pp	within run-to-run noise
Latent prediction MSE delta	2.8%	inside noise

int8 PTQ (qnnpack) plus a TorchScript freeze keeps the latent predictions at parity: rank-preservation moves -0.04 percentage points and mean latent-prediction MSE changes by 2.8 percent. int8 here is a parity-preserving size and ARM-latency optimization, not an x86 speedup: on x86 the qnnpack int8 kernels are slower than fp32 (x86 wants the fbgemm backend), so the latency benefit is ARM/qnnpack-only, on the device. The Jetson Orin control loop targets 20 Hz device-side, where int8 plus qnnpack helps ARM latency; that absolute figure is a device measurement and is not reproducible on x86. Parity uses planning-correct latent metrics; an object-detection score has no meaning for a world model that predicts latent dynamics, so it is not reported.

Closed-loop navigation (sidewalk sim), 50 episodes, identical per-episode seeds

Metric	Kinematic-MPC baseline	Latent WM + CEM (warm-start)	Finding
Cross-track RMSE (m)	0.250	0.800	baseline lower (better) by 220%
Goal success rate	0.80	0.20	baseline higher
Collision rate	0.00	0.00	obstacle-free run

This run isolates the dynamics question on an obstacle-free path with actuator lag (motor_tau=0.30): can the learned model track the hidden lag better than the lag-blind kinematic-MPC baseline? Collision rate is 0.00 for both by construction, because there are no obstacles to hit, so the gap shows up in cross-track RMSE and goal success. The comparison is paired: each episode uses the same seed for both planners (eval/closed_loop.py), so the path and start pose are identical and the difference comes from the planner alone. The best learned planner tracks at 0.800 m RMSE versus the baseline's 0.250 m, about 3.2x worse. The baseline wins. This is the negative finding, characterized below, not a win for the world model.

Why the learned planner loses: planner-method comparison

Method	WM RMSE (m)	Baseline RMSE (m)	Success	Collision	motor_tau
From-scratch CEM	1.030	0.570	0.00	0.58	0.38
Warm-start CEM	0.800	0.250	0.20	0.00	0.30
Pessimistic-ensemble CEM	1.092	0.256	0.15	0.00	0.30

All three CEM variants over the learned model fall below the kinematic-MPC baseline in their respective regimes. The root cause is the same each time: the planner exploits model error off-distribution, steering into confident-but-wrong regions of the latent model. Warm-start (initialize the CEM mean from the analytic baseline's action sequence and refine with small std) is the best of the three and stays closest to the baseline, but does not catch it. Pessimistic ensemble scoring does not help, because the ensemble heads agree even off-distribution, so there is nothing for the pessimism to penalize.

How it went, week by week

The lab notes in docs/lab-notes/ are dated journal entries. They show the real path, including the week 2 dead-end, not a cleaned-up after-the-fact story.

• Week 0 (docs/lab-notes/2026-03-23-week0-reading-jepa.md): joined, shadowed the senior engineer, read the JEPA and latent-world-model literature (LeCun's JEPA position paper, Assran et al. I-JEPA 2023, Hafner et al. PlaNet/Dreamer, Hansen et al. TD-MPC2), wrote reading notes, and framed the problem. Notes in docs/01-literature-review.md, citations in docs/references.bib.
• Week 1 (docs/lab-notes/2026-03-30-week1-baseline-and-data.md): built the sidewalk sim and the data pipeline, implemented the kinematic-MPC baseline, and fixed honest metrics (cross-track RMSE, success rate, collision rate). Baseline: experiments/2026-03-30-baseline-mpc/.
• Week 2 (docs/lab-notes/2026-04-06-week2-world-model-and-bad-planning.md): implemented the JEPA-style world model (online encoder, ensemble predictor, EMA target encoder, latent predictive loss) and trained it. The first planning attempts were far worse than the baseline: the latent collapsed under a prediction-only loss, latent rollouts drifted, and the from-scratch CEM exploited model error over too long a horizon. The bad run is kept at experiments/2026-04-06-wm-cem-v1-bad/ (RMSE ~1.03 m vs a 0.57 m baseline at motor_tau=0.38, with collisions on most episodes).
• Week 3 (docs/lab-notes/2026-04-13-week3-fixing-the-planner.md): added the VICReg variance and covariance terms (anti-collapse), a multi-step rollout-consistency loss, and an ensemble-variance uncertainty penalty in the CEM cost so the planner stops steering into untrustworthy latent regions. Cut the effective horizon, recalibrated the cost weights, and added warm-start. These fixes removed the collapse and drift and made the model accurate open-loop at 3.3 cm, but they did not close the planning gap: the best learned planner still tracks at 0.800 m RMSE versus the baseline's 0.250 m. A pessimistic-ensemble variant was tried too and did not help. The CEM keeps exploiting model error off-distribution. This is the exploitation wall, characterized, not a win. Headline run: experiments/2026-04-13-wm-cem/.
• Week 4 (docs/lab-notes/2026-04-17-week4-edge-optimization.md): the headline is the planner speedup. Vectorizing the 512 CEM rollouts from a per-sample Python loop into one batched forward pass took the planner from 0.27 to 5.1 FPS, an 18.6x portable, hardware-independent gain. Applied dynamic int8 quantization and a TorchScript freeze at latent parity; the int8 latency benefit is ARM/qnnpack-only (slower on x86), so it is a size and device-latency optimization, not an x86 speedup. The Jetson Orin control loop targets 20 Hz device-side. Wrote a C++/libtorch inference shim for the ROS 2 control loop and ran the planner in sim and on the prototype robot. Bench: experiments/2026-04-17-jetson-bench/.

What this is built on

The world model is JEPA-style. It is not a new objective; it reuses the I-JEPA recipe (predict in representation space with an EMA target encoder) and VICReg anti-collapse terms, with a CEM planner on top in the PlaNet / TD-MPC2 line of plan-over-a-learned-model. Full citations in docs/references.bib; reading notes in docs/01-literature-review.md.

• LeCun, 2022. A Path Towards Autonomous Machine Intelligence (the JEPA position).
• Assran et al., CVPR 2023. I-JEPA. The online encoder, predictor, EMA target design used here.
• Hafner et al., ICML 2019 (PlaNet) and ICLR 2020 (Dreamer). CEM over a learned latent model.
• Hansen et al., ICLR 2024. TD-MPC2. Receding-horizon MPC over a learned latent world model, and the case for an uncertainty signal against model exploitation.
• Bardes et al., ICLR 2022. VICReg. The variance/covariance anti-collapse terms.
• Rubinstein, 1999. The cross-entropy method, the planner's optimizer.

Install

Targets Python 3.9.6. The code is device-agnostic: training/trainer.py auto-selects CUDA when a GPU is present and falls back to CPU otherwise. Training and eval for the headline numbers ran on an RTX 4090; the CPU wheel is enough for CI and the laptop smoke run. Every module starts with from __future__ import annotations so modern type-hint syntax works as annotations on 3.9.

# uv (preferred)
uv venv --python 3.9.6
uv pip install -e .

# or plain pip
python3.9 -m venv .venv && source .venv/bin/activate
pip install -e .

Pinned versions: torch 2.8.0 (CPU), numpy 2.0.2. See requirements.txt and pyproject.toml. The deploy C++ shim additionally needs libtorch (see deploy/cpp/CMakeLists.txt).

Quickstart

The fast path proves the whole pipeline links up on a laptop CPU in seconds using configs/smoke.yaml. The numbers it produces are directional only and are not the results above.

# smoke: collect, train, evaluate end-to-end on CPU in seconds
python scripts/collect_data.py      --config configs/smoke.yaml
python scripts/train_world_model.py --config configs/smoke.yaml
python scripts/run_baseline.py      --config configs/smoke.yaml
python scripts/evaluate.py          --config configs/smoke.yaml

# tests (CPU, tiny, seeded)
pytest -q

The full pipeline that produces the headline numbers (long, needs the full data budget):

python scripts/collect_data.py      --config configs/default.yaml
python scripts/train_world_model.py --config configs/world_model.yaml
python scripts/run_baseline.py      --config configs/default.yaml
python scripts/evaluate.py          --config configs/default.yaml
python scripts/quantize_and_bench.py --config configs/quant.yaml   # vectorization speedup + int8 parity
python scripts/make_results.py                                     # regenerate docs/03-results.md

Planner-method comparison (why the learned planner loses, three CEM variants over the same latent model; all below baseline, recorded in experiments/2026-04-14-ablations/):

python scripts/evaluate.py --config configs/world_model.yaml --planner from_scratch_cem
python scripts/evaluate.py --config configs/world_model.yaml --planner warm_start_cem
python scripts/evaluate.py --config configs/world_model.yaml --planner pessimistic_ensemble_cem

Make targets

make wraps the common flows. See the Makefile for the exact commands.

Target	Does
make install	Create the venv and install in editable mode
make smoke	Run the full pipeline on configs/smoke.yaml (CPU, seconds)
make test	Run the pytest suite
make lint	Run ruff (pinned to py39)
make baseline	Evaluate the kinematic-MPC baseline on configs/default.yaml
make train	Train the world model on configs/world_model.yaml
make evaluate	Closed-loop WM-vs-baseline comparison
make bench	int8 quantize, parity check, FPS benchmark
make results	Regenerate docs/03-results.md from experiments/*

Layout

src/latent_nav/
  config.py            typed config dataclasses, yaml loader, global seeding
  envs/                unicycle sidewalk env, obstacle field, lidar raycast
  data/                transition collection and the multi-step sequence dataset
  models/              encoder, ensemble predictor, JEPA world model, losses
  planning/            CEM planner, cost terms, kinematic-MPC baseline
  training/            trainer (EMA, checkpoint/resume, early stop), LR schedule, logging
  eval/                closed-loop episode runner and WM-vs-baseline comparison
  deploy/              int8 PTQ, FPS/parity benchmark, ROS 2 shim
  metrics.py           cross-track RMSE, success rate, collision rate
deploy/cpp/            libtorch inference shim and CMake build for the ROS 2 loop
scripts/               argparse CLIs (each takes --config PATH)
configs/               yaml configs mirroring the dataclasses
experiments/           committed run summaries (metrics.jsonl + config snapshots)
docs/                  literature review, architecture, results, dated lab notes
tests/                 pytest, CPU, tiny, seeded

Architecture and design reasoning: docs/02-architecture.md. The key design choices, in one line each: predict in latent space to skip task-irrelevant pixel and lidar detail; EMA target plus VICReg to stop representation collapse; an ensemble predictor whose head disagreement penalizes the planner for steering into untrustworthy latent regions (it curbs the exploitation but does not eliminate it, and the baseline still wins); an explicit CEM planner (debuggable, no extra training instability); a short horizon with replanning every step because latent rollouts drift.

Honesty notes

• The world model is JEPA-style and cites the papers it copies. No claim of inventing JEPA.
• The closed-loop result is a negative finding: the learned planner loses to the analytic baseline (0.800 m vs 0.250 m RMSE in sim), characterized across three CEM variants. It is presented as a characterized result, not a WM win. Closed-loop quality was measured in sim and in field runs (week 4), not at deployment scale.
• The 18.6x speedup is a portable, hardware-independent CEM-vectorization gain measured on x86. int8 is a parity result (rank -0.04 pp, latent-MSE 2.8%); the int8 latency benefit is ARM/qnnpack-only and is slower on x86, so it is not claimed as an x86 speedup. The Jetson Orin 20 Hz figure is a device-side target.
• Parity uses planning-correct latent metrics (a rank-preservation proxy and latent MSE delta). The proxy is labeled as a proxy, not as closed-loop success. An object-detection score is deliberately not used; it does not fit a world model.
• Closed-loop numbers come from the full configs; deploy numbers from x86. CI runs the smoke config, which is directional only and is not comparable. docs/03-results.md states this plainly.
• The week 2 dead-end (collapse, then planner-exploits-model) is preserved in the dated lab notes and in the committed bad run (experiments/2026-04-06-wm-cem-v1-bad/, the prediction-only loss and the no-uncertainty slow CEM described there). It is evidence the fixes were earned, not assumed.

Limitations

• Simulation, not real concrete. The dynamics are a noisy unicycle, single seed family (paired per-episode seeds, seed 0), 50 episodes. Not a large external benchmark, and not enough seeds for confidence intervals on the deltas.
• The field runs are a sanity check, not a field study.
• The world model is JEPA-style: it adapts published ideas (LeCun 2022; Assran et al. 2023; Hafner et al.; Hansen et al. TD-MPC2), not a new method.
• The uncertainty penalty uses ensemble head variance as an epistemic proxy. An ensemble of 4 is cheap and uncalibrated; it curbs the planner exploiting model error but does not eliminate it, and the baseline still wins.
• The deploy numbers are the portable CEM-vectorization speedup and int8 latent parity, measured on x86. The Jetson Orin figure is a device-side target, not an x86-reproducible number. Quantization is post-training only; no QAT.

Full list with the reasoning is in docs/03-results.md.

What I'd do next

• Constrain the planner harder to in-distribution actions (a trust region or a behavior-cloning prior around the analytic baseline) so CEM cannot walk the rollout off-distribution in the first place.
• Calibrate the epistemic uncertainty (for example with a small held-out calibration set) so the pessimism penalty has signal off-distribution, instead of raw head variance.
• Add a longer-horizon latent value estimate so the receding-horizon planner is less myopic on tight obstacle fields (TD-MPC2-style value bootstrap).
• Collect on-robot transitions and fine-tune the encoder so the sim-to-real gap is measured rather than assumed.
• Multiple seeds and confidence intervals on every reported number.

License

MIT. See LICENSE.