登录/ 注册
 找回密码
 立即注册
关闭

CSDN热搜

程序园 精品问答 技术交流 资源下载
返回列表 发新帖
首页 › 业界区 › 业界 ›再不学就晚了！RDT × LeRobot与RDKS100部署详解 ...

再不学就晚了！RDT × LeRobot与RDKS100部署详解

雌鲳签 2025-9-29 12:50:14
作者：SkyXZ
CSDN：SkyXZ～-CSDN博客
博客园：SkyXZ - 博客园
机械臂：LeRobot-SO101           数采机：MacBook-Pro Python3.10
开发机：Ubuntu 22.04, Cuda12.4,8 × NVIDIA A100-SXM4-40GB
开发板：RDK OS 4.0.2 Based on Ubuntu 22.04, Python 3.10.12, OpenExplore 3.2.0
相关资料：

  • LeRobot Doc：https://huggingface.co/docs/lerobot/main/en/index
  • RDT 170M&1B：https://github.com/thu-ml/RoboticsDiffusionTransformer
  • RDT on RDKS100：RDT on Double RDK S100P 全流程文档
一、环境安装&机械臂配置


  • 所有代码已上传至GitHub：GitHub - xiongqi123123/LeRobot-VLA: Classic VLA for LeRobot
环境安装

       我们首先完成LeRobot环境的安装，我们默认使用conda作为环境管理，先运行以下命令创建一个Python3.10的虚拟环境
  1. conda create -y -n lerobot python=3.10
复制代码
        接着便可以在环境中运行以下命令来配置lerobot所需要的依赖(使用的lerobot源码为我修改之后的，仅添加了本地serverclient，其他部分与官方源码一致)
  1. # step:0 安装编译依赖
  2. sudo apt-get install cmake build-essential python-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev pkg-config
  3. # step:1 激活环境
  4. conda activate lerobot
  5. # step:2 安装ffmpeg
  6. conda install ffmpeg -c conda-forge
  7. # 以下两种方式任选其一：
  8. # step:3 从源码安装lerobot
  9. git clone https://github.com/xiongqi123123/LeRobot-VLA.git
  10. cd LeRobot-VLA/lerobot
  11. pip install -e .
  12. # step:3 从PyPI 安装
  13. pip install lerobot
  14. # 要安装附加功能，请使用以下之一
  15. pip install 'lerobot[all]'          # All available features
  16. pip install 'lerobot[aloha,pusht]'  # Specific features (Aloha & Pusht)
  17. pip install 'lerobot[feetech]'      # Feetech motor support
复制代码
机械臂安装


  • 官方安装教程：SO-101 - Hugging Face 机器学习平台
        由于SO-ARM101机械臂默认不提供机械臂上的相机安装位置，因此我们在原始的机械臂夹爪部分自行设计添加了一个相机固定的位置(安装孔位与夹爪是对齐的)，我们使用的相机是亚博智能的1080P高清免驱摄像头，打印文件已保存进仓库中：
1.png

        拿到刚拆封的机械臂配件我们首先将Follower和Leader臂的物料进行区分，要注意的是Follower机械臂使用的是12个ST-3215-C001（7.4V）1:345齿轮比的电机，而Leader臂不同关机使用的电机型号有所不同，不同关节的电机型号区分如下图及上表：
2.png

        接下来我们便可以按照官方提供的如下3D演示动画安装Follower和Leader了，下面是关节一到关节五以及夹爪的安装实例，除了夹爪之外，前面五个关节的安装方法均一致，仅需注意Leader臂的电机型号：
机械臂电机配置

        在完成机械臂的安装后我们便可以开始对机械臂两个臂的电机进行配置设置其对应的ID了，新版的LeRobot提供了CLI命令可以直接运行对应的任务，我们首先将两个机械臂的串口钱全部接上电脑并运行如下命令，接着按照提示拔出其中一个串口线按下回车即可知道拔出的串口号是多少（实际就是记录插上的所有串口，然后再和拔出后的进行对比就可以知道哪个串口少了...），示例输出如下图：
  1. lerobot-find-port
复制代码
3.png

        由于LeRobot使用的是总线舵机，每个电机都通过总线上的唯一 ID 进行识别，而全新电机通常带有一个默认 ID 为 1，所以为了让电机和控制器之间正常通信，我们需要为每个电机设置一个唯一的、不同的 ID，在确定了自己的两个机械臂分别对应的串口后我们便可以一个臂一个臂一个电机一个电机的进行配置啦，在新版的LeRobot库不需要重复多次运行电机配置命令，我们只需要运行以下命令并依次将不同的舵机线插到控制板按回车即可,具体可以参考以下视频：
  1. lerobot-setup-motors \
  2.     --robot.type=so101_follower \
  3.     --robot.port=/dev/tty.usbmodem585A0076841  # <- paste here the port found at previous step
复制代码
        接下来仅需运行这个脚本或者是如下命令即可完成数据的转换：
  1. lerobot-calibrate \
  2.     --robot.type=so101_follower \
  3.     --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
  4.     --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
复制代码
        我们使用RoboTwin修改后的RDT版本，这个版本使用比较简单快速（更多关于RDT的信息请见：RDT on Double RDK S100P 全流程文档），由于RDT默认是双臂任务，而我们采集的LeRobot是单臂数据，且我们只采集了两个摄像头的画面跟RDT默认的三个摄像头不匹配，因此如果直接训练的话肯定会报索引不匹配的错误，因此我们还需要对数据集加载的部分进行修改，首先是修改action的归一化部分，我们直接在加载的时候对数据除以[[180, 180, 180, 180, 180, 180]]，然后便是将LeRobot的单臂映射到RDT的右臂部分的动作维度并将左臂整个给剔除同时把RDT默认加载的右臂图像用Ground图像进行替代，请使用以下完成了修改的代码替换原本的代码中的RDT/data/hdf5_vla_dataset.py
  1. lerobot-find-cameras opencv
复制代码
开始训练

训练环境配置

        我们首先来安装RDT训练所需的环境，此我们进入RDT目录下依次安装如下包即可：
  1. lerobot-teleoperate \
  2.     --robot.type=so101_follower \
  3.     --robot.port=/dev/tty.usbmodem5AB90671801 \
  4.     --robot.id=my_awesome_follower_arm \
  5.     --teleop.type=so101_leader \
  6.     --teleop.port=/dev/tty.usbmodem5AB90671501 \
  7.     --teleop.id=my_awesome_leader_arm \
复制代码
        除了上述依赖之外我们还需要安装flash_attn用来加速，为了避免网络连接问题我们手动下载预编译的wheel 文件，下载连接为：https://github.com/Dao-AILab/flash-attention/releases，我们需要根据我们实际安装的torch及cuda版本来选择对应的版本，然后我们还需要根据我们下载的PyTorch是如何编译的来选择对应的cxx11abi 是 TRUE 还是 FALSE。
  1. lerobot-teleoperate \
  2.     --robot.type=so101_follower \
  3.     --robot.port=/dev/tty.usbmodem5AB90671801 \
  4.     --robot.id=my_awesome_follower_arm \
  5.     --robot.cameras="{ arm: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}, front: {type: opencv, index_or_path: 2, width: 1920, height: 1080, fps: 30}}" \
  6.     --teleop.type=so101_leader \
  7.     --teleop.port=/dev/tty.usbmodem5AB90671501 \
  8.     --teleop.id=my_awesome_leader_arm \
  9.     --display_data=true
复制代码
        接着我们便可以根据输出判断我们要下载的是哪个版本啦，如下图所示我们当前的PyTorch是CXX11_ABI = 0，因此我们要下载的是cxx11abiFALSE 的 .whl 文件
  1. lerobot-record \
  2.     --robot.type=so101_follower \
  3.     --robot.port=/dev/tty.usbmodem5AB90671801 \
  4.     --robot.id=my_awesome_follower_arm \
  5.     --robot.cameras="{ arm: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}, front: {type: opencv, index_or_path: 2, width: 1920, height: 1080, fps: 30}}" \
  6.     --teleop.type=so101_leader \
  7.     --teleop.port=/dev/tty.usbmodem5AB90671501 \
  8.     --teleop.id=my_awesome_leader_arm \
  9.     --display_data=true \
  10.     --dataset.repo_id=skyxz/blackmarker_scence1 \
  11.     --dataset.num_episodes=5 \
  12.     --dataset.single_task="Grab the black marker and put it in the bin" \
  13.     --dataset.push_to_hub=False \
  14.     --dataset.episode_time_s=15 \
  15.     --dataset.reset_time_s=5
复制代码
        如果上述步骤跟我的一样的话，那么大家需要下载并安装的应该是如下的版本，下载后安装即可：
  1. (RoboTwin) qi.xiong@A100-Test:~/Data_Qi/LeRobot/skyxz/blackmarker_scence1$ python3 -c "import h5py; f=h5py.File('/home/qi.xiong/Data_Qi/RDT/processed_data/place_dual_shoes-demo_clean-300/episode_3/episode_3.hdf5','r'); f.visit(print); f.close()"
  2. action
  3. observations
  4. observations/images
  5. observations/images/cam_high
  6. observations/images/cam_left_wrist
  7. observations/images/cam_right_wrist
  8. observations/left_arm_dim
  9. observations/qpos
  10. observations/right_arm_dim
复制代码
预训练模型下载

        RDT分为了1B版本（SigLip+DIT+Adaptor）以及单独的170M（DiT）两个版本，其中的区别仅在于最后的DiT的hidden_size和depth维度区别，170M相比于1B的版本直接减半了，如果需要在RDKS100上部署的话请参考接下来的RDT170M模型版本以保证可行的性能，在训练之前还需按照如下的步骤下载预训练模型及完成训练环境的安装；
  1. #!/usr/bin/env python3
  2. """
  3. LeRobot到RDT数据转换脚本
  5. LeRobot机器人结构：
  6. - 5个关节 (shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll)
  7. - 1个夹爪 (gripper)
  8. - 总计：6个自由度 (6DOF)
  10. 维度映射（匹配RDT训练代码）：
  11. - left_arm_dim = 0 (单臂机器人，左臂不存在)
  12. - right_arm_dim = 6 (5关节 + 1夹爪，映射到RDT的right_arm部分)
  13. - 状态向量：6维 [joint1, joint2, joint3, joint4, joint5, gripper]
  14. - RDT索引映射：right_arm_joint_0_pos到right_arm_joint_5_pos (索引0-5)
  15. """
  17. import sys
  18. import os
  19. import h5py
  20. import numpy as np
  21. import cv2
  22. import argparse
  23. import yaml
  24. import json
  25. import subprocess
  26. from pathlib import Path
  27. import pandas as pd
  28. import torch
  30. current_dir = os.path.dirname(__file__)
  31. sys.path.append(os.path.join(current_dir, ".."))
  32. from models.multimodal_encoder.t5_encoder import T5Embedder
  34. def extract_frames_from_video(video_path, output_dir, episode_idx):
  35.     if not os.path.exists(video_path):
  36.         print(f"  No video file: {video_path}")
  37.         return []
  38.    
  39.     temp_dir = os.path.join(output_dir, f"temp_frames_{episode_idx}")
  40.     if not os.path.exists(temp_dir):
  41.         os.makedirs(temp_dir)
  42.    
  43.     output_pattern = os.path.join(temp_dir, "frame_%04d.jpg")
  44.    
  45.     try:
  46.         cmd = [
  47.             'ffmpeg', '-i', video_path,
  48.             '-vf', 'fps=30',
  49.             '-q:v', '2',
  50.             output_pattern,
  51.             '-y'
  52.         ]
  53.         
  54.         result = subprocess.run(cmd, capture_output=True, text=True)
  55.         
  56.         if result.returncode != 0:
  57.             print(f"  Failed to extract frames with ffmpeg: {result.stderr}")
  58.             return []
  59.         
  60.         frames = []
  61.         frame_files = sorted([f for f in os.listdir(temp_dir) if f.endswith('.jpg')])
  62.         
  63.         for frame_file in frame_files:
  64.             frame_path = os.path.join(temp_dir, frame_file)
  65.             frame = cv2.imread(frame_path)
  66.             if frame is not None:
  67.                 frame_resized = cv2.resize(frame, (640, 480))
  68.                 frames.append(frame_resized)
  69.         
  70.         print(f"  Successfully extracted {len(frames)} frames")
  71.         
  72.         for frame_file in frame_files:
  73.             os.remove(os.path.join(temp_dir, frame_file))
  74.         os.rmdir(temp_dir)
  75.         
  76.         return frames
  77.         
  78.     except Exception as e:
  79.         print(f"  Error extracting frames: {e}")
  80.         return []
  82. def load_lerobot_episode(data_dir, episode_idx, output_dir):
  83.     """加载LeRobot的单个episode数据
  84.    
  85.     LeRobot数据结构：
  86.     - action: 6维 [shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper]
  87.     - observation.state: 6维 [shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper]
  88.     - 图像: 高位相机 + 手臂相机
  89.     """
  90.     parquet_path = os.path.join(data_dir, "data/chunk-000", f"episode_{episode_idx:06d}.parquet")
  91.     if not os.path.exists(parquet_path):
  92.         print(f"Episode {episode_idx} parquet file does not exist: {parquet_path}")
  93.         return None
  94.    
  95.     df = pd.read_parquet(parquet_path)
  96.    
  97.     actions = []
  98.     qpos = []
  99.    
  100.     for i in range(len(df)):
  101.         action = df['action'].iloc[i]
  102.         state = df['observation.state'].iloc[i]
  103.         
  104.         if isinstance(action, np.ndarray):
  105.             actions.append(action.astype(np.float32))
  106.         else:
  107.             actions.append(np.array(action, dtype=np.float32))
  108.             
  109.         if isinstance(state, np.ndarray):
  110.             qpos.append(state.astype(np.float32))
  111.         else:
  112.             qpos.append(np.array(state, dtype=np.float32))
  113.    
  114.     high_cam_path = os.path.join(data_dir, "videos/chunk-000/observation.images.high", f"episode_{episode_idx:06d}.mp4")
  115.     arm_cam_path = os.path.join(data_dir, "videos/chunk-000/observation.images.arm", f"episode_{episode_idx:06d}.mp4")
  116.    
  117.     print(f"  Extracting high camera frames...")
  118.     high_images = extract_frames_from_video(high_cam_path, output_dir, episode_idx)
  119.    
  120.     print(f"  Extracting arm camera frames...")
  121.     arm_images = extract_frames_from_video(arm_cam_path, output_dir, episode_idx)
  122.    
  123.     target_frames = len(df)
  124.     if len(high_images) > target_frames:
  125.         high_images = high_images[:target_frames]
  126.     if len(arm_images) > target_frames:
  127.         arm_images = arm_images[:target_frames]
  128.    
  129.     while len(high_images) < target_frames and high_images:
  130.         high_images.append(high_images[-1])
  131.     while len(arm_images) < target_frames and arm_images:
  132.         arm_images.append(arm_images[-1])
  133.    
  134.     return {
  135.         'actions': np.array(actions),
  136.         'qpos': np.array(qpos),
  137.         'high_images': high_images,
  138.         'arm_images': arm_images,
  139.         'episode_length': len(df)
  140.     }
  142. def images_encoding(imgs):
  143.     if not imgs:
  144.         return [], 0
  145.         
  146.     encode_data = []
  147.     padded_data = []
  148.     max_len = 0
  149.    
  150.     for i in range(len(imgs)):
  151.         success, encoded_image = cv2.imencode(".jpg", imgs[i])
  152.         if success:
  153.             jpeg_data = encoded_image.tobytes()
  154.             encode_data.append(jpeg_data)
  155.             max_len = max(max_len, len(jpeg_data))
  156.         else:
  157.             print(f"  Image encoding failed: {i}")
  158.             empty_data = b""
  159.             encode_data.append(empty_data)
  160.    
  161.     for i in range(len(imgs)):
  162.         padded_data.append(encode_data[i].ljust(max_len, b"\0"))
  163.    
  164.     return encode_data, max_len
  166. def load_task_instructions(data_dir):
  167.     tasks_file = os.path.join(data_dir, "meta/tasks.jsonl")
  168.     if not os.path.exists(tasks_file):
  169.         print(f"Warning: tasks file not found: {tasks_file}")
  170.         return None
  171.    
  172.     instructions = []
  173.     with open(tasks_file, 'r') as f:
  174.         for line in f:
  175.             if line.strip():
  176.                 task_data = json.loads(line.strip())
  177.                 instructions.append(task_data["task"])
  178.    
  179.     print(f"  加载了 {len(instructions)} 个任务指令")
  180.     return instructions
  182. def encode_language_instruction(instruction_text, t5_embedder, device):
  183.     try:
  184.         text_embeds, attn_mask = t5_embedder.get_text_embeddings([instruction_text])
  185.         
  186.         valid_embeds = text_embeds[0][attn_mask[0]].float()
  187.         return valid_embeds.cpu().numpy()
  188.         
  189.     except Exception as e:
  190.         print(f"  Language encoding failed: {e}")
  191.         return np.zeros((1, 4096))
  193. def convert_lerobot_to_rdt(data_dir, output_dir, episode_num, gpu=0, no_language=False):
  194.     if not os.path.exists(output_dir):
  195.         os.makedirs(output_dir)
  196.    
  197.     print(f"Start converting LeRobot data to RDT format...")
  198.     print(f"Data source: {data_dir}")
  199.     print(f"Output directory: {output_dir}")
  200.     print(f"Processing episode number: {episode_num}")
  201.     print(f"GPU device: {gpu}")
  202.    
  203.     scene_name = os.path.basename(data_dir)
  204.    
  205.     instructions = None
  206.     if not no_language:
  207.         instructions = load_task_instructions(data_dir)
  208.    
  209.     t5_embedder = None
  210.     if not no_language and instructions:
  211.         try:
  212.             print(f"  Initializing T5 encoder...")
  213.             t5_model_path = "/home/qi.xiong/Data_Qi/t5-v1_1-xxl"
  214.             if not os.path.exists(t5_model_path):
  215.                 print(f"  Warning: T5 model path does not exist: {t5_model_path}")
  216.                 print(f"  Will skip language processing")
  217.                 no_language = True
  218.             else:
  219.                 t5_embedder = T5Embedder(
  220.                     from_pretrained=t5_model_path,
  221.                     device=f"cuda:{gpu}" if torch.cuda.is_available() else "cpu",
  222.                     model_max_length=120,
  223.                     use_offload_folder=None,
  224.                 )
  225.                 print(f"  T5 encoder initialized successfully")
  226.         except Exception as e:
  227.             print(f"  T5 encoder initialization failed: {e}")
  228.             print(f"  Will skip language processing")
  229.             no_language = True
  230.    
  231.     for i in range(episode_num):
  232.         print(f"Processing episode {i}...")
  233.         
  234.         episode_data = load_lerobot_episode(data_dir, i, output_dir)
  235.         if episode_data is None:
  236.             print(f"Skipping episode {i}")
  237.             continue
  238.         
  239.         episode_output_dir = os.path.join(output_dir, f"episode_{i}")
  240.         if not os.path.exists(episode_output_dir):
  241.             os.makedirs(episode_output_dir)
  242.         
  243.         hdf5_path = os.path.join(episode_output_dir, f"episode_{i}.hdf5")
  244.         
  245.         with h5py.File(hdf5_path, "w") as f:
  246.             f.create_dataset("action", data=episode_data['actions'])
  247.             
  248.             obs = f.create_group("observations")
  249.             obs.create_dataset("qpos", data=episode_data['qpos'])
  250.             
  251.             image = obs.create_group("images")
  252.             
  253.             if episode_data['high_images']:
  254.                 print(f"  Encoding high camera images...")
  255.                 high_enc, len_high = images_encoding(episode_data['high_images'])
  256.                 if high_enc and len_high > 0:
  257.                     image.create_dataset("cam_high", data=high_enc, dtype=f"S{len_high}")
  258.                     print(f"  Saved high camera images: {len(episode_data['high_images'])} frames")
  259.                 else:
  260.                     print(f"  Warning: High camera images encoding failed")
  261.             
  262.             if episode_data['arm_images']:
  263.                 print(f"  Encoding arm camera images...")
  264.                 arm_enc, len_arm = images_encoding(episode_data['arm_images'])
  265.                 if arm_enc and len_arm > 0:
  266.                     image.create_dataset("cam_right_wrist", data=arm_enc, dtype=f"S{len_arm}")
  267.                     print(f"  Saved arm camera images: {len(episode_data['arm_images'])} frames")
  268.                 else:
  269.                     print(f"  Warning: Arm camera images encoding failed")
  270.             
  271.             # 添加机器人维度信息（LeRobot: 5个关节 + 1个夹爪）
  272.             # 根据process_data.py的逻辑，每个时间步都需要记录维度信息
  273.             # LeRobot是单臂机器人，只有右臂：5个关节 + 1个夹爪 = 6维
  274.             # 左臂：0维（单臂机器人）
  275.             
  276.             # 为每个时间步记录维度信息
  277.             left_arm_dim = [0] * len(episode_data['actions'])  # 左臂0维（单臂机器人）
  278.             right_arm_dim = [6] * len(episode_data['actions'])  # 右臂6维（5关节+1夹爪）
  279.             
  280.             obs.create_dataset("left_arm_dim", data=np.array(left_arm_dim))
  281.             obs.create_dataset("right_arm_dim", data=np.array(right_arm_dim))
  282.         
  283.         print(f"  Episode {i} converted successfully: {hdf5_path}")
  284.         print(f"  Data length: {episode_data['episode_length']}")
  285.         print(f"  Action shape: {episode_data['actions'].shape}")
  286.         print(f"  Qpos shape: {episode_data['qpos'].shape}")
  287.         print(f"  High camera frames: {len(episode_data['high_images'])}")
  288.         print(f"  Arm camera frames: {len(episode_data['arm_images'])}")
  289.         
  290.         if not no_language and t5_embedder and instructions:
  291.             print(f"  Processing language instructions...")
  292.             try:
  293.                 instruction = instructions[0]
  294.                
  295.                 language_features = encode_language_instruction(instruction, t5_embedder, f"cuda:{gpu}")
  296.                
  297.                 instructions_dir = os.path.join(episode_output_dir, "instructions")
  298.                 if not os.path.exists(instructions_dir):
  299.                     os.makedirs(instructions_dir)
  300.                
  301.                 lang_embed_path = os.path.join(instructions_dir, "lang_embed_0.pt")
  302.                 torch.save(torch.from_numpy(language_features), lang_embed_path)
  303.                
  304.                 print(f"  Language instruction encoded successfully: {instruction}")
  305.                 print(f"  Language features saved to: {lang_embed_path}")
  306.                 print(f"  Language features shape: {language_features.shape}, data type: {language_features.dtype}")
  307.                
  308.             except Exception as e:
  309.                 print(f"  Language instruction processing failed: {e}")
  310.    
  311.     print(f"\nConversion completed! Processed {episode_num} episodes")
  312.     print(f"Output directory: {output_dir}")
  314. def main():
  315.     parser = argparse.ArgumentParser(description="Convert LeRobot data to RDT format")
  316.     parser.add_argument("--data_dir", type=str, required=True,
  317.                        help="LeRobot data directory path")
  318.     parser.add_argument("--output_dir", type=str, required=True,
  319.                        help="Output directory path")
  320.     parser.add_argument("--episode_num", type=int, default=10,
  321.                        help="Number of episodes to process")
  322.     parser.add_argument("--gpu", type=int, default=0,
  323.                        help="GPU device ID")
  324.     parser.add_argument("--no_language", action="store_true",
  325.                        help="Skip language processing")
  326.    
  327.     args = parser.parse_args()
  328.    
  329.     if not os.path.exists(args.data_dir):
  330.         print(f"Error: Data directory does not exist: {args.data_dir}")
  331.         return
  332.    
  333.     meta_file = os.path.join(args.data_dir, "meta/info.json")
  334.     if not os.path.exists(meta_file):
  335.         print(f"Error: Meta information file not found: {meta_file}")
  336.         return
  337.    
  338.     try:
  339.         subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
  340.         print("ffmpeg is available, will use ffmpeg to extract video frames")
  341.     except (subprocess.CalledProcessError, FileNotFoundError):
  342.         print("Warning: ffmpeg is not available, image data may not be extracted correctly")
  343.         print("Please install ffmpeg: conda install -c conda-forge ffmpeg=6.1")
  344.         return
  345.    
  346.     with open(meta_file, 'r') as f:
  347.         meta_info = yaml.safe_load(f)
  348.    
  349.     total_episodes = meta_info.get('total_episodes', 10)
  350.     if args.episode_num > total_episodes:
  351.         print(f"Warning: Requested episode number ({args.episode_num}) exceeds available number ({total_episodes})")
  352.         args.episode_num = total_episodes
  353.    
  354.     convert_lerobot_to_rdt(
  355.         args.data_dir,
  356.         args.output_dir,
  357.         args.episode_num,
  358.         args.gpu,
  359.         args.no_language
  360.     )
  362. if __name__ == "__main__":
  363.     main()
复制代码
RDT1B微调

        要训练1B的版本我们需要修改RDT/configs/base.yaml文件中model类下的rdt参数
  1. # 法一 bash process_data_rdt.sh data_dir=${1} output_dir=${2} episode_num=${3} gpu_id=${4}
  2. bash process_data_rdt.sh /home/qi.xiong/Data_Qi/LeRobot/skyxz/redmarker_scence4 /home/qi.xiong/DualArm/RoboTwin/policy/RDT-LeRobot/processed_data/redmarker_scence4 5 0
  3. # 法二 python scripts/process_data_lerobot.py --data_dir --output_dir --episode_num --gpu
  4. python3 scripts/process_data_lerobot.py --data_dir /home/qi.xiong/Data_Qi/LeRobot/skyxz/redmarker_scence4 --output_dir /home/qi.xiong/DualArm/RoboTwin/policy/RDT-LeRobot/processed_data/redmarker_scence4
复制代码
        接着生成训练的参数yml文件，并将pretrained_model_name_or_path指向我们先前下载的1B模型
  1. import os
  2. import fnmatch
  3. import json
  5. import h5py
  6. import yaml
  7. import cv2
  8. import numpy as np
  10. from configs.state_vec import STATE_VEC_IDX_MAPPING
  12. class HDF5VLADataset:
  13.     """
  14.     This class is used to sample episodes from the embododiment dataset
  15.     stored in HDF5.
  16.     """
  18.     def __init__(self, model_config_path) -> None:
  19.         # [Modify] The path to the HDF5 dataset directory
  20.         # Each HDF5 file contains one episode
  21.         with open(model_config_path, "r") as f:
  22.             model_config = yaml.safe_load(f)
  23.         HDF5_DIR = model_config["data_path"]
  24.         self.DATASET_NAME = "agilex"
  26.         self.file_paths = []
  27.         for root, _, files in os.walk(HDF5_DIR):
  28.             for filename in fnmatch.filter(files, "*.hdf5"):
  29.                 file_path = os.path.join(root, filename)
  30.                 self.file_paths.append(file_path)
  32.         # Load the config
  33.         with open("configs/base.yaml", "r") as file:
  34.             config = yaml.safe_load(file)
  35.         self.CHUNK_SIZE = config["common"]["action_chunk_size"]
  36.         self.IMG_HISORY_SIZE = config["common"]["img_history_size"]
  37.         self.STATE_DIM = config["common"]["state_dim"]
  39.         # Get each episode's len (use original length, not standardized length)
  40.         episode_lens = []
  41.         for file_path in self.file_paths:
  42.             try:
  43.                 with h5py.File(file_path, "r") as f:
  44.                     qpos = f["observations"]["qpos"][:]
  45.                     num_steps = qpos.shape[0]
  46.                     episode_lens.append(num_steps)
  47.             except Exception as e:
  48.                 print(f"Warning: Could not read {file_path}: {e}")
  49.                 episode_lens.append(0)
  50.         self.episode_sample_weights = np.array(episode_lens) / np.sum(episode_lens)
  52.     def __len__(self):
  53.         return len(self.file_paths)
  55.     def get_dataset_name(self):
  56.         return self.DATASET_NAME
  58.     def get_item(self, index: int = None, state_only=False):
  59.         """Get a training sample at a random timestep.
  61.         Args:
  62.             index (int, optional): the index of the episode.
  63.                 If not provided, a random episode will be selected.
  64.             state_only (bool, optional): Whether to return only the state.
  65.                 In this way, the sample will contain a complete trajectory rather
  66.                 than a single timestep. Defaults to False.
  68.         Returns:
  69.            sample (dict): a dictionary containing the training sample.
  70.         """
  71.         while True:
  72.             if index is None:
  73.                 file_path = np.random.choice(self.file_paths, p=self.episode_sample_weights)
  74.             else:
  75.                 file_path = self.file_paths[index]
  76.             valid, sample = (self.parse_hdf5_file(file_path)
  77.                              if not state_only else self.parse_hdf5_file_state_only(file_path))
  78.             if valid:
  79.                 return sample
  80.             else:
  81.                 index = np.random.randint(0, len(self.file_paths))
  83.     def parse_hdf5_file(self, file_path):
  84.         """[Modify] Parse a hdf5 file to generate a training sample at
  85.             a random timestep.
  87.         Args:
  88.             file_path (str): the path to the hdf5 file
  90.         Returns:
  91.             valid (bool): whether the episode is valid, which is useful for filtering.
  92.                 If False, this episode will be dropped.
  93.             dict: a dictionary containing the training sample,
  94.                 {
  95.                     "meta": {
  96.                         "dataset_name": str,    # the name of your dataset.
  97.                         "#steps": int,          # the number of steps in the episode,
  98.                                                 # also the total timesteps.
  99.                         "instruction": str      # the language instruction for this episode.
  100.                     },
  101.                     "step_id": int,             # the index of the sampled step,
  102.                                                 # also the timestep t.
  103.                     "state": ndarray,           # state[t], (1, STATE_DIM).
  104.                     "state_std": ndarray,       # std(state[:]), (STATE_DIM,).
  105.                     "state_mean": ndarray,      # mean(state[:]), (STATE_DIM,).
  106.                     "state_norm": ndarray,      # norm(state[:]), (STATE_DIM,).
  107.                     "actions": ndarray,         # action[t:t+CHUNK_SIZE], (CHUNK_SIZE, STATE_DIM).
  108.                     "state_indicator", ndarray, # indicates the validness of each dim, (STATE_DIM,).
  109.                     "cam_high": ndarray,        # external camera image, (IMG_HISORY_SIZE, H, W, 3)
  110.                                                 # or (IMG_HISORY_SIZE, 0, 0, 0) if unavailable.
  111.                     "cam_high_mask": ndarray,   # indicates the validness of each timestep, (IMG_HISORY_SIZE,) boolean array.
  112.                                                 # For the first IMAGE_HISTORY_SIZE-1 timesteps, the mask should be False.
  113.                     "cam_left_wrist": ndarray,  # left wrist camera image, (IMG_HISORY_SIZE, H, W, 3).
  114.                                                 # or (IMG_HISORY_SIZE, 0, 0, 0) if unavailable.
  115.                     "cam_left_wrist_mask": ndarray,
  116.                     "cam_right_wrist": ndarray, # right wrist camera image, (IMG_HISORY_SIZE, H, W, 3).
  117.                                                 # or (IMG_HISORY_SIZE, 0, 0, 0) if unavailable.
  118.                                                 # If only one wrist, make it right wrist, plz.
  119.                     "cam_right_wrist_mask": ndarray
  120.                 } or None if the episode is invalid.
  121.         """
  122.         with h5py.File(file_path, "r") as f:
  123.             qpos = f["observations"]["qpos"][:]
  124.             left_arm_dim = f["observations"]["left_arm_dim"][:]
  125.             right_arm_dim = f["observations"]["right_arm_dim"][:]
  126.             num_steps = qpos.shape[0]
  127.             action_dim = qpos
  128.             # [Optional] We drop too-short episode
  129.             # if num_steps < 128:
  130.             #     return False, None
  132.             # [Optional] We skip the first few still steps
  133.             EPS = 1e-2
  134.             # Get the idx of the first qpos whose delta exceeds the threshold
  135.             qpos_delta = np.abs(qpos - qpos[0:1])
  136.             indices = np.where(np.any(qpos_delta > EPS, axis=1))[0]
  137.             if len(indices) > 0:
  138.                 first_idx = indices[0]
  139.             else:
  140.                 raise ValueError("Found no qpos that exceeds the threshold.")
  142.             # We randomly sample a timestep
  143.             step_id = np.random.randint(first_idx - 1, num_steps)
  145.             # Load the instruction
  146.             dir_path = os.path.dirname(file_path)
  148.             # with open(os.path.join(dir_path, 'instruction.json'), 'r') as f_instr:
  149.             #     instruction_dict = json.load(f_instr)
  150.             # # We have 1/3 prob to use original instruction,
  151.             # # 1/3 to use simplified instruction,
  152.             # # and 1/3 to use expanded instruction.
  153.             # instruction_type = np.random.choice([
  154.             #     'instruction', 'expanded_instruction'])
  155.             # instruction = instruction_dict[instruction_type]
  156.             # if isinstance(instruction, list):
  157.             #    instruction = np.random.choice(instruction)
  159.             # You can also use precomputed language embeddings (recommended)
  160.             # instruction = "path/to/lang_embed.pt"
  161.             instructions_path = os.path.join(dir_path, "instructions")
  162.             instructions_names = []
  164.             for filename in os.listdir(instructions_path):
  165.                 # 检查文件名是否以.pt结尾
  166.                 if filename.endswith(".pt"):
  167.                     instructions_names.append(os.path.join(instructions_path, filename))
  168.             instruction = np.random.choice(instructions_names)
  169.             # print(f"choose {instruction} file as instruction.")
  170.             # Assemble the meta
  171.             meta = {
  172.                 "dataset_name": self.DATASET_NAME,
  173.                 "#steps": num_steps,
  174.                 "step_id": step_id,
  175.                 "instruction": instruction,
  176.             }
  178.             # Rescale gripper to [0, 1]
  179.             # qpos = qpos / np.array([[1 for i in range(left_arm_dim[0] + 1 + right_arm_dim[0] + 1)]])
  180.             # target_qpos = f["action"][step_id:step_id + self.CHUNK_SIZE] / np.array(
  181.             #     [[1 for i in range(left_arm_dim[0] + 1 + right_arm_dim[0] + 1)]])
  183.             qpos = qpos / np.array(
  184.             #    [[1, 1, 1, 1, 1, 1, 4.7908, 1, 1, 1, 1, 1, 1, 4.7888]]
  185.             [[180, 180, 180, 180, 180, 180]]
  186.             )
  187.             target_qpos = f['action'][step_id:step_id + self.CHUNK_SIZE] / np.array(
  188.             #    [[1, 1, 1, 1, 1, 1, 11.8997, 1, 1, 1, 1, 1, 1, 13.9231]]
  189.             [[180, 180, 180, 180, 180, 180]]
  190.             )
  192.             # Parse the state and action
  193.             state = qpos[step_id:step_id + 1]
  194.             state_std = np.std(qpos, axis=0)
  195.             state_mean = np.mean(qpos, axis=0)
  196.             state_norm = np.sqrt(np.mean(qpos**2, axis=0))
  197.             actions = target_qpos
  198.             if actions.shape[0] < self.CHUNK_SIZE:
  199.                 # Pad the actions using the last action
  200.                 actions = np.concatenate(
  201.                     [
  202.                         actions,
  203.                         np.tile(actions[-1:], (self.CHUNK_SIZE - actions.shape[0], 1)),
  204.                     ],
  205.                     axis=0,
  206.                 )
  208.             # Fill the state/action into the unified vector
  210.             def fill_in_state(values):
  211.                 # Target indices corresponding to your state space
  212.                 # In this example: 6 joints + 1 gripper for each arm
  213.                 UNI_STATE_INDICES =  [
  214.                     STATE_VEC_IDX_MAPPING[f"right_arm_joint_{i}_pos"] for i in range(6)
  215.                 # ] + [
  216.                     # STATE_VEC_IDX_MAPPING["right_gripper_open"]
  217.                 ]
  218.                 uni_vec = np.zeros(values.shape[:-1] + (self.STATE_DIM, ))
  219.                 uni_vec[..., UNI_STATE_INDICES] = values
  220.                 return uni_vec
  222.             state = fill_in_state(state)
  223.             state_indicator = fill_in_state(np.ones_like(state_std))
  224.             state_std = fill_in_state(state_std)
  225.             state_mean = fill_in_state(state_mean)
  226.             state_norm = fill_in_state(state_norm)
  227.             # If action's format is different from state's,
  228.             # you may implement fill_in_action()
  229.             actions = fill_in_state(actions)
  231.             # Parse the images
  232.             def parse_img(key):
  233.                 imgs = []
  234.                 for i in range(max(step_id - self.IMG_HISORY_SIZE + 1, 0), step_id + 1):
  235.                     img_bits = f["observations"]["images"][key][i]
  236.                     img = cv2.imdecode(np.frombuffer(img_bits, np.uint8), cv2.IMREAD_COLOR)
  237.                     imgs.append(img)
  238.                 imgs = np.stack(imgs)
  239.                 if imgs.shape[0] < self.IMG_HISORY_SIZE:
  240.                     # Pad the images using the first image
  241.                     imgs = np.concatenate(
  242.                         [
  243.                             np.tile(
  244.                                 imgs[:1],
  245.                                 (self.IMG_HISORY_SIZE - imgs.shape[0], 1, 1, 1),
  246.                             ),
  247.                             imgs,
  248.                         ],
  249.                         axis=0,
  250.                     )
  251.                 return imgs
  253.             # `cam_high` is the external camera image
  254.             cam_high = parse_img("cam_high")
  255.             # For step_id = first_idx - 1, the valid_len should be one
  256.             valid_len = min(step_id - (first_idx - 1) + 1, self.IMG_HISORY_SIZE)
  257.             cam_high_mask = np.array([False] * (self.IMG_HISORY_SIZE - valid_len) + [True] * valid_len)
  258.             # cam_left_wrist = parse_img("cam_left_wrist")
  259.             # cam_left_wrist_mask = cam_high_mask.copy()
  260.             cam_left_wrist = np.zeros((self.IMG_HISORY_SIZE, 0, 0, 0))#parse_img('cam_right_wrist')
  261.             cam_left_wrist_mask = np.array([False] * self.IMG_HISORY_SIZE)#cam_high_mask.copy()
  262.             cam_right_wrist = parse_img("cam_right_wrist")
  263.             cam_right_wrist_mask = cam_high_mask.copy()  # 使用相同的掩码逻辑
  265.             # Return the resulting sample
  266.             # For unavailable images, return zero-shape arrays, i.e., (IMG_HISORY_SIZE, 0, 0, 0)
  267.             # E.g., return np.zeros((self.IMG_HISORY_SIZE, 0, 0, 0)) for the key "cam_left_wrist",
  268.             # if the left-wrist camera is unavailable on your robot
  269.             return True, {
  270.                 "meta": meta,
  271.                 "state": state,
  272.                 "state_std": state_std,
  273.                 "state_mean": state_mean,
  274.                 "state_norm": state_norm,
  275.                 "actions": actions,
  276.                 "state_indicator": state_indicator,
  277.                 "cam_high": cam_high,
  278.                 "cam_high_mask": cam_high_mask,
  279.                 "cam_left_wrist": cam_left_wrist,
  280.                 "cam_left_wrist_mask": cam_left_wrist_mask,
  281.                 "cam_right_wrist": cam_right_wrist,
  282.                 "cam_right_wrist_mask": cam_right_wrist_mask,
  283.             }
  285.     def parse_hdf5_file_state_only(self, file_path):
  286.         """[Modify] Parse a hdf5 file to generate a state trajectory.
  288.         Args:
  289.             file_path (str): the path to the hdf5 file
  291.         Returns:
  292.             valid (bool): whether the episode is valid, which is useful for filtering.
  293.                 If False, this episode will be dropped.
  294.             dict: a dictionary containing the training sample,
  295.                 {
  296.                     "state": ndarray,           # state[:], (T, STATE_DIM).
  297.                     "action": ndarray,          # action[:], (T, STATE_DIM).
  298.                 } or None if the episode is invalid.
  299.         """
  300.         with h5py.File(file_path, "r") as f:
  301.             qpos = f["observations"]["qpos"][:]
  302.             left_arm_dim = f["observations"]["left_arm_dim"][:]
  303.             right_arm_dim = f["observations"]["right_arm_dim"][:]
  305.             num_steps = qpos.shape[0]
  306.             # [Optional] We drop too-short episode
  307.             # if num_steps < 128:
  308.             # return False, None
  310.             # [Optional] We skip the first few still steps
  311.             EPS = 1e-2
  312.             # Get the idx of the first qpos whose delta exceeds the threshold
  313.             qpos_delta = np.abs(qpos - qpos[0:1])
  314.             indices = np.where(np.any(qpos_delta > EPS, axis=1))[0]
  315.             if len(indices) > 0:
  316.                 first_idx = indices[0]
  317.             else:
  318.                 raise ValueError("Found no qpos that exceeds the threshold.")
  320.             # Rescale gripper to [0, 1]
  321.             # qpos = qpos / np.array([[1 for i in range(left_arm_dim[0] + right_arm_dim[0] + 2)]])
  322.             # target_qpos = f["action"][:] / np.array([[1 for i in range(left_arm_dim[0] + right_arm_dim[0] + 2)]])
  323.             
  324.             qpos = qpos / np.array(
  325.             #    [[1, 1, 1, 1, 1, 1, 4.7908, 1, 1, 1, 1, 1, 1, 4.7888]]
  326.             [[180, 180, 180, 180, 180, 180]]
  327.             )
  328.             target_qpos = f['action'][first_idx - 1:] / np.array(
  329.             #    [[1, 1, 1, 1, 1, 1, 11.8997, 1, 1, 1, 1, 1, 1, 13.9231]]
  330.             [[180, 180, 180, 180, 180, 180]]
  331.             )
  332.             # Parse the state and action
  333.             state = qpos[first_idx - 1:]
  334.             action = target_qpos[first_idx - 1:]
  335.             
  336.             # Standardize trajectory length to avoid batch size mismatch
  337.             # Use a fixed length (e.g., 128) or pad/truncate to match
  338.             target_length = 128  # You can adjust this value
  339.             if state.shape[0] > target_length:
  340.                 # Truncate to target length
  341.                 state = state[:target_length]
  342.                 action = action[:target_length]
  343.             elif state.shape[0] < target_length:
  344.                 # Pad with the last state/action
  345.                 pad_length = target_length - state.shape[0]
  346.                 state = np.concatenate([state, np.tile(state[-1:], (pad_length, 1))], axis=0)
  347.                 action = np.concatenate([action, np.tile(action[-1:], (pad_length, 1))], axis=0)
  349.             # Fill the state/action into the unified vector
  350.             def fill_in_state(values):
  351.                 # Target indices corresponding to your state space
  352.                 # In this example: 6 joints + 1 gripper for each arm
  353.                 UNI_STATE_INDICES =  [
  354.                     STATE_VEC_IDX_MAPPING[f"right_arm_joint_{i}_pos"] for i in range(6)
  355.                 # ] + [
  356.                     # STATE_VEC_IDX_MAPPING["right_gripper_open"]
  357.                 ]
  358.                 uni_vec = np.zeros(values.shape[:-1] + (self.STATE_DIM, ))
  359.                 uni_vec[..., UNI_STATE_INDICES] = values
  360.                 return uni_vec
  362.             state = fill_in_state(state)
  363.             action = fill_in_state(action)
  365.             # Return the resulting sample
  366.             return True, {"state": state, "action": action}
  368. if __name__ == "__main__":
  369.     ds = HDF5VLADataset()
  370.     for i in range(len(ds)):
  371.         print(f"Processing episode {i}/{len(ds)}...")
  372.         ds.get_item(i)
复制代码
        接着直接开始训练：
  1. # step1：安装torch、torchvision
  2. pip install torch==2.1.0 torchvision==0.16.0  --index-url https://download.pytorch.org/whl/cu121
  3. # step2：安装packaging
  4. pip install packaging==24.0
  5. # step3：安装其他依赖
  6. pip install -r requirements.txt
复制代码
RDT170M微调

        要训练170M的版本我们需要修改RDT/configs/base.yaml文件中model类下的rdt参数
  1. $ python3 -c "import torch ; print(torch.__config__.show())"
复制代码
        接着生成训练的参数yml文件，并将pretrained_model_name_or_path指向我们先前下载的170M模型
  1. PyTorch built with:
  2.   - GCC 9.3
  3.   - C++ Version: 201703
  4.   - Intel(R) oneAPI Math Kernel Library Version 2022.2-Product Build 20220804 for Intel(R) 64 architecture applications
  5.   - Intel(R) MKL-DNN v3.4.2 (Git Hash 1137e04ec0b5251ca2b4400a4fd3c667ce843d67)
  6.   - OpenMP 201511 (a.k.a. OpenMP 4.5)
  7.   - LAPACK is enabled (usually provided by MKL)
  8.   - NNPACK is enabled
  9.   - CPU capability usage: AVX2
  10.   - CUDA Runtime 12.1
  11.   - NVCC architecture flags: -gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_90,code=sm_90
  12.   - CuDNN 90.1  (built against CUDA 12.4)
  13.   - Magma 2.6.1
  14.   - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=12.1, CUDNN_VERSION=9.1.0, CXX_COMPILER=/opt/rh/devtoolset-9/root/usr/bin/c++, CXX_FLAGS= -D_GLIBCXX_USE_CXX11_ABI=0 -fabi-version=11 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wsuggest-override -Wno-psabi -Wno-error=pedantic -Wno-error=old-style-cast -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=2.4.1, USE_CUDA=ON, USE_CUDNN=ON, USE_CUSPARSELT=1, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_GLOO=ON, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=1, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF, USE_ROCM_KERNEL_ASSERT=OFF,
复制代码
        接着直接开始训练：
  1. pip3 install flash_attn-2.7.2.post1+cu12torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
复制代码
        在双卡A100-40GB上BS16即单卡8BS的显存占用及训练速度参考如下，按照RDT论文中的说法关注overall_avg_sample_mse指标，RDT170M和1B的版本在数据量仅有100条的时候均能在7000步左右实现拟合指标下降到0.0001量级
4.png

5.png

四、实际评测

        在训练完成之后我们便可以开始实机评测啦，由于我们目前需要在Mac上连接LeRobot机械臂进行控制，因此我们在实际使用A100或者是RDKS100部署推理的时候还需要完成两端之间的通信代码，我们在这里就用最简单的Socket来实现ServerClient，还有其他更优的ZMQ等方式就不在这里呈现了，具体实现的代码如下，这份代码需要同时放到本地及推理服务器端进行调用
  1. # 法一： 直接运行我仓库中写好的脚本
  2. cd weights/RDT
  3. bash _download.sh
  4. # 法二： 手动下载
  5. export HF_ENDPOINT=https://hf-mirror.com # 国内镜像，加速下载
  6. huggingface-cli download google/t5-v1_1-xxl --local-dir t5-v1_1-xxl
  7. huggingface-cli download google/siglip-so400m-patch14-384 --local-dir siglip-so400m-patch14-384
  8. huggingface-cli download robotics-diffusion-transformer/rdt-1b --local-dir rdt-1b
  9. huggingface-cli download robotics-diffusion-transformer/rdt-170m --local-dir rdt-170m
复制代码
服务器单卡部署

       接着我们便可以完成我们服务器上的推理代码了，我们参考RDT中的RDT/scripts/agilex_model.py来完成我们的lerobot_rdt_server，我们在这个代码中完成对RoboticDiffusionTransformerModel类的修改同时使用类中的step执行推理，并集成我们的ServerClient来接收本地电脑发来的机械臂observation数据，具体的代码实现如下：
  1. rdt:
  2. # 1B: num_head 32 hidden_size 2048 depth: 28
  3. # 170M: num_head 32 hidden_size 1024 depth: 14
  4. hidden_size: 2048
  5. depth: 28
  6. num_heads: 32
  7. cond_pos_embed_type: multimodal
复制代码
        新版的LeRobot使用函数来替换了老的observation获取函数，同时其中的数据结构也进行了更改，因此我们在record的基础上直接复制一份使用最小的状态获取及机械臂控制实例完成本地的数据传输及通信控制代码，具体代码如下，文件位置为lerobot/src/lerobot/record_rdt.py
  1. bash generate.sh RDT1B_LeRobot
  2. # Generated on 2025-08-28 17:14:20
  3. model: RDT1B_LeRobot
  4. data_path: training_data/RDT1B_LeRobot
  5. checkpoint_path: checkpoints/RDT1B_LeRobot
  6. pretrained_model_name_or_path: ../weights/RDT/rdt-1b
  7. cuda_visible_device: '0,1,2,3'
  8. train_batch_size: 16
  9. sample_batch_size: 32
  10. max_train_steps: 10000
  11. checkpointing_period: 2500
  12. sample_period: 100
  13. checkpoints_total_limit: 40
  14. learning_rate: 0.0001
  15. dataloader_num_workers: 8
  16. state_noise_snr: 40
  17. gradient_accumulation_steps: 1
复制代码
        接着按照你的实际摄像头串口配置修改main()中的对应配置即可，接着在依次运行服务端即本地客户端即可看到机械臂开始运动完成任务：
  1. bash finetune.sh RDT1B_LeRobot
复制代码
RDKS100部署

        目前RDKS100上仅支持RDT170M的部署，接下来我们参考文档：RDT on Double RDK S100P 全流程文档，一步步的完成LeRobot的RDT的上板流程，我们首先使用仓库中RDT目录下的脚本RDT/export_all.py来完成RDT中所有ONNX模型的导出，这段脚本有如下的参数可以进行配置，大家仅需按照自己的实际情况配置即可
  1. rdt:
  2. # 1B: num_head 32 hidden_size 2048 depth: 28
  3. # 170M: num_head 32 hidden_size 1024 depth: 14
  4. hidden_size: 1024
  5. depth: 14
  6. num_heads: 32
  7. cond_pos_embed_type: multimodal
复制代码
        接下来我们执行这段脚本之后在RDT的目录下会生成脚本导出的ONNX模型以及dump出的校准数据和文件
  1. bash generate.sh RDT170M_LeRobot
  2. # Generated on 2025-08-28 17:14:20
  3. model: RDT170M_LeRobot
  4. data_path: training_data/RDT170M_LeRobot
  5. checkpoint_path: checkpoints/RDT170M_LeRobot
  6. pretrained_model_name_or_path: ../weights/RDT/rdt-170m
  7. cuda_visible_device: '0,1'
  8. train_batch_size: 16
  9. sample_batch_size: 32
  10. max_train_steps: 10000
  11. checkpointing_period: 2500
  12. sample_period: 100
  13. checkpoints_total_limit: 40
  14. learning_rate: 0.0001
  15. dataloader_num_workers: 8
  16. state_noise_snr: 40
  17. gradient_accumulation_steps: 1
复制代码
        接着我们使用RDKS100算法工具链标准交付的docker环境，来完成BPU模型的量化和编译，docker的安装挂载和使用命令参考以下：
  1. bash finetune.sh RDT170M_LeRobot
复制代码
        在Docker内输入bash build_all.sh后便会自动开启编译并导出可以在RDKS100板端运行的HBM模型
6.png

        以下是DiT、img_adaptor、lang_adaptor以及state_adaptor模型量化过程中的部分指标参考：
  1. import socket
  2. import numpy as np
  3. import zlib
  4. import json
  5. import base64
  6. import time
  7. from typing import Any
  8. import torch
  10. class NumpyEncoder(json.JSONEncoder):
  11.     """Enhanced json encoder for numpy types and PyTorch tensors with array reconstruction info"""
  12.     def default(self, obj):
  13.         if isinstance(obj, np.ndarray):
  14.             return {
  15.                 '__numpy_array__': True,
  16.                 'data': base64.b64encode(obj.tobytes()).decode('ascii'),
  17.                 'dtype': str(obj.dtype),
  18.                 'shape': obj.shape
  19.             }
  20.         elif torch is not None and isinstance(obj, torch.Tensor):
  21.             # 将 PyTorch Tensor 转换为 numpy 数组
  22.             numpy_array = obj.cpu().detach().numpy()
  23.             return {
  24.                 '__numpy_array__': True,
  25.                 'data': base64.b64encode(numpy_array.tobytes()).decode('ascii'),
  26.                 'dtype': str(numpy_array.dtype),
  27.                 'shape': numpy_array.shape
  28.             }
  29.         elif isinstance(obj, (np.integer, np.floating, np.bool_)):
  30.             return obj.item()
  31.         return super().default(obj)
  33. def numpy_to_json(data: Any) -> str:
  34.     return json.dumps(data, cls=NumpyEncoder)
  36. def json_to_numpy(json_str: str) -> Any:
  37.     def hook(dct):
  38.         if '__numpy_array__' in dct:
  39.             data = base64.b64decode(dct['data'])
  40.             return np.frombuffer(data, dtype=dct['dtype']).reshape(dct['shape'])
  41.         return dct
  42.     return json.loads(json_str, object_hook=hook)
  44. class CommonUtils:
  45.     @staticmethod
  46.     def serialize(data: Any) -> bytes:
  47.         return zlib.compress(numpy_to_json(data).encode('utf-8'))
  49.     @staticmethod
  50.     def deserialize(data: bytes) -> Any:
  51.         return json_to_numpy(zlib.decompress(data).decode('utf-8'))
  53. def send_all(sock, payload):
  54.     sock.sendall(len(payload).to_bytes(8, 'big') + payload)
  56. def recv_all(sock) :
  57.     length_bytes = sock.recv(8)
  58.     if not length_bytes:
  59.         return None
  60.     length = int.from_bytes(length_bytes, 'big')
  61.     buf = b''
  62.     while len(buf) < length:
  63.         chunk = sock.recv(length - len(buf))
  64.         if not chunk:
  65.             return None
  66.         buf += chunk
  67.     return buf
  69. class ServerClient:
  70.     def __init__(self, host='localhost', port=5000, is_server=True):
  71.         self.host, self.port, self.is_server = host, port, is_server
  72.         self.utils = CommonUtils()
  73.         self._connect()
  75.     def _connect(self):
  76.         if self.is_server:
  77.             self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  78.             self.sock.bind((self.host, self.port))
  79.             self.sock.listen(1)
  80.             print(f"[ServerClient] Listening on {self.host}:{self.port}")
  81.             self.conn, addr = self.sock.accept()
  82.             print(f"[ServerClient] Connected by {addr}")
  83.         else:
  84.             while True:
  85.                 try:
  86.                     self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  87.                     self.sock.connect((self.host, self.port))
  88.                     self.conn = self.sock
  89.                     print(f"[ServerClient] Connected to {self.host}:{self.port}")
  90.                     break
  91.                 except (ConnectionRefusedError, OSError):
  92.                     print("[ServerClient] Waiting for server...")
  93.                     time.sleep(2)
  95.     def send(self, data):
  96.         payload = self.utils.serialize(data)
  97.         try:
  98.             send_all(self.conn, payload)
  99.         except (BrokenPipeError, ConnectionResetError, OSError):
  100.             print("[ServerClient] Connection lost. Reconnecting...")
  101.             self._connect()
  102.             send_all(self.conn, payload)
  104.     def receive(self):
  105.         try:
  106.             buf = recv_all(self.conn)
  107.             return self.utils.deserialize(buf) if buf else None
  108.         except (BrokenPipeError, ConnectionResetError, OSError):
  109.             print("[ServerClient] Connection lost. Reconnecting...")
  110.             self._connect()
  111.             return None
  113.     def close(self):
  114.         self.conn.close()
  115.         self.sock.close()
  117. class Client:
  118.     def __init__(self, host='127.0.0.1', port=5000):
  119.         self.host, self.port = host, port
  120.         self.utils = CommonUtils()
  121.         self.connect()
  123.     def connect(self):
  124.         while True:
  125.             try:
  126.                 self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  127.                 self.sock.connect((self.host, self.port))
  128.                 print(f"[Client] Connected to {self.host}:{self.port}")
  129.                 break
  130.             except (ConnectionRefusedError, OSError):
  131.                 print("[Client] Waiting for server...")
  132.                 time.sleep(2)
  134.     def send(self, data):
  135.         payload = self.utils.serialize(data)
  136.         try:
  137.             send_all(self.sock, payload)
  138.         except (BrokenPipeError, ConnectionResetError, OSError):
  139.             print("[Client] Connection lost. Reconnecting...")
  140.             self.connect()
  141.             send_all(self.sock, payload)
  143.     def receive(self):
  144.         try:
  145.             buf = recv_all(self.sock)
  146.             return self.utils.deserialize(buf) if buf else None
  147.         except (BrokenPipeError, ConnectionResetError, OSError):
  148.             print("[Client] Connection lost. Reconnecting...")
  149.             self.connect()
  150.             return None
  152.     def close(self):
  153.         self.sock.close()
  154.         print("[Client] Closed.")
  156. class Server:
  157.     def __init__(self, host='0.0.0.0', port=5000):
  158.         self.host, self.port = host, port
  159.         self.utils = CommonUtils()
  160.         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  161.         self.sock.bind((self.host, self.port))
  162.         self.sock.listen(1)
  163.         print(f"[Server] Listening on {self.host}:{self.port}")
  164.         self._wait_client()
  166.     def _wait_client(self):
  167.         print("[Server] Waiting for client...")
  168.         self.conn, addr = self.sock.accept()
  169.         print(f"[Server] Connected by {addr}")
  171.     def send(self, data: Any):
  172.         payload = self.utils.serialize(data)
  173.         try:
  174.             send_all(self.conn, payload)
  175.         except (BrokenPipeError, ConnectionResetError, OSError):
  176.             print("[Server] Client disconnected. Waiting new client...")
  177.             self._wait_client()
  178.             send_all(self.conn, payload)
  180.     def receive(self):
  181.         try:
  182.             buf = recv_all(self.conn)
  183.             return self.utils.deserialize(buf) if buf else None
  184.         except (BrokenPipeError, ConnectionResetError, OSError):
  185.             print("[Server] Client disconnected. Waiting new client...")
  186.             self._wait_client()
  187.             return None
  189.     def close(self):
  190.         self.conn.close()
  191.         self.sock.close()
  192.         print("[Server] Closed.")
复制代码
        接着我们运行以下命令下载已经量化编译好的SigLip到编译产物结果文件夹BPU_RDT_Policy中并将这个文件夹和test_data文件夹以及LeRobot-VLA仓库中RDKS_ModelRun/RDT路径下的所有文件复制到我们的RDKS100板端如图所示：
  1. import os, sys
  3. import numpy as np
  4. import torch
  5. from PIL import Image
  6. from torchvision import transforms
  7. import yaml
  9. from pathlib import Path
  11. # get current workspace
  12. current_file = Path(__file__)
  13. sys.path.append(os.path.join(current_file.parent.parent, "models"))
  14. sys.path.append(os.path.join(current_file.parent.parent, "models"))
  15. sys.path.append(os.path.join(current_file.parent.parent))  
  16. from configs.state_vec import STATE_VEC_IDX_MAPPING
  17. from multimodal_encoder.siglip_encoder import SiglipVisionTower
  18. from multimodal_encoder.t5_encoder import T5Embedder
  19. from rdt_runner import RDTRunner
  20. from server_client import ServerClient
  22. # The indices that the raw vector should be mapped to in the unified action vector
  23. AGILEX_STATE_INDICES = [
  24.     STATE_VEC_IDX_MAPPING[f"right_arm_joint_{i}_pos"] for i in range(6)
  25. ]
  27. # Create the RDT model
  28. def create_model(args, **kwargs):
  29.     model = RoboticDiffusionTransformerModel(args, **kwargs)
  30.     pretrained = kwargs.get("pretrained", None)
  31.     if pretrained is not None and os.path.isfile(pretrained):
  32.         model.load_pretrained_weights(pretrained)
  34.     return model
  36. class RoboticDiffusionTransformerModel(object):
  37.     """A wrapper for the RDT model, which handles
  38.     1. Model initialization
  39.     2. Encodings of instructions
  40.     3. Model inference
  41.     """
  43.     def __init__(
  44.         self,
  45.         args,
  46.         device="cuda",
  47.         dtype=torch.bfloat16,
  48.         image_size=None,
  49.         control_frequency=25,
  50.         pretrained=None,
  51.         pretrained_vision_encoder_name_or_path=None,
  52.     ):
  53.         self.args = args
  54.         self.dtype = dtype
  55.         self.image_size = image_size
  56.         self.device = device
  57.         self.control_frequency = control_frequency
  58.         # We do not use the text encoder due to limited GPU memory
  59.         # self.text_tokenizer, self.text_model = self.get_text_encoder(pretrained_text_encoder_name_or_path)
  60.         self.image_processor, self.vision_model = self.get_vision_encoder(pretrained_vision_encoder_name_or_path)
  61.         self.policy = self.get_policy(pretrained)
  63.         self.reset()
  65.     def get_policy(self, pretrained):
  66.         """Initialize the model."""
  67.         # Initialize model with arguments
  68.         if pretrained is None or os.path.isfile(pretrained):
  69.             img_cond_len = (self.args["common"]["img_history_size"] * self.args["common"]["num_cameras"] *
  70.                             self.vision_model.num_patches)
  72.             _model = RDTRunner(
  73.                 action_dim=self.args["common"]["state_dim"],
  74.                 pred_horizon=self.args["common"]["action_chunk_size"],
  75.                 config=self.args["model"],
  76.                 lang_token_dim=self.args["model"]["lang_token_dim"],
  77.                 img_token_dim=self.args["model"]["img_token_dim"],
  78.                 state_token_dim=self.args["model"]["state_token_dim"],
  79.                 max_lang_cond_len=self.args["dataset"]["tokenizer_max_length"],
  80.                 img_cond_len=img_cond_len,
  81.                 img_pos_embed_config=[
  82.                     # No initial pos embed in the last grid size
  83.                     # since we've already done in ViT
  84.                     (
  85.                         "image",
  86.                         (
  87.                             self.args["common"]["img_history_size"],
  88.                             self.args["common"]["num_cameras"],
  89.                             -self.vision_model.num_patches,
  90.                         ),
  91.                     ),
  92.                 ],
  93.                 lang_pos_embed_config=[
  94.                     # Similarly, no initial pos embed for language
  95.                     ("lang", -self.args["dataset"]["tokenizer_max_length"]),
  96.                 ],
  97.                 dtype=self.dtype,
  98.             )
  99.         else:
  100.             _model = RDTRunner.from_pretrained(pretrained)
  102.         return _model
  104.     def get_text_encoder(self, pretrained_text_encoder_name_or_path):
  105.         text_embedder = T5Embedder(
  106.             from_pretrained=pretrained_text_encoder_name_or_path,
  107.             model_max_length=self.args["dataset"]["tokenizer_max_length"],
  108.             device=self.device,
  109.         )
  110.         tokenizer, text_encoder = text_embedder.tokenizer, text_embedder.model
  111.         return tokenizer, text_encoder
  113.     def get_vision_encoder(self, pretrained_vision_encoder_name_or_path):
  114.         vision_encoder = SiglipVisionTower(vision_tower=pretrained_vision_encoder_name_or_path, args=None)
  115.         image_processor = vision_encoder.image_processor
  116.         return image_processor, vision_encoder
  118.     def reset(self):
  119.         """Set model to evaluation mode."""
  120.         device = self.device
  121.         weight_dtype = self.dtype
  122.         self.policy.eval()
  123.         # self.text_model.eval()
  124.         self.vision_model.eval()
  126.         self.policy = self.policy.to(device, dtype=weight_dtype)
  127.         # self.text_model = self.text_model.to(device, dtype=weight_dtype)
  128.         self.vision_model = self.vision_model.to(device, dtype=weight_dtype)
  130.     def load_pretrained_weights(self, pretrained=None):
  131.         if pretrained is None:
  132.             return
  133.         print(f"Loading weights from {pretrained}")
  134.         filename = os.path.basename(pretrained)
  135.         if filename.endswith(".pt"):
  136.             checkpoint = torch.load(pretrained)
  137.             self.policy.load_state_dict(checkpoint["module"])
  138.         elif filename.endswith(".safetensors"):
  139.             from safetensors.torch import load_model
  141.             load_model(self.policy, pretrained)
  142.         else:
  143.             raise NotImplementedError(f"Unknown checkpoint format: {pretrained}")
  145.     def encode_instruction(self, instruction, device="cuda"):
  146.         """Encode string instruction to latent embeddings.
  148.         Args:
  149.             instruction: a string of instruction
  150.             device: a string of device
  152.         Returns:
  153.             pred: a tensor of latent embeddings of shape (text_max_length, 512)
  154.         """
  155.         tokens = self.text_tokenizer(instruction, return_tensors="pt", padding="longest",
  156.                                      truncation=True)["input_ids"].to(device)
  158.         tokens = tokens.view(1, -1)
  159.         with torch.no_grad():
  160.             pred = self.text_model(tokens).last_hidden_state.detach()
  162.         return pred
  164.     def _format_joint_to_state(self, joints):
  165.         """
  166.         Format the joint proprioception into the unified action vector.
  168.         Args:
  169.             joints (torch.Tensor): The joint proprioception to be formatted.
  170.                 qpos ([B, N, 14]).
  172.         Returns:
  173.             state (torch.Tensor): The formatted vector for RDT ([B, N, 128]).
  174.         """
  175.         # Rescale the gripper to the range of [0, 1]
  176.         joints = joints / torch.tensor(
  177.             [[[180, 180, 180, 180, 180, 180]]],
  178.             device=joints.device,
  179.             dtype=joints.dtype,
  180.         )
  182.         B, N, _ = joints.shape
  183.         state = torch.zeros(
  184.             (B, N, self.args["model"]["state_token_dim"]),
  185.             device=joints.device,
  186.             dtype=joints.dtype,
  187.         )
  188.         # Fill into the unified state vector
  189.         state[:, :, AGILEX_STATE_INDICES] = joints
  190.         # Assemble the mask indicating each dimension's availability
  191.         state_elem_mask = torch.zeros(
  192.             (B, self.args["model"]["state_token_dim"]),
  193.             device=joints.device,
  194.             dtype=joints.dtype,
  195.         )
  196.         state_elem_mask[:, AGILEX_STATE_INDICES] = 1
  197.         return state, state_elem_mask
  199.     def _unformat_action_to_joint(self, action):
  200.         """
  201.         Unformat the unified action vector into the joint action to be executed.
  203.         Args:
  204.             action (torch.Tensor): The unified action vector to be unformatted.
  205.                 ([B, N, 128])
  207.         Returns:
  208.             joints (torch.Tensor): The unformatted robot joint action.
  209.                 qpos ([B, N, 14]).
  210.         """
  211.         action_indices = AGILEX_STATE_INDICES
  212.         joints = action[:, :, action_indices]
  214.         # Rescale the gripper back to the action range
  215.         # Note that the action range and proprioception range are different
  216.         # for Mobile ALOHA robot
  217.         joints = joints * torch.tensor(
  218.             [[[180, 180, 180, 180, 180, 180]]],
  219.             device=joints.device,
  220.             dtype=joints.dtype,
  221.         )
  223.         return joints
  225.     @torch.no_grad()
  226.     def step(self, proprio, images, text_embeds):
  227.         """
  228.         Predict the next action chunk given the
  229.         proprioceptive states, images, and instruction embeddings.
  231.         Args:
  232.             proprio: proprioceptive states
  233.             images: RGB images, the order should be
  234.                 [ext_{t-1}, right_wrist_{t-1}, left_wrist_{t-1},
  235.                 ext_{t}, right_wrist_{t}, left_wrist_{t}]
  236.             text_embeds: instruction embeddings
  238.         Returns:
  239.             action: predicted action
  240.         """
  241.         device = self.device
  242.         dtype = self.dtype
  244.         # The background image used for padding
  245.         background_color = np.array([int(x * 255) for x in self.image_processor.image_mean],
  246.                                     dtype=np.uint8).reshape(1, 1, 3)
  247.         background_image = (np.ones(
  248.             (
  249.                 self.image_processor.size["height"],
  250.                 self.image_processor.size["width"],
  251.                 3,
  252.             ),
  253.             dtype=np.uint8,
  254.         ) * background_color)
  256.         # Preprocess the images by order and encode them
  257.         image_tensor_list = []
  258.         for image in images:
  259.             if image is None:
  260.                 # Replace it with the background image
  261.                 image = Image.fromarray(background_image)
  262.             else:
  263.                 # Convert numpy array to PIL Image if needed
  264.                 if isinstance(image, np.ndarray):
  265.                     image = Image.fromarray(image)
  267.             if self.image_size is not None:
  268.                 image = transforms.Resize(self.image_size)(image)
  270.             if self.args["dataset"].get("auto_adjust_image_brightness", False):
  271.                 pixel_values = list(image.getdata())
  272.                 average_brightness = sum(sum(pixel) for pixel in pixel_values) / (len(pixel_values) * 255.0 * 3)
  273.                 if average_brightness <= 0.15:
  274.                     image = transforms.ColorJitter(brightness=(1.75, 1.75))(image)
  276.             if self.args["dataset"].get("image_aspect_ratio", "pad") == "pad":
  278.                 def expand2square(pil_img, background_color):
  279.                     width, height = pil_img.size
  280.                     if width == height:
  281.                         return pil_img
  282.                     elif width > height:
  283.                         result = Image.new(pil_img.mode, (width, width), background_color)
  284.                         result.paste(pil_img, (0, (width - height) // 2))
  285.                         return result
  286.                     else:
  287.                         result = Image.new(pil_img.mode, (height, height), background_color)
  288.                         result.paste(pil_img, ((height - width) // 2, 0))
  289.                         return result
  291.                 image = expand2square(image, tuple(int(x * 255) for x in self.image_processor.image_mean))
  292.             image = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
  293.             image_tensor_list.append(image)
  295.         image_tensor = torch.stack(image_tensor_list, dim=0).to(device, dtype=dtype)
  297.         image_embeds = self.vision_model(image_tensor).detach()
  298.         image_embeds = image_embeds.reshape(-1, self.vision_model.hidden_size).unsqueeze(0)
  300.         # Prepare the proprioception states and the control frequency
  301.         # Convert numpy array to tensor if needed
  302.         if isinstance(proprio, np.ndarray):
  303.             # Copy the array to make it writable
  304.             proprio = torch.from_numpy(proprio.copy())
  305.         
  306.         joints = proprio.to(device).unsqueeze(0)  # (1, 1, 14)
  307.         states, state_elem_mask = self._format_joint_to_state(joints)  # (1, 1, 128), (1, 128)
  308.         states, state_elem_mask = states.to(device, dtype=dtype), state_elem_mask.to(device, dtype=dtype)
  309.         states = states[:, -1:, :]  # (1, 1, 128)
  310.         ctrl_freqs = torch.tensor([self.control_frequency]).to(device)
  312.         text_embeds = text_embeds.to(device, dtype=dtype)
  314.         # Predict the next action chunk given the inputs
  315.         trajectory = self.policy.predict_action(
  316.             lang_tokens=text_embeds,
  317.             lang_attn_mask=torch.ones(text_embeds.shape[:2], dtype=torch.bool, device=text_embeds.device),
  318.             img_tokens=image_embeds,
  319.             state_tokens=states,
  320.             action_mask=state_elem_mask.unsqueeze(1),
  321.             ctrl_freqs=ctrl_freqs,
  322.         )
  323.         trajectory = self._unformat_action_to_joint(trajectory).to(torch.float32)
  325.         return trajectory
  327. class LERobotRDTServer:
  328.     def __init__(self, pretrained_vision_encoder_name_or_path, pretrained, args, lang_model):
  329.         self.policy = create_model(
  330.             args=args,
  331.             dtype=torch.bfloat16,
  332.             pretrained=pretrained,
  333.             pretrained_vision_encoder_name_or_path=pretrained_vision_encoder_name_or_path,
  334.             control_frequency=30,
  335.         )
  336.         self.server = ServerClient(host="0.0.0.0", port=5002, is_server=True)
  337.         
  338.         # Load and debug language embeddings
  339.         self.lang_embeddings = torch.load(lang_model)
  340.         print(f"Loaded language embeddings shape: {self.lang_embeddings.shape}")
  341.         print(f"Model expects tokenizer_max_length: {self.policy.args['dataset']['tokenizer_max_length']}")
  342.         print(f"Model lang_token_dim: {self.policy.args['model']['lang_token_dim']}")
  343.         
  344.         # Check if dimensions match
  345.         expected_seq_len = self.policy.args["dataset"]["tokenizer_max_length"]
  346.         expected_hidden_dim = self.policy.args["model"]["lang_token_dim"]
  347.         
  348.         # Handle different embedding formats
  349.         if len(self.lang_embeddings.shape) == 2:
  350.             # Format: [seq_len, hidden_dim]
  351.             actual_seq_len, actual_hidden_dim = self.lang_embeddings.shape
  352.             if actual_seq_len != expected_seq_len:
  353.                 print(f"WARNING: Sequence length mismatch! Expected {expected_seq_len}, got {actual_seq_len}")
  354.             if actual_hidden_dim != expected_hidden_dim:
  355.                 print(f"WARNING: Hidden dimension mismatch! Expected {expected_hidden_dim}, got {actual_hidden_dim}")
  356.         elif len(self.lang_embeddings.shape) == 3:
  357.             # Format: [batch_size, seq_len, hidden_dim]
  358.             actual_batch, actual_seq_len, actual_hidden_dim = self.lang_embeddings.shape
  359.             if actual_seq_len != expected_seq_len:
  360.                 print(f"WARNING: Sequence length mismatch! Expected {expected_seq_len}, got {actual_seq_len}")
  361.             if actual_hidden_dim != expected_hidden_dim:
  362.                 print(f"WARNING: Hidden dimension mismatch! Expected {expected_hidden_dim}, got {actual_hidden_dim}")
  363.         else:
  364.             print(f"WARNING: Unexpected embedding shape: {self.lang_embeddings.shape}")
  365.         
  366.     def run(self):
  367.         print("LERobot RDT Server started, waiting for messages...")
  368.         try:
  369.             while True:
  370.                 print("Waiting for RDT data...")
  371.                 rdt_data = self.server.receive()
  372.                 print(f"Received RDT data, message_id: {rdt_data['message_id']}")
  373.                
  374.                 # Perform inference
  375.                 # Ensure language embeddings have correct shape
  376.                 if len(self.lang_embeddings.shape) == 2:
  377.                     # [seq_len, hidden_dim] -> [1, seq_len, hidden_dim]
  378.                     text_embeds = self.lang_embeddings.unsqueeze(0)
  379.                 else:
  380.                     # Already [batch_size, seq_len, hidden_dim]
  381.                     text_embeds = self.lang_embeddings
  382.                
  383.                 action = self.policy.step(
  384.                     proprio=rdt_data["proprio"],
  385.                     images=rdt_data["images"],
  386.                     text_embeds=text_embeds,
  387.                 )
  388.                
  389.                 # Prepare response - use 'actions' key to match client expectation
  390.                 message_id = rdt_data["message_id"]
  391.                 action_data = {
  392.                     "message_id": message_id,
  393.                     "actions": action,  # Changed from 'action' to 'actions'
  394.                 }
  395.                
  396.                 # Send response
  397.                 print(f"send action data, action_data: {action_data}")
  398.                 self.server.send(action_data)
  399.                 print(f"Sent action data for message_id: {message_id}")
  400.                
  401.         except KeyboardInterrupt:
  402.             print("\nServer stopped by user")
  403.             self.server.close()
  404.         except Exception as e:
  405.             print(f"Error in server loop: {e}")
  406.             self.server.close()
  407.             raise
  409. if __name__ == "__main__":
  410.     path_to_rdt_model_wights = "/home/qi.xiong/DualArm/RoboTwin/policy/RDT/checkpoints/RDT_LeRobot/checkpoint-7500/pytorch_model/mp_rank_00_model_states.pt"
  411.     path_to_vision_encoder_model = "/home/qi.xiong/DualArm/RoboTwin/policy/weights/RDT/siglip-so400m-patch14-384"
  412.     lang_model = "/home/qi.xiong/DualArm/RoboTwin/policy/RDT/scripts/lerobot_rdt_data/greenmarker_scene1/episode_4/instructions/lang_embed_0.pt"
  413.     with open("/home/qi.xiong/DualArm/RoboTwin/policy/RDT/configs/base.yaml", "r") as fp:
  414.         config = yaml.safe_load(fp)
  415.     rdt_server = LERobotRDTServer(path_to_vision_encoder_model, path_to_rdt_model_wights, config, lang_model)
  416.     rdt_server.run()
复制代码
7.png

        接着我们直接运行以下命令即可首先使用校准测试数据验证模型推理是否成功，若无报错即证明模型可正常推理无异常，若运行过程中缺少某个依赖，直接安装即可：
  1. import logging
  2. import time
  3. from dataclasses import asdict, dataclass
  4. from pathlib import Path
  5. from pprint import pformat
  6. from collections import deque
  7. import torch
  8. from PIL import Image
  9. import numpy as np
  10. from lerobot.cameras import (  # noqa: F401
  11.     CameraConfig,  # noqa: F401
  12. )
  13. from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
  14. from lerobot.configs import parser
  15. from lerobot.configs.policies import PreTrainedConfig
  16. from lerobot.datasets.image_writer import safe_stop_image_writer
  17. from lerobot.datasets.lerobot_dataset import LeRobotDataset
  18. from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
  19. from lerobot.datasets.video_utils import VideoEncodingManager
  20. from lerobot.policies.factory import make_policy
  21. from lerobot.policies.pretrained import PreTrainedPolicy
  22. from lerobot.robots import (  # noqa: F401
  23.     Robot,
  24.     RobotConfig,
  25.     bi_so100_follower,
  26.     hope_jr,
  27.     koch_follower,
  28.     make_robot_from_config,
  29.     so100_follower,
  30.     so101_follower,
  31. )
  32. from lerobot.robots.so101_follower.so101_follower import SO101Follower
  33. from lerobot.robots.so101_follower.config_so101_follower import SO101FollowerConfig
  34. from lerobot.teleoperators import (  # noqa: F401
  35.     Teleoperator,
  36.     TeleoperatorConfig,
  37.     bi_so100_leader,
  38.     homunculus,
  39.     koch_leader,
  40.     make_teleoperator_from_config,
  41.     so100_leader,
  42.     so101_leader,
  43. )
  44. from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop
  45. from lerobot.utils.control_utils import (
  46.     init_keyboard_listener,
  47.     is_headless,
  48.     predict_action,
  49.     sanity_check_dataset_name,
  50.     sanity_check_dataset_robot_compatibility,
  51. )
  52. from lerobot.utils.robot_utils import busy_wait
  53. from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
  54. from server_client import *
  56. debug_save_img = True
  58. _action_queue = deque([], maxlen=64)
  59. _message_id = 0
  60. _last_cam_high = None  
  61. _last_cam_right_wrist = None  
  63. @safe_stop_image_writer
  64. def record_loop(
  65.     robot: Robot,
  66.     client: Client,
  67. ):
  68.     global _action_queue, _message_id, _last_cam_high, _last_cam_right_wrist
  70.     observation = robot.get_observation()
  72.     cam_high = observation['high']
  73.     cam_right_wrist = observation['arm']
  74.     image_arrs = [
  75.         _last_cam_high,
  76.         _last_cam_right_wrist,
  77.         None,
  78.         cam_high,
  79.         cam_right_wrist,
  80.         None
  81.     ]
  83.     images = [arr if arr is not None else None
  84.             for arr in image_arrs]
  85.     joint_positions = [observation[key] for key in observation.keys() if key.endswith('.pos')]
  86.     proprio = torch.tensor(joint_positions, dtype=torch.float32).unsqueeze(0)
  88. ###################Debug图像########################
  89.     if debug_save_img:
  90.         imgs_to_show = [cam_high, cam_right_wrist, _last_cam_high, _last_cam_right_wrist]
  91.         if all(img is not None for img in imgs_to_show):
  92.             pil_imgs = []
  93.             for img in imgs_to_show:
  94.                 if img.dtype != np.uint8:
  95.                     img = np.clip(img, 0, 1)
  96.                     img = (img * 255).astype(np.uint8)
  97.                 if img.ndim == 2:
  98.                     img = np.stack([img]*3, axis=-1)  
  99.                 elif img.shape[-1] == 1:
  100.                     img = np.repeat(img, 3, axis=-1)
  101.                 pil_imgs.append(Image.fromarray(img))
  102.             w, h = pil_imgs[0].size
  103.             for i in range(4):
  104.                 if pil_imgs[i].size != (w, h):
  105.                     pil_imgs[i] = pil_imgs[i].resize((w, h))
  106.             new_img = Image.new('RGB', (w*2, h*2))
  107.             new_img.paste(pil_imgs[0], (0, 0))       # 左上：新high
  108.             new_img.paste(pil_imgs[1], (w, 0))       # 右上：新wrist
  109.             new_img.paste(pil_imgs[2], (0, h))       # 左下：老high
  110.             new_img.paste(pil_imgs[3], (w, h))       # 右下：老wrist
  112.             debug_save_path = "debug_2x2.png"
  113.             new_img.save(debug_save_path)
  114.             print(f"Have been saved at: {debug_save_path}")
  115.             # new_img.show()
  116. ###################Debug图像########################
  117.     rdt_data = {
  118.         'message_id': _message_id,
  119.         'proprio': proprio,
  120.         'images': images,
  121.         'text_embeds': ""
  122.     }
  123.     client.send(rdt_data)
  124.     _message_id += 1
  125.     print(f"send new rdt data done, message_id: {_message_id-1}")
  126.     action_data = client.receive()
  127.     if action_data is None:
  128.         print("ERROR: Server returned None. Is the RDT server running?")
  129.         print("Please start the RDT server first!")
  130.         raise ConnectionError("Failed to receive response from RDT server")
  131.     actions = action_data['actions']
  132.     action_message_id = action_data["message_id"]
  133.     print(f"receive actions done, message_id: {action_message_id}")
  134.     # print(f"receive actions contents: {actions}")
  135.     actions_array = np.array(actions)
  136.     if len(actions_array.shape) == 3:
  137.         action_sequence = actions_array[0, :, :]  # 取第一个batch的所有时间步
  138.     else:
  139.         print(f"action shape should be 3 dim, but get {actions_array.shape} ")
  140.    
  141.     joint_names = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
  142.     for step_idx in range(0, len(action_sequence), 4):  # 64个动作隔4个执行一次动作
  143.         action_values = action_sequence[step_idx]
  144.         action_dict = {f"{joint}.pos": float(action_values[i]) for i, joint in enumerate(joint_names)}
  145.         sent_action = robot.send_action(action_dict)
  146.         time.sleep(0.1)  
  148.     _last_cam_high = cam_high
  149.     _last_cam_right_wrist = cam_right_wrist
  151. def main():
  152.     robot = SO101Follower(SO101FollowerConfig(
  153.         port="/dev/tty.usbmodem5AB90671801",
  154.         id="my_awesome_follower_arm",
  155.         cameras={
  156.             "arm": OpenCVCameraConfig(index_or_path=0, width=1920, height=1080, fps=30),
  157.             "high": OpenCVCameraConfig(index_or_path=2, width=1920, height=1080, fps=30)
  158.         }
  159.     ))
  160.    
  161.     robot.connect()
  162.     client = Client(host="localhost", port=5002)
  164.     try:
  165.         while True:
  166.             record_loop(
  167.                 robot,
  168.                 client
  169.             )
  170.             time.sleep(0.1)
  171.     except KeyboardInterrupt:
  172.         pass
  174.     robot.disconnect()
  176. if __name__ == "__main__":
  177.     main()
复制代码
        测试正常后我们便可以仿照上面服务器推理一样先启动板端推理代码，将其作为一个板端推理Server，接着运行我们本地的机械臂控制代码即可：
  1. # 服务端
  2. python3 RDT/scripts/lerobot_rdt_server.py
  3. # 客户端
  4. python3 lerobot/src/lerobot/record_rdt.py
复制代码
        RDKS100板端性能占用参考如下（多多支持我的Dtop:Jetson有Jtop,Linux有Htop,RDK也有Dtop！ - SkyXZ - 博客园哈哈哈哈哈哈哈哈哈
8.png

        RDKS100可直接连接LeRobot控制机械臂，但由于写文档的时候手上没有多余的摄像头，头部的摄像头只能使用Apple的连续互通实现，若使用RDK直接连接LeRobot的话将无法当问头部摄像头，因此只能绕个弯用本地Mac来控制机械臂了...

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！
不学晚了RDTLeRobotRDKS100

使用道具 举报
煅汾付
15 小时前

举报

谢谢分享，辛苦了
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册
发帖
雌鲳签
15 小时前

0

粉丝关注

17

主题发布

板块介绍填写区域，请于后台编辑
微信扫一扫