AVDC

AVDC

AVDC 训练视频策略的官方代码库

新闻：我们已经发布了另一个用于运行 Meta-World 和 iTHOR 实验的代码库在此！

https://github.com/flow-diffusion/flow-diffusion.github.io/assets/43379407/9aa380df-0ff7-4c41-af2d-d67d23c53e72

本代码库包含了我们工作中提出的视频策略训练代码 通过密集对应学习从无动作视频中行动 Ko Po-Chen、 Mao Jiayuan、 Du Yilun、 Sun Shao-Hua、 Joshua B. Tenenbaum 项目网站 | 论文 | arXiv | 实验代码库

@article{Ko2023Learning,
  title={{Learning to Act from Actionless Videos through Dense Correspondences}},
  author={Ko, Po-Chen and Mao, Jiayuan and Du, Yilun and Sun, Shao-Hua and Tenenbaum, Joshua B},
  journal={arXiv:2310.08576},
  year={2023},
}

更新

• 2023/10/21：支持自定义任务名称和任意数量的视频（移除了实验中遗留的任务/视频数量限制）
• 2024/01/02：发布了另一个用于 Meta-World 和 iTHOR 实验的代码库在此。
• 2024/01/03：更新了 DDIM 采样和无分类器引导的参数。

入门

我们建议使用 conda 创建一个安装了 pytorch 的新环境。

conda create -n avdc python=3.9
conda activate avdc
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

接下来，克隆代码库并安装依赖

git clone https://github.com/flow-diffusion/AVDC
cd AVDC
pip install -r requirements.txt

数据集结构

本代码库在 datasets/ 中包含示例数据集结构。

pytorch 数据集类在 flowdiffusion/datasets.py 中定义

训练模型

对于 Meta-World 实验，运行

cd flowdiffusion
python train_mw.py --mode train
# 或 python train_mw.py -m train

或使用 accelerate 运行

accelerate launch train_mw.py

对于 iTHOR 实验，运行 train_thor.py 而不是 train_mw.py 对于桥梁实验，运行 train_bridge.py 而不是 train_mw.py

训练好的模型应该保存在 ../results 文件夹中

要恢复训练，可以使用 -c --checkpoint_num 参数。

# 这将恢复使用第一个检查点（应命名为 model-1.pt）的训练
python train_mw.py --mode train -c 1

推理

使用以下参数进行推理 -p --inference_path：指定输入图像路径 -t --text：指定任务描述文本 -n sample_steps 可选，测试时采样使用的步数。如果指定值小于 100，将使用 DDIM 采样。 -g guidance_weight 可选，用于无分类器引导的权重。设为正值以开启无分类器引导。

例如：

python train_mw.py --mode inference -c 1 -p ../examples/assembly.png -t assembly -g 2 -n 20

预训练模型

我们还提供了实验中描述的模型检查点，如下。 Meta-World | iTHOR | Bridge

下载并将 .pt 文件放在 results/[environment] 文件夹中。最终的目录结构应为 results/{mw, thor, bridge}/model-[x].pt，例如 results/mw/model-24.pt

或使用 download.sh

./download.sh metaworld
# ./download.sh ithor
# ./download.sh bridge

之后，您可以使用参数 -c [x] 恢复训练或使用我们的检查点进行推理。例如：

python train_mw.py --mode train -c 24

或

python train_mw.py --mode inference -c 24 -p ../examples/assembly.png -t assembly

致谢

本代码库修改自以下代码库： imagen-pytorch guided-diffusion