zpointsun / DiBO-AntMorphology

DiBO Ant Morphology

Final task-specific DiBO checkpoint for AntMorphology-Exact-v0 , released with Training Diffusion Language Models for Black-Box Optimization (ICML 2026 Spotlight). The paper is also available on the Hugging Face paper page .

DiBO is implemented in the DiBO code repository . This artifact is the final checkpoint after the DA, SFT, and RL stages.

Artifact

• Checkpoint: dibo_ant_morphology_final.pt
• Format: PyTorch checkpoint dictionary with the model state dict at model
• Base model: GSAI-ML/LLaDA-8B-Instruct , revision 08b83a6feb34df1a6011b80c3c00c7563e963b07

This repository does not mirror the base model as a standalone Transformers model directory. Initialize the pinned base model with the DiBO loader, which installs the required delimiter tokens and resizes its embeddings, then load this checkpoint's complete DiBO model state dict. It is not compatible with a generic AutoModel.from_pretrained() call on this repository.

Download and load

From a checkout of the released DiBO code:

hf download zpointsun/DiBO-AntMorphology dibo_ant_morphology_final.pt \
  --revision v1.0.0 --local-dir checkpoints/ant

import torch
from src.model.dllm import DEFAULT_MODEL_ID, LLADA_MODEL_REVISION, load_model_and_tokenizer

assert DEFAULT_MODEL_ID == "GSAI-ML/LLaDA-8B-Instruct"
assert LLADA_MODEL_REVISION == "08b83a6feb34df1a6011b80c3c00c7563e963b07"

model, tokenizer = load_model_and_tokenizer(DEFAULT_MODEL_ID, device="cuda")
checkpoint = torch.load("checkpoints/ant/dibo_ant_morphology_final.pt", map_location="cuda")
model.load_state_dict(checkpoint["model"], strict=True)
model.eval()

load_model_and_tokenizer adds |design-start| , |design-end| , |label-start| , and |label-end| before the checkpoint is loaded. These tokens and the embedding resize are required for compatible loading.

For task evaluation, install the oracle environment described in the installation guide and use the released evaluation entry point, for example:

python eval.py --tasks AntMorphology-Exact-v0 \
  --checkpoint_path checkpoints/ant/dibo_ant_morphology_final.pt \
  --seeds <SEED> --max_attempts 1000

Limitations

The checkpoint requires the pinned LLaDA base model, the DiBO tokenizer setup, and a CUDA-capable PyTorch environment for practical inference. Direct morphology evaluation additionally requires the Design-Bench data cache and the MuJoCo-based oracle environment described in the DiBO installation guide.

Related DiBO checkpoints

• TFBind8
• TFBind10
• D'Kitty Morphology

Citation

If you find DiBO helpful, please cite:

@article{sun2026training,
  title={Training diffusion language models for black-box optimization},
  author={Sun, Zipeng and Chen, Can and Yuan, Ye and Wu, Haolun and Gu, Jiayao and Pal, Christopher and Liu, Xue},
  journal={arXiv preprint arXiv:2603.17919},
  year={2026}
}