Equilibrium Planner

Code for ICML 2025 paper Closed-Loop Long-Horizon Robotic Planning via Equilibrium Sequence Modeling.

Requirements

• Linux with python ≥ 3.9 and CUDA ≥ 11.8.

• Install torch ≥ 2.3.0, transformers ≥ 4.43.1, flash-attn ≥ 2.6.1 and bitsandbytes ≥ 0.43.0.

• Prepare base model:

  • Download Meta-Llama-3-8B-Instruct to your MODEL_DIR in utils/predefine.py.

  • Move the customized llama model to MODEL_DIR/Meta-Llama-3-8B-Instruct and overwrite the config file:

    mv llama3/* MODEL_DIR/Meta-Llama-3-8B-Instruct/

• Prepare dataset:

  • Download VirtualHome-Env and unzip it to your RAW_DATA_DIR in utils/predefine.py.

Getting Started

• Preprocess dataset:

  python prepare_dataset.py

• Training for equilibrium planner:

  CUDA_VISIBLE_DEVICES=0 python run_finetune.py

• Training for world model:

  CUDA_VISIBLE_DEVICES=0 python run_finetune_worldmodel.py

• Generation for equilibrium planner in different settings:

  • Equilibrium planner with environmental feedback:

    CUDA_VISIBLE_DEVICES=0 python run_generate.py

  • Equilibrium planner with world model feedback:

    CUDA_VISIBLE_DEVICES=0 python run_generate.py --world_model_feedback

  • Equilibrium planner with both environmental and world model feedback:

    CUDA_VISIBLE_DEVICES=0 python run_generate.py --mixed_feedback

• Evaluation in the VirtualHome graph simulator:

  python evaluation.py

Results

• Performance on VirtualHome-Env without correction (only world model feedback):

  

• Performance on VirtualHome-Env with up to 10 corrections (both environmental and world model feedback):

  

• Effectiveness of different types of feedback:

  

Citation

If you find this code useful, please consider citing:

@inproceedings{li2025closed,
  title={Closed-Loop Long-Horizon Robotic Planning via Equilibrium Sequence Modeling},
  author={Li, Jinghan and Sun, Zhicheng and Mu, Yadong},
  booktitle={International Conference on Machine Learning},
  year={2025}
}

Acknowledgement

• Our implementation in llama3 is based on transformers and implementation in evolving_graph is based on virtualhome-1.0.0.