Simplify README: clear per-env table, download code, clean structure

README.md

@@ -7,126 +7,73 @@ tags:
 - model-based-control
 - jepa
 - reinforcement-learning
-- latent-planning
-- iterative-control
 ---
 
-# LePlanner: Amortized Iterative Control in Frozen Latent World Models
+# LePlanner
 
-This repository contains the trained controllers, frozen world models, latent caches, and evaluation results for **LePlanner**:
+Official checkpoints for **LePlanner: Amortized Iterative Control in Frozen Latent World Models**.
 
-> **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.
+[[Paper](link-to-paper)] [[Code](https://github.com/SaltedLemon/lejepa_control)]
 
 ## 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.
+| Environment | Success | Checkpoint |
+|---|---|---|
+| PushT | 98% (rh=1) | `checkpoints/ablations/ah_hold0.5/controller.pt` |
+| Reacher | 100% | `checkpoints/environment_controllers/reacher/controller.pt` |
+| TwoRooms | 100% | `checkpoints/environment_controllers/tworooms/controller.pt` |
 
-## Quick Start
+## Download a Controller
 
 ```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",
+# Choose your environment: "pusht", "reacher", or "tworooms"
+env = "reacher"
+
+# Download controller
+controller_path = hf_hub_download(
+    "nottygian/Leplanner",
+    f"checkpoints/environment_controllers/{env}/controller.pt",
     repo_type="model",
 )
 
-# Load a frozen world model
+# Download matching world model
 wm_path = hf_hub_download(
-    "SaltedLemon/lejepa-control-pusht",
-    "checkpoints/base_world_models/reacher/weights.pt",
+    "nottygian/Leplanner",
+    f"checkpoints/base_world_models/{env}/weights.pt",
     repo_type="model",
 )
 ```
 
-The training code is available at [`github.com/SaltedLemon/lejepa_control`](https://github.com/SaltedLemon/lejepa_control).
+> **Note:** For PushT, the best controller is at `checkpoints/ablations/ah_hold0.5/controller.pt`.
 
 ## 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
+├── environment_controllers/    # Trained LePlanner controllers
+│   ├── reacher/controller.pt   # 100% success
+│   └── tworooms/controller.pt  # 100% success
+├── ablations/
+│   └── ah_hold0.5/controller.pt  # PushT: 98% success
+├── base_world_models/          # Frozen LeWM world models
+│   ├── pusht/weights.pt
+│   ├── reacher/weights.pt
+│   └── tworooms/weights.pt
+└── density/                    # Behavior density models
+
+code/                           # Training & evaluation scripts
+results/eval/results.jsonl     # Raw evaluation logs
 ```
 
-## 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]},
+  journal = {[Venue]},
   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-*`.