学习笔记——基于Ubuntu22.04的RoboticsDiffusionTransformer ManiSkill仿真

原创于 2025-12-06 17:35:11 发布 · 730 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#linux

RoboticsDiffusionTransformer（简称RDT）开源地址：https://github.com/csq-ing/RoboticsDiffusionTransformer?tab=readme-ov-file

本文展示了在ManiSkill中运行RDT模型跑五个基准任务：PegInsertionSide、PickCube、StackCube、PlugCharger 和 PushCube的具体步骤。

一.环境配置

1.基础环境配置

这一步基本复制原文

但是直接用pip install flash-attn --no-build-isolation下载的是最新版本的flash-attn，而RDT中有用到以前版本的一个函数，如果下载最新版的会报错因此建议用

pip install flash-attn==2.3.0 --no-build-isolation

指定版本。

总体运行步骤是这样的：

git clone git@github.com:thu-ml/RoboticsDiffusionTransformer.git
cd RoboticsDiffusionTransformer

conda create -n rdt python=3.10.0
conda activate rdt

pip install torch==2.1.0 torchvision==0.16.0  --index-url https://download.pytorch.org/whl/cu121

pip install packaging==24.0

pip install flash-attn==2.3.0 --no-build-isolation

pip install -r requirements.txt

建议一步一步运行。

2.仿真环境配置

对应原文

conda activate rdt
pip install --upgrade mani_skill
sudo apt-get install libvulkan1
sudo apt-get install vulkan-tools
vulkaninfo

最后一句vulkaninfo是用来测试libvulkan1是否下载成功的。

然后从这里下载Hugging Face 仓库微调后的模型权重。下载后直接放在克隆到本地的RoboticsDiffusionTransformer文件夹中就可以了。

二.预计算配置

如果GPU显存不够（大概需要40多G）建议使用源代码中推荐的预计算，即提前将语言指令转换成模型能理解的数学向量（嵌入），并把这个结果保存到硬盘上。这样，在实际运行机器人时，就无需再进行这个转换过程了。

1.下载预计算语言嵌入文件

点击这里，就可以下载五个基准任务：PegInsertionSide、PickCube、StackCube、PlugCharger 和 PushCube的预计算语言嵌入文件了。下载完后也是放在克隆到本地的RoboticsDiffusionTransformer文件夹中。

2.代码修改

（1）.将eval_sime中val_rdt_maniskill.py

policy = create_model(
    args=config, 
    dtype=torch.bfloat16,
    pretrained=pretrained_path,
    pretrained_text_encoder_name_or_path=pretrained_text_encoder_name_or_path,
    pretrained_vision_encoder_name_or_path=pretrained_vision_encoder_name_or_path
)

改为

policy = create_model(
    args=config, 
    dtype=torch.bfloat16,
    pretrained=pretrained_path,
    pretrained_text_encoder_name_or_path=None,  # 关键：设为 None 表示不加载 T5
    pretrained_vision_encoder_name_or_path=pretrained_vision_encoder_name_or_path
)

（2）.将scripts中maniskill_model.py

    def get_text_encoder(self, pretrained_text_encoder_name_or_path):
        text_embedder = T5Embedder(from_pretrained=pretrained_text_encoder_name_or_path, 
                                   model_max_length=self.args["dataset"]["tokenizer_max_length"], 
                                   device=self.device)
        tokenizer, text_encoder = text_embedder.tokenizer, text_embedder.model
        return tokenizer, text_encoder

改为

   def get_text_encoder(self, pretrained_text_encoder_name_or_path):
        # 新增：如果传入 None，则不加载文本编码器
        if pretrained_text_encoder_name_or_path is None:
            return None, None  
        # 返回空的 tokenizer 和 model
        
        text_embedder = T5Embedder(from_pretrained=pretrained_text_encoder_name_or_path, 
                                   model_max_length=self.args["dataset"]["tokenizer_max_length"], 
                                   device=self.device)
        tokenizer, text_encoder = text_embedder.tokenizer, text_embedder.model
        return tokenizer, text_encoder

将

    def reset(self):
        """Set model to evaluation mode.
        """
        device = self.device
        weight_dtype = self.dtype
        self.policy.eval()
        self.text_model.eval()
        self.policy = self.policy.to(device, dtype=weight_dtype)
        if self.text_model is not None:
            self.text_model = self.text_model.to(device, dtype=weight_dtype)

        self.vision_model = self.vision_model.to(device, dtype=weight_dtype)

改为

    def reset(self):
        """Set model to evaluation mode.
        """
        device = self.device
        weight_dtype = self.dtype
        self.policy.eval()
        #self.text_model.eval()
        if self.text_model is not None:
            self.text_model.eval()
        self.vision_model.eval()

        self.policy = self.policy.to(device, dtype=weight_dtype)
        if self.text_model is not None:
            self.text_model = self.text_model.to(device, dtype=weight_dtype)

        self.vision_model = self.vision_model.to(device, dtype=weight_dtype)

三.保存视频

如果需要保存视频，可以按一下步骤将仿真过程保存为mp4视频文件。

以下步骤都在eval_sime中val_rdt_maniskill.py 中处理。

1.导入视频处理库

在代码顶部添加 imageio 库（用于合成视频）和 os 库（用于创建目录）：

import imageio
import os

2.修改参数解析，添加视频保存相关参数

在 parse_args 函数中添加 --save-video（控制是否保存）和 --record-dir（保存目录）参数：

def parse_args(args=None):
    parser = argparse.ArgumentParser()
    # ... 已有的其他参数 ...
    # 添加视频保存参数
    parser.add_argument("--save-video", action="store_true", help="Whether to save simulation videos")
    parser.add_argument("--record-dir", type=str, default="simulation_videos/rdt", help="Directory to save videos")
    return parser.parse_args()

3.创建视频保存目录

在循环开始前（for episode in tqdm.trange(total_episodes): 之前），创建视频保存目录：

# 处理视频保存目录
if args.save_video:
    os.makedirs(args.record_dir, exist_ok=True)
    print(f"Videos will be saved to: {args.record_dir}")

4.视频合成

在

print(f"Trial {episode+1} finished, success: {info['success']}, steps: {global_steps}")

之前，添加视频合成代码：

    # 保存视频（如果启用了 --save-video）
    if args.save_video and len(video_frames) > 0:
        # 生成视频文件名（包含任务名、episode编号、成功状态）
        video_filename = f"{args.record_dir}/{env_id}_episode_{episode+1}_success_{info['success']}.mp4"
        # 保存视频（帧率设为30）
        imageio.mimsave(video_filename, video_frames, fps=30)
        print(f"Video saved to: {video_filename}")

四.运行评估脚本

如果需要保存仿真视频，运行

python -m eval_sim.eval_rdt_maniskill \
  --env-id StackCube-v1 \
  --pretrained_path ./mp_rank_00_model_states.pt \
  --save-video \
  --record-dir "simulation_videos/rdt"

如果不需要保存，删除后面两行，运行

python -m eval_sim.eval_rdt_maniskill \
  --env-id StackCube-v1 \
  --pretrained_path ./mp_rank_00_model_states.pt \

其中：

参数： --env-id StackCube-v1指定要评估的仿真环境（任务），如果要运行其他任务，只需要将它修改它就可以了，比如： --env-id PickCube-v1。

参数： --record-dir "simulation_videos/rdt"指定视频的保存目录，如果要将视频保存到其他地方只需要修改它就可以了。

五.其他报错

1.huggingface版本过高

现象：

ImportError: cannot import name '__version__' from partially initialized module 'huggingface_hub' (most likely due to a circular import) (/root/miniconda3/envs/rdt/lib/python3.10/site-packages/huggingface_hub/__init__.py)

解决方法：

# 激活 rdt 环境（确保在正确环境中操作）
conda activate rdt

# 卸载两个库
pip uninstall -y huggingface_hub diffusers

# 安装指定版本（确保兼容性）
pip install "huggingface_hub==0.18.0" "diffusers==0.20.0"

# 验证库是否正确安装
pip list | grep -E "huggingface_hub|diffusers"

2.内存不足

创建交换文件：

# 1. 创建 50GB 空文件（/swapfile 为交换文件路径，可自定义）
sudo fallocate -l 50G /swapfile
# 若 fallocate 报错，用 dd 命令替代（速度稍慢，但兼容性更好）：
# sudo dd if=/dev/zero of=/swapfile bs=1G count=50

# 2. 设置文件权限（仅 root 可读写，避免安全风险）
sudo chmod 600 /swapfile

# 启用指定的交换文件（路径需与创建的文件一致）
sudo swapon /swapfile

# 验证是否启用成功（若输出包含/swapfile，说明启用成功）
swapon --show
# 或通过free命令查看（Swap行的total会显示交换文件大小）
free -h

如果需要设置开机自动启用

设置开机自动启用

sudo nano /etc/fstab

文件末尾添加一行

/swapfile  swap  swap  defaults  0  0

Ctrl+O确认保存，回车，Ctrl+X退出编辑器

执行以下命令验证配置是否正确（无报错则说明配置有效）：
sudo mount -a