File size: 5,323 Bytes
d5b1c6f af55769 d5b1c6f af55769 01a9e74 af55769 01a9e74 af55769 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | ---
license: other
license_name: mixed-apache-2.0-and-mit
tags:
- robotics
- world-model
- model-based-control
- jepa
- reinforcement-learning
- latent-planning
- iterative-control
---
# LePlanner: Amortized Iterative Control in Frozen Latent World Models
This repository contains the trained controllers, frozen world models, latent caches, and evaluation results for **LePlanner**:
> **LePlanner: Amortized Iterative Control in Frozen Latent World Models**
> [Authors], [Year].
> [Paper link — arXiv / venue]
LePlanner is an amortized iterative controller that plans through a **frozen JEPA world model** (LeWM) over a small number of learned refinement steps, trained with an **arrival–hold** objective. It matches or exceeds test-time search planners (CEM) while requiring an order-of-magnitude fewer predictor invocations and 8–15× lower wall-clock per decision.
## Results
| Environment | Method | Success | Notes |
|---|---|---|---|
| PushT | LePlanner (rh=1) | **98%** | arrival+hold checkpoint |
| PushT | LePlanner (rh=5) | **92%** | |
| PushT | CEM (rh=5) | 90% | baseline search planner |
| Reacher | LePlanner (rh=1) | **100%** | 50 episodes, seed 42 |
| TwoRooms | LePlanner (rh=1) | **100%** | 50 episodes, seed 42 |
LePlanner uses **~16× fewer predictor rows per decision** than CEM and **8–15× lower wall-clock** per decision.
## Quick Start
```python
from huggingface_hub import hf_hub_download
# Load a controller checkpoint
ckpt_path = hf_hub_download(
"SaltedLemon/lejepa-control-pusht",
"checkpoints/environment_controllers/reacher/controller.pt",
repo_type="model",
)
# Load a frozen world model
wm_path = hf_hub_download(
"SaltedLemon/lejepa-control-pusht",
"checkpoints/base_world_models/reacher/weights.pt",
repo_type="model",
)
```
The training code is available at [`github.com/SaltedLemon/lejepa_control`](https://github.com/SaltedLemon/lejepa_control).
## Repository Structure
```
checkpoints/
├── environment_controllers/ # Trained LePlanner controllers (per-environment)
│ ├── reacher/
│ │ ├── controller.pt # 100% success, rh=1
│ │ └── density.pt # Behavior density model
│ └── tworooms/
│ ├── controller.pt # 100% success, rh=1
│ └── density.pt
├── controller/ # PushT iterative controller (phase-1)
├── ablations/ # PushT training ablations
│ └── ah_hold0.5/
│ └── controller.pt # Headline PushT checkpoint (94-96%)
├── base_world_models/ # Frozen LeWM world models
│ ├── pusht/ # quentinll/lewm-pusht (MIT)
│ ├── reacher/ # quentinll/lewm-reacher (MIT)
│ └── tworooms/ # quentinll/lewm-tworooms (MIT)
├── planner_*/ # Recursive & cross-attention planners
├── manifold_transfer/ # Latent transfer experiments
└── density/ # PushT behavior density model
code/ # Training & evaluation scripts
├── planner.py # LePlanner controller
├── solver.py # Training/eval harness
├── losses.py # Arrival-hold, action-Gaussian, support losses
└── scripts/ # train_planner.py, eval_planner.py, ...
latents/ # Pre-computed latent caches
├── reacher/ # Latents, actions, episode metadata
└── tworoom_cls/
results/ # Raw evaluation logs
├── eval/results.jsonl # PushT eval across planners
└── eval_planner/results.jsonl
```
## Key Checkpoints
| File | Description |
|---|---|
| `checkpoints/ablations/ah_hold0.5/controller.pt` | **Headline PushT controller** — 98% at rh=1, 92% at rh=5 |
| `checkpoints/environment_controllers/reacher/controller.pt` | Reacher controller — 100% |
| `checkpoints/environment_controllers/tworooms/controller.pt` | TwoRooms controller — 100% |
| `checkpoints/base_world_models/pusht/weights.pt` | Frozen PushT world model (LeWM) |
| `checkpoints/base_world_models/reacher/weights.pt` | Frozen Reacher world model (LeWM) |
| `checkpoints/base_world_models/tworooms/weights.pt` | Frozen TwoRooms world model (LeWM) |
## Training & Evaluation
The controller and planner checkpoints are trained in this project; the LeWM world model is **always frozen** while those planners are trained.
Training code: [`github.com/SaltedLemon/lejepa_control`](https://github.com/SaltedLemon/lejepa_control)
Each controller checkpoint embeds its training arguments, action statistics, step, and validation profile.
## Citation
If you use this code or these checkpoints, please cite:
```bibtex
@article{leplanner,
title = {LePlanner: Amortized Iterative Control in Frozen Latent World Models},
author = {[Authors]},
journal = {[Venue / arXiv]},
year = {[Year]},
url = {https://huggingface.co/SaltedLemon/lejepa-control-pusht}
}
```
## License
- **Controller/planner code and checkpoints:** as specified in the linked `lejepa_control` repository.
- **LeWM world model weights** (`base_world_models/pusht`, `reacher`, `tworooms`): MIT licensed upstream from `quentinll/lewm-*`.
|