Equilibrium Planner

May 5, 2025 · View on GitHub

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

framework

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):

    table1

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

    table2

  • Effectiveness of different types of feedback:

    table3

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