SaltedLemon/lejepa-control-pusht

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)	94–96%	arrival+hold checkpoint
PushT	LePlanner (rh=5)	88–90%
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

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.

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 — 94–6% at rh=1
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

Each controller checkpoint embeds its training arguments, action statistics, step, and validation profile.

Citation

If you use this code or these checkpoints, please cite:

@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-*.