概述
本仓库包含了Motif的PyTorch代码,用于在NetHack游戏中训练AI代理,其奖励函数源自大语言模型的偏好。
作者:Martin Klissarov* & Pierluca D'Oro*, Shagun Sodhani, Roberta Raileanu, Pierre-Luc Bacon, Pascal Vincent, Amy Zhang 和 Mikael Henaff
Motif从NetHack游戏交互数据集中获取带有说明的观察对,引出大语言模型(LLM)对这些对的偏好。它自动地将LLM的常识提炼为奖励函数,用于通过强化学习训练代理。
为便于比较,我们在pickle文件motif_results.pkl
中提供了训练曲线,该文件包含以任务为键的字典。对于每个任务,我们提供了Motif和基线方法在多个种子下的时间步和平均回报列表。
如下图所示,Motif包含三个阶段:
- 数据集标注:利用LLM对带说明的观察对的偏好创建标注后的数据集;
- 奖励训练:使用标注后的数据集对和LLM的偏好作为监督信号训练奖励函数;
- 强化学习训练:使用Motif的奖励函数训练代理。
我们通过提供必要的数据集、命令和原始结果,详细说明了每个阶段,以便重现论文中的实验。
我们通过NetHack学习环境评估Motif在具有挑战性、开放式和程序生成的NetHack游戏中的表现。我们研究了Motif如何主要生成符合人类直觉的行为,这些行为可以通过提示修改轻松调整,以及其扩展性。
要安装整个流程所需的依赖项,只需运行pip install -r requirements.txt
。
使用Llama 2进行数据集标注
在第一阶段,我们使用一个带有说明(即游戏中的消息)的观察对数据集,这些观察对是由经过强化学习训练以最大化游戏分数的代理收集的。
我们在本仓库中提供了该数据集。
我们将不同部分存储在motif_dataset_zipped
目录中,可以使用以下命令解压缩。
cat motif_dataset_zipped/motif_dataset_part_* > motif_dataset.zip; unzip motif_dataset.zip; rm motif_dataset.zip
我们提供的数据集包含了Llama 2模型给出的一组偏好,存储在preference/
目录中,使用了论文中描述的不同提示。
包含标注的.npy
文件名遵循模板llama{size}b_msg_{instruction}_{version}
,其中size
是来自集合{7,13,70}
的LLM大小,instruction
是引入给LLM的提示中的指令,来自集合{defaultgoal, zeroknowledge, combat, gold, stairs}
,version
是要使用的提示模板版本,来自集合{default, reworded}
。
以下是可用标注的摘要:
标注 | 论文中的用例 |
---|---|
llama70b_msg_defaultgoal_default | 主要实验 |
llama70b_msg_combat_default | 引导向_怪物杀手_行为 |
llama70b_msg_gold_default | 引导向_黄金收集者_行为 |
llama70b_msg_stairs_default | 引导向_下降者_行为 |
llama7b_msg_defaultgoal_default | 扩展实验 |
llama13b_msg_defaultgoal_default | 扩展实验 |
llama70b_msg_zeroknowledge_default | 零知识提示实验 |
llama70b_msg_defaultgoal_reworded | 提示重写实验 |
为创建标注,我们使用vLLM和Llama 2的聊天版本。如果你想生成自己的标注或重现我们的标注过程,请确保按照官方说明下载模型(获取模型权重可能需要几天时间)。
标注脚本假设数据集将使用n-annotation-chunks
参数分成不同的块进行标注。这允许根据可用资源进行并行处理,并且对重启/抢占具有鲁棒性。要使用单个块运行(即处理整个数据集),并使用默认提示模板和任务规范进行标注,请运行以下命令。
python -m scripts.annotate_pairs_dataset --directory motif_dataset \
--prompt-version default --goal-key defaultgoal \
--n-annotation-chunks 1 --chunk-number 0 \
--llm-size 70 --num-gpus 8
请注意,默认行为是通过将标注附加到指定配置的文件来恢复标注过程,除非通过--ignore-existing
标志另有说明。也可以使用--custom-annotator-string
标志手动选择为标注创建的'.npy'文件的名称。可以使用单个32GB内存的GPU进行--llm-size 7
和--llm-size 13
的标注。
你可以使用8个GPU的节点进行--llm-size 70
的标注。这里我们提供使用NVIDIA V100s 32G GPU对100k对数据集进行标注的粗略时间估计,这应该能大致重现我们的大多数结果(我们的结果是使用500k对获得的)。
模型 | 所需标注资源 |
---|---|
Llama 2 7b | 约32 GPU小时 |
Llama 2 13b | 约40 GPU小时 |
Llama 2 70b | 约72 GPU小时 |
奖励训练
在第二阶段,我们通过交叉熵将大语言模型的偏好提炼为奖励函数。要使用默认超参数启动奖励训练,请使用以下命令。
python -m scripts.train_reward --batch_size 1024 --num_workers 40 \
--reward_lr 1e-5 --num_epochs 10 --seed 777 \
--dataset_dir motif_dataset --annotator llama70b_msg_defaultgoal_default \
--experiment standard_reward --train_dir train_dir/reward_saving_dir
奖励函数将通过位于--dataset_dir
中的annotator
的标注进行训练。然后,结果函数将保存在train_dir
下的--experiment
子文件夹中。
强化学习训练
最后,我们通过强化学习使用得到的奖励函数训练一个智能体。要在NetHackScore-v1
任务上使用默认超参数训练智能体,结合内在和外在奖励进行实验,可以使用以下命令。
python -m scripts.main --algo APPO --env nle_fixed_eat_action --num_workers 24 \
--num_envs_per_worker 20 --batch_size 4096 --reward_scale 0.1 --obs_scale 255.0 \
--train_for_env_steps 2_000_000_000 --save_every_steps 10_000_000 \
--keep_checkpoints 5 --stats_avg 1000 --seed 777 --reward_dir train_dir/reward_saving_dir/standard_reward/ \
--experiment standard_motif --train_dir train_dir/rl_saving_dir \
--extrinsic_reward 0.1 --llm_reward 0.1 --reward_encoder nle_torchbeast_encoder \
--root_env NetHackScore-v1 --beta_count_exponent 3 --eps_threshold_quantile 0.5
要更改任务,只需修改--root_env
参数。下表明确列出了与论文中呈现的实验相匹配所需的值。NetHackScore-v1
任务的extrinsic_reward
值为0.1
,而其他所有任务的值为10.0
,以激励智能体达到目标。
环境 | root_env |
---|---|
得分 | NetHackScore-v1 |
楼梯 | NetHackStaircase-v1 |
楼梯(第3层) | NetHackStaircaseLvl3-v1 |
楼梯(第4层) | NetHackStaircaseLvl4-v1 |
神谕 | NetHackOracle-v1 |
神谕-清醒 | NetHackOracleSober-v1 |
此外,如果你只想使用来自大语言模型的内在奖励而不使用环境奖励来训练智能体,只需设置--extrinsic_reward 0.0
。在仅使用内在奖励的实验中,我们只在智能体死亡时终止回合,而不是在智能体达到目标时终止。这些修改后的环境列在下表中。
环境 | root_env |
---|---|
楼梯(第3层)- 仅内在奖励 | NetHackStaircaseLvl3Continual-v1 |
楼梯(第4层)- 仅内在奖励 | NetHackStaircaseLvl4Continual-v1 |
可视化你的强化学习智能体
我们还提供了一个脚本来可视化你训练的强化学习智能体。这可以提供对其行为的重要洞察,同时还会生成每个回合的顶级消息,有助于理解它试图优化的目标。你只需运行以下命令即可。
python -m scripts.visualize --train_dir train_dir/rl_saving_dir --experiment standard_motif
引用
如果你在我们的工作基础上进行研究或发现它有用,请使用以下bibtex引用。
@article{klissarovdoro2023motif,
title={Motif: Intrinsic Motivation From Artificial Intelligence Feedback},
author={Klissarov, Martin and D'Oro, Pierluca and Sodhani, Shagun and Raileanu, Roberta and Bacon, Pierre-Luc and Vincent, Pascal and Zhang, Amy and Henaff, Mikael},
year={2023},
month={9},
journal={arXiv preprint arXiv:2310.00166}
}
许可证
Motif的大部分内容采用CC-BY-NC许可,但项目的某些部分采用单独的许可条款:sample-factory采用MIT许可。