使用语言学习建模世界
该存储库包含了论文使用语言学习建模世界的代码。我们介绍了Dynalang,一种利用多种语言类型来解决任务的代理,通过多模态世界模型使用语言来预测未来。
欲了解更多信息,请访问:
快速开始
安装依赖项:
pip install -e .
🏠 HomeGrid
安装HomeGrid环境并运行示例训练脚本,任务名称是 (homegrid_task
, homegrid_future
, homegrid_dynamics
, 或 homegrid_corrections
):
pip install homegrid
sh scripts/run_homegrid.sh homegrid_task EXP_NAME GPU_IDS SEED
👾 Messenger
- 为Messenger环境安装依赖项:
sudo apt-get install \
libsdl-image1.2-dev libsdl-mixer1.2-dev libsdl-ttf2.0-dev \
libsdl1.2-dev libsmpeg-dev subversion libportmidi-dev ffmpeg \
libswscale-dev libavformat-dev libavcodec-dev libfreetype6-dev
- 克隆Messenger环境并安装:
git clone https://github.com/ahjwang/messenger-emma
pip install -e messenger-emma
-
从Google Drive下载预嵌入的句子,并将
pkl
文件放到dynalang/embodied/envs/data
目录中。 -
运行示例训练脚本,选择s1, s2, 或 s3:
sh scripts/run_messenger_s1.sh EXP_NAME GPU_IDS SEED
🛋 VLN
- 安装VLN的依赖项。由于VLN-CE依赖于旧版本的Habitat模拟器,我们提供了一个带有固定依赖项版本的conda yml文件,并建议安装相同版本的JAX。
conda create -n dynalang-vln python=3.8
pip install "jax[cuda11_cudnn82]==0.4.8" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install torch torchvision
conda env update -f env_vln.yml
conda install -c aihabitat -c conda-forge habitat-sim=0.1.7 headless
- 将我们的VLN-CE和habitat-lab分叉克隆到此项目存储库:
git clone https://github.com/jlin816/VLN-CE VLN_CE
git clone https://github.com/jlin816/habitat-lab habitat_lab
- 将Matterport3D数据下载到VLN-CE目录中(需要Python 2.7):附加说明
conda create -n py27 python=2.7
conda activate py27
python scripts/download_mp.py --task habitat -o VLN_CE/data/scene_datasets/mp3d/
cd VLN_CE/data/scene_datasets
unzip mp3d/v1/tasks/mp3d_habitat.zip
conda deactivate
场景应现在已填充到目录下 VLN_CE/data/scene_datasets/mp3d/{scene_dir}/...
。
- 将Room-to-Room数据集下载到VLN_CE目录:更多说明
mkdir -p VLN_CE/data/datasets
cd VLN_CE/data/datasets
gdown https://drive.google.com/uc?id=1kQ_at68wiK2vAmlWbjJ4EDrLtaM0nfkR
unzip R2R_VLNCE_v1-3_preprocessed.zip
-
从Google Drive下载预嵌入的句子,并将
pkl
文件放到dynalang/embodied/envs/data
目录中。 -
运行示例训练脚本:
sh scripts/run_vln.sh EXP_NAME GPU_IDS SEED
💬 LangRoom
LangRoom必须从langroom
分支运行,该分支包含了支持LangRoom字典动作空间进行同步移动和对话的更改。安装LangRoom环境并运行示例训练脚本:
git checkout langroom
pip install langroom
sh run_langroom.sh EXP_NAME GPU_IDS SEED
📑 文本预训练和微调
安装附加的依赖项并运行示例训练脚本,在TinyStories上进行训练。默认情况下,脚本使用Messenger S2任务的观察和行动空间初始化一个模型,同时监控训练时回放缓冲器上的损失。如果您希望执行相同操作,则需要传递指向先前Messenger S2代理的 .../episodes/
目录(包含 .npz
块)的路径。
pip install datasets
sh scripts/pretrain_text.sh EXP_NAME GPU_IDS SEED roneneldan/TinyStories /PATH/TO/EVAL/REPLAY/EPISODES
要将预训练的世界模型加载到代理中进行微调,运行训练脚本并使用以下标志(load_wm_ckpt
仅加载世界模型的权重,并从头开始初始化代理的其余部分):
python dynalang/train.py \
--load_wm_ckpt True \
--run.from_checkpoint /path/to/pretraining/checkpoint.pkl \
...
训练配置提示
- 请注意,要使用我们在论文中使用的相同批量大小,您可能需要使用多个GPU。为此,运行脚本时提供GPU ID列表,并额外指定用于训练世界模型和运行策略的GPU(例如
sh run_messenger_s3.sh EXP_NAME 0,1,2,3 SEED --jax.train_devices 0,1,2,3 --jax.policy_devices 0
)。 - 您可能需要使用
parallel
脚本来加快训练速度。有关示例,请参阅scripts/run_vln.sh
。
致谢
此存储库改编自 DreamerV3。
引用
@article{lin2023learning,
title={Learning to Model the World with Language},
author={Jessy Lin and Yuqing Du and Olivia Watkins and Danijar Hafner and Pieter Abbeel and Dan Klein and Anca Dragan},
year={2023},
eprint={2308.01399},
archivePrefix={arXiv},
}