WKM

Agent Planning with World Knowledge Model

📄arXiv • 🌐Web • 𝕏 Blog • 🤗 HF • 🎧NotebookLM Audio

🌻Acknowledgement

Our code of the training module is referenced and adapted from LLaMA-Factory, while the code of the inference module is implemented based on ETO. Various baseline codes are sourced from ReAct, Reflexion, NAT, ETO. We use LangChain with open models via Fastchat. Thanks for their great contributions!

alt text

🌟Overview

Recent endeavors towards directly using large language models (LLMs) as agent models to execute interactive planning tasks have shown commendable results. Despite their achievements, however, they still struggle with brainless trial-and-error in global planning and generating hallucinatory actions in local planning due to their poor understanding of the "real" physical world. Imitating humans' world knowledge model which provides global prior knowledge before the task and maintains local dynamic knowledge during the task, in this paper, we introduce parametric World Knowledge Model (WKM) to facilitate agent planning. Concretely, we steer the agent model to self-synthesize knowledge from both expert and sampled trajectories. Then we develop WKM, providing prior task knowledge to guide the global planning and dynamic state knowledge to assist the local planning. Experimental results on three complex real-world simulated datasets with three state-of-the-art open-source LLMs, Mistral-7B, Gemma-7B, and Llama-3-8B, demonstrate that our method can achieve superior performance compared to various strong baselines. Besides, we analyze to illustrate that WKM can effectively alleviate the blind trial-and-error and hallucinatory action issues, providing strong support for the agent's understanding of the world. Other interesting findings include: 1) our instance-level task knowledge can generalize better to unseen tasks, 2) weak WKM can guide strong agent model planning 3) unified WKM training has promising potential for further development

🔧Installation

git clone https://github.com/zjunlp/WKM
cd WKM
pip install -r requirements.txt

📚World Knowledge Build

To build the task knowledge

python worldknowledgebuild.py \     --dataset_path your/rejected and chosen/data/pair \     --task your/task \     --gen task_knowledge \     --model_name your/model/name \     --output_path your/output/path

To build the state knowledge

python worldknowledgebuild.py \     --dataset_path your/rejected and chosen/data/pair \     --task your/task \     --gen state_knowledge \     --model_name your/model/name \     --output_path your/output/path

After your get taskknowledge and stateknowledge, process the data to train format

python traindataprocess.py \     --task alfworld \     --file_path your/path \     --mode model_type     --output_path your/output/path

And use the state knowledege train data to build state knowledge cache base

python statebasebuild.py \     --statefilepath your/state/knowledge/path \     --stateactionpairpath path/to/store/stateaction/pair \     --vectorcachepath path/to/store/vector/cache

Our training data has been uploaded to huggingface.

📉Model Training

Use LLama-Factory to train the agent model and world model

CUDAVISIBLEDEVICES=0,1,2,3 accelerate launch \     --configfile ./examples/accelerate/singleconfig.yaml \     src/train_bash.py \     --ddp_timeout 180000000 \     --stage sft \     --do_train \     --modelnameor_path /base/model/path \     --dataset_dir ./data \     --dataset traindatafor_agent \     --template model_template \     --finetuning_type lora \     --loratarget qproj,v_proj \     --outputdir ../lora/peftmodel_name \     --overwrite_cache \     --perdevicetrainbatchsize 4\     --gradientaccumulationsteps 2 \     --lrschedulertype cosine \     --logging_steps 1 \     --save_steps 1000 \     --learning_rate 1e-4 \     --numtrainepochs 3 \     --plot_loss \     --fp16 \     --cutoff_len 2048 \     --save_safetensors False \     --overwriteoutputdir \     --trainonprompt False

🧐Evaluation

To evaluate the task, you should first lanuch a local API server with fastchat. Our lora model adapter weights can be downloaded from here.

cd .src/eval agent_model api server
 python -u -m fastchat.serve.model_worker \     --model-path /path/peft/agent_model \     --port 21020 \      --worker-address http://localhost:21020 \     --max-gpu-memory 31GiB \     --dtype float16 worldknowledgemodel api server
 python -u -m fastchat.serve.model_worker \     --model-path /path/peft/world_model \     --port 21021 \      --worker-address http://localhost:21021 \     --max-gpu-memory 31GiB \     --dtype float16

Evaluate the task

python -m evalagent.etomultimainprobs \     --agent_config fastchat \     --agentmodelname agent_model \     --worldmodelname world_model \     --exp_config alfworld \     --exp_name eval \     --split test

🚩Citation

Please cite our repository if you use WKM in your work. Thanks!

@article{DBLP:journals/corr/abs-2405-14205,
  author       = {Shuofei Qiao and
                  Runnan Fang and
                  Ningyu Zhang and
                  Yuqi Zhu and
                  Xiang Chen and
                  Shumin Deng and
                  Yong Jiang and
                  Pengjun Xie and
                  Fei Huang and
                  Huajun Chen},
  title        = {Agent Planning with World Knowledge Model},
  journal      = {CoRR},
  volume       = {abs/2405.14205},
  year         = {2024},
  url          = {https://doi.org/10.48550/arXiv.2405.14205},
  doi          = {10.48550/ARXIV.2405.14205},
  eprinttype    = {arXiv},
  eprint       = {2405.14205},
  timestamp    = {Wed, 19 Jun 2024 08:52:49 +0200},
  biburl       = {https://dblp.org/rec/journals/corr/abs-2405-14205.bib},
  bibsource    = {dblp computer science bibliography, https://dblp.org}
}