latent-memory-palace

KO | EN

gitlite — search

#video #proxy #web #agents #animation #machine-learning #devops #kotlin #redis #github

latent-memory-palace

(★ 9)

Official codebase for Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference

File Explorer

ZIP Download

- - __init__.py
  - autoregressive_policy.py
  - eval.py
  - train.py
- - __init__.py
  - diffusion_policy.py
  - eval.py
  - train.py
- - __init__.py
  - eval_policy.py
  - lmp_buffers.py
  - lmp_policy.py
  - lmp_vae.py
  - lmp_wam.py
  - train_policy.py
  - train_vae.py
  - train_wam.py
  - utils.py
- - __init__.py
  - train.py
  - vae_policy.py
- - __init__.py
  - train.py
  - vqbet.py
- callback.py
- file_utils.py
- optimization.py
- torch_utils.py
- train_utils.py
- video_utils.py
- websocket_client.py
- websocket_server.py
- - d3il_aligning.yaml
  - d3il_avoiding.yaml
  - d3il_sorting.yaml
  - d3il_sorting4.yaml
  - d3il_sorting6.yaml
  - d3il_stacking.yaml
  - droid_rlds.yaml
  - droid_rlds_webvideo.yaml
  - droid_zarr.yaml
  - libero_90.yaml
  - libero_all_actions.yaml
  - libero_book_caddy.yaml
  - libero_bowl_drawer.yaml
  - libero_moka_moka.yaml
  - libero_mug_mug.yaml
  - libero_soup_cheese.yaml
  - pointmaze_large.yaml
  - pointmaze_medium.yaml
  - pusht.yaml
  - robomimic_all_actions.yaml
  - robomimic_can.yaml
  - robomimic_lift.yaml
  - robomimic_square.yaml
  - robomimic_tool_hang.yaml
  - robomimic_transport.yaml
  - spiral.yaml
- - fast.yaml
  - lmp.yaml
  - oat.yaml
  - vqvae.yaml
- ap_fast.yaml
- ap_lmp.yaml
- ap_oat.yaml
- ap_vqvae.yaml
- default_optimizer.yaml
- default_wandb.yaml
- dp.yaml
- dp_droid.yaml
- dp_droid_finetune.yaml
- dp_libero.yaml
- lmp_policy.yaml
- lmp_policy_droid.yaml
- lmp_policy_droid_finetune.yaml
- lmp_policy_libero.yaml
- lmp_vae.yaml
- lmp_wam.yaml
- lmp_wam_libero.yaml
- tokenizer.yaml
- vae.yaml
- vqbet.yaml
- - __init__.py
  - compute_normalizer_stats.py
  - convert_zarr_dataset.py
  - droid_rlds.py
  - droid_zarr.py
  - utils.py
- - data_loader.py
  - normalizer.py
  - obs_utils.py
  - sampler.py
  - wrappers.py
- __init__.py
- d3il.py
- dataset.py
- libero.py
- ogbench.py
- pusht.py
- robomimic.py
- spiral.py
- __init__.py
- libero.py
- ogbench.py
- pusht.py
- robomimic.py
- wrappers.py
- - __init__.py
  - base.py
  - train.py
- adaln_attention.py
- attention.py
- language.py
- obs_encoders.py
- pos_embeds.py
- transforms.py
- utils.py
- vision.py
.gitignore
pyproject.toml
README.md
requirements.txt

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

🌐 [Website](https://weirdlabuw.github.io/lmp/) &nbsp;•&nbsp; 📄 [Paper](https://arxiv.org/abs/2607.08724) &nbsp;•&nbsp; 🎬 [Talk](https://www.youtube.com/watch?v=56GeTq7IPhg)

Latent Memory Palace (LMP) instantiates iterative and adaptive reasoning for control policies through variational inference with an autoregressive latent distribution. This codebase provides the official implementation of the policy (LMP-π) and the action tokenizer (LMP-tok).

## 🗂️ Code Structure
The codebase is structured as follows:

| Directory | Description |
|-----------|-------------|
| `algorithms/latent_memory_palace/` | LMP policy (LMP-π) and action tokenizer (LMP-tok) |
| `configs/` | Experiment configs |
| `datasets/` | Dataset wrappers for DROID, LIBERO, D3IL, and Robomimic |
| `environments/` | Environment wrappers for DROID, LIBERO, D3IL, and Robomimic |
| `models/` | Shared model components (e.g. observation encoders) |

## 🛠️ Setup
Install the repository as a package: 
```bash
pip install -e .
```

Install simulated environments:
```bash
# LIBERO
pip install git+https://github.com/zchuning/LIBERO.git@1.5.1
# D3IL (lightweight fork with gym interface)
pip install git+https://github.com/zchuning/d3il_lite.git
# Robomimic
pip install git+https://github.com/zchuning/robomimic.git@0.4.0
```

To set up [DROID](https://droid-dataset.github.io/), follow [these instructions](https://droid-dataset.github.io/droid/) to configure the robot and download the dataset. Then, install the following package to enable training on RLDS:
```bash
pip install git+https://github.com/zchuning/dlimp.git
```

## 🤗 Pretrained Checkpoints
We release the following pretrained checkpoints on Hugging Face Hub. Each repo contains `config.json` and `models.pt`. The policy checkpoints can be used for evaluation or finetuning. The tokenizer checkpoints can be used for training downstream autoregressive policies. Check out the instructions in the following sections for usage commands.

| Checkpoint | Repo |
|------------|------|
| DROID policy | [`zchuning/lmp-droid`](https://huggingface.co/zchuning/lmp-droid) |
| LIBERO-90 policy | [`zchuning/lmp-libero`](https://huggingface.co/zchuning/lmp-libero) |
| Robomimic action tokenizer | [`zchuning/lmp-tokenizer-robomimic`](https://huggingface.co/zchuning/lmp-tokenizer-robomimic) |
| LIBERO action tokenizer | [`zchuning/lmp-tokenizer-libero`](https://huggingface.co/zchuning/lmp-tokenizer-libero) |

To download a checkpoint:
```bash
hf download zchuning/lmp-droid --local-dir ckpts/lmp-droid
```

## 🚀 Training

We use [Hydra](https://hydra.cc/) to configure training scripts. Use `--config-name <config name>` to override the default config file (located in `configs`) and `<key>=<value>` to override individual parameters.

### Single-Task Robomimic / D3IL Policy
To train a single-task policy (e.g. robomimic_square), edit `configs/task/dataset/robomimic_square.yaml` to point to the data directory and target Zarr buffer directory. Then run:
```bash
torchrun --standalone --nproc_per_node 4 \
    algorithms/latent_memory_palace/train_policy.py \
    task=robomimic_square
```

### Multi-Task LIBERO Policy
To train a language-conditioned multitask policy on LIBERO, edit `configs/task/dataset/libero_90.yaml` to point to the data directory and target Zarr buffer directory. Then run training with `lmp_policy_libero.yaml`:
```bash
torchrun --nnodes 4 --nproc_per_node 8 \
    algorithms/latent_memory_palace/train_policy.py \
    --config-name lmp_policy_libero.yaml
```

### Multi-Task DROID Policy
To train a language-conditioned multitask policy on DROID, edit `configs/task/dataset/droid_rlds.yaml` to point to the data directory. Then run training with `lmp_policy_droid.yaml`:
```bash
torchrun --nnodes 4 --nproc_per_node 8 \
    algorithms/latent_memory_palace/train_policy.py \
    --config-name lmp_policy_droid.yaml
```

To finetune the pretrained DROID policy, collect expert demonstrations using the DROID interface and convert them into an RLDS dataset. Then create a new task config (e.g. `configs/task/droid_clean_table.yaml`) and run:
```bash
torchrun --nnodes 4 --nproc_per_node 8 \
    algorithms/latent_memory_palace/train_policy.py \
    --config-name lmp_policy_droid_finetune.yaml \
    pretrain_checkpoint_path=path/to/pretrained/checkpoint \
    task=droid_clean_table
```

### Action Tokenizer & Autoregressive Policy

To train an action tokenizer (e.g. on the combined robomimic dataset), edit `configs/task/dataset/robomimic_all_actions.yaml` to point to the data directory and target Zarr buffer directory, then run:

```bash
torchrun --standalone --nproc_per_node 4 \
    algorithms/latent_memory_palace/train_vae.py \
    task=robomimic_all_actions
```

To train an autoregressive policy using the pretrained tokenizer, run:
```bash
torchrun --standalone --nproc_per_node 4 \
    algorithms/autoregressive_policy/train.py \
    --config-name ap_lmp.yaml \
    task=robomimic_square \
    tokenizer.ckpt_dir=/path/to/vae/checkpoint
```

## 📊 Evaluation
We provide evaluation scripts for simulated environments and the real-world DROID platform.

### Evaluating on Simulated Environments
Use `algorithms/latent_memory_palace/eval_policy.py` to generate rollout videos alongside the number of reasoning steps:
```bash
python algorithms/latent_memory_palace/eval_policy.py \
    --logdir /path/to/checkpoint \
    --task libero_90 \
    --temperature 1.0 \
    --batch_size 64
```
By default, this evaluates a LIBERO multitask policy on the `libero-bowl-bowl` task. To evaluate on different tasks, edit `hdf5_path` in `configs/task/libero_90.yaml` to point to the new eval task. This edit is not required for single-task policies.

### Evaluating on DROID Platform
To evaluate a DROID policy, serve the policy on the GPU/server machine and drive the robot from the DROID machine over websockets.

On the server machine, start the policy server:
```bash
python common/websocket_server.py --logdir ckpts/lmp-droid --port 8000
```

On the robot machine, connect and run episodes:
```bash
python common/websocket_client.py \
    --server_host <MODEL_MACHINE_IP> --server_port 8000 \
    --instruction "pick up the block and put it in the bowl" \
    --right_controller
```
Control: press **A** to start an episode, **B** to delete the last episode. During an episode, hold **trigger** to enable robot motion, press **A** to mark success, **B** to mark failure. Episode videos are saved under `runs/<timestamp>/`.

## 📈 Baselines
Use the following commands to train baseline policies:
```bash
# Diffusion Policy
torchrun --standalone --nproc_per_node 4 algorithms/diffusion_policy/train.py task=robomimic_square

# VQ-BeT
torchrun --standalone --nproc_per_node 4 algorithms/vqbet/train.py task=robomimic_square

# Gaussian VAE
torchrun --standalone --nproc_per_node 4 algorithms/vae_policy/train.py task=robomimic_square
```

## ⚙️ Key Hyperparameters

| Parameter | Description |
|-----------|-------------|
| `model.seq_len` | Max latent sequence length |
| `model.vocab_size` | Discrete latent vocabulary size |
| `model.sigma_max/min` | Decoder noise schedule range |
| `clip_range` | PPO-style clip range (applied exponentially per step) |
| `free_nats_ratio` | Minimum normalized KL per timestep |
| `alpha` | Prior mixing: `alpha * learned_prior + (1-alpha) * uniform` |
| `kl_coef` / `ll_coef` | Loss term weights |

## 📝 Bibtex
If you find this code useful, please cite:

```
@misc{zhu2026lmp,
    author    = {Zhu, Chuning and Xu, Eva and Barreiros, Jose and Srinivasan, Krishnan and Shah, Paarth and Gupta, Abhishek},
    title     = {Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference},
    year      = {2026},
}
```

latent-memory-palace

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

latent-memory-palace

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

Download README