vhgyvyu/ADWM-checkpoints

ADWM checkpoints

Trained world models and latent adapters for ADWM (paper), a latent diffusion world model for off-policy evaluation of LLM agents.

One folder per world model, named as in the code's configs/world_models.yaml:

Folder	Environment	Behavior data	Adapter trained with
hotpotqa	HotpotQA	ReAct-HotpotQA-SFT	ReAct-HotpotQA-PRM
scienceworld	ScienceWorld	sw-llama-sft	sw-llama-eto
alfworld	ALFWorld	Llama-3.1-8B-Instruct	leap-alf-iter3
alfworld-mixed	ALFWorld	Llama-3.1-8B-Instruct, leap-alf-iter0, leap-alf-iter2	leap-alf-iter3
webshop-mixed	WebShop	Llama-3.1-8B-Instruct, leap-ws-iter0, leap-ws-iter2	leap-ws-iter1

Each folder holds:

• wm_best.pt: the world model (model state dict, action_vocab, training args and env_config). Load it with adwm.models.checkpoint.load_checkpoint.
• adapter.pt: the latent-to-soft-token adapter G_psi used for continuation guidance (state_dict, d_z, d_llm, K).

Usage

From the root of the code repository:

hf download vhgyvyu/ADWM-checkpoints --local-dir checkpoints
hf download vhgyvyu/ADWM-data --repo-type dataset --local-dir data
python scripts/pipeline.py ope --launcher print   # commands for every evaluation cell

The behavior data and ground truth are in the companion dataset vhgyvyu/ADWM-data.

Citation

@article{liu2026adwm,
  title   = {Autoregressive Diffusion World Models for Off-Policy Evaluation of LLM Agents},
  author  = {Liu, Kaixuan and Xiong, Guojun and Zhang, Weinan and Tang, Shengpu},
  journal = {arXiv preprint arXiv:2606.05558},
  year    = {2026}
}