- 🌻Acknowledgement
- 🌟Overview
- 🔧Installation
- 📚World Knowledge Build
- 📉Model Training
- 🧐Evaluation
- 🚩Citation
Our code of training module is referenced and adapted from LLaMA-Factory, while the code of inference module is implemented based on ETO. Various baseline codes use ReAct, Reflexion, NAT, ETO. We use LangChain with open models via Fastchat. Thanks for their great contributions!
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 our 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:
- our instance-level task knowledge can generalize better to unseen tasks,
- weak WKM can guide strong agent model planning
- unified WKM training has promising potential for further development
git clone https://github.com/zjunlp/WKM
cd WKM
pip install -r requirements.txt
To build the task knowledge
python world_knowledge_build.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 world_knowledge_build.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 task_knowledge and state_knowledge, process the data to train format
python train_data_process.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 state_base_build.py \
--state_file_path your/state/knowledge/path \
--state_action_pair_path path/to/store/state_action/pair \
--vector_cache_path path/to/store/vector/cache
Our training data has been uploaded to huggingface.
Use LLama-Factory to train the agent model and world model
CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
--config_file ./examples/accelerate/single_config.yaml \
src/train_bash.py \
--ddp_timeout 180000000 \
--stage sft \
--do_train \
--model_name_or_path /base/model/path \
--dataset_dir ./data \
--dataset train_data_for_agent \
--template model_template \
--finetuning_type lora \
--lora_target q_proj,v_proj \
--output_dir ../lora/peft_model_name \
--overwrite_cache \
--per_device_train_batch_size 4\
--gradient_accumulation_steps 2 \
--lr_scheduler_type cosine \
--logging_steps 1 \
--save_steps 1000 \
--learning_rate 1e-4 \
--num_train_epochs 3 \
--plot_loss \
--fp16 \
--cutoff_len 2048 \
--save_safetensors False \
--overwrite_output_dir \
--train_on_prompt False
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
# world_knowledge_model 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 eval_agent.eto_multi_main_probs \
--agent_config fastchat \
--agent_model_name agent_model \
--world_model_name world_model \
--exp_config alfworld \
--exp_name eval \
--split test
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}
}
We will offer long-term maintenance to fix bugs and solve issues. So if you have any problems, please put issues to us.