Fourier-Lerobot——把斯坦福人形动作策略iDP3封装进了Lerobot(含我司七月的idp3落地实践)

前言

近期在抠lerobot源码时,看到其封装了ALOHA ACT、diffusion policy、π0时,我就在想,lerobot其实可以再封装下idp3

  1. 我甚至考虑是否从我联合带的那十几个具身研究生中选几个同学做下这事,对他们也是很好的历练
  2. 然当25年3.18日晚上,我把lerobot抠的差不多了「比如此文《LeRobot源码剖析——对机器人各个动作策略的统一封装:包含ALOHA ACT、Diffusion Policy、VLA模型π0」之后
    却发现傅利叶fork了lerobot,并在fork的fourier-lerobot中「GitHub地址详见:FFTAI/fourier-lerobot」,把idp3封装了进去,实在是卷啊..

再加之工厂机械臂开发订单之外,我司「七月在线」近期接到的B端人形开发订单越来越多了(且还有多个人形开发需求正在并行推进中)

而对于其中一个人形开发订单,我(们)准备把ipd3作为备选,既然fourier把ipd3封装进了lerobot,那这个fourier-lerobot便是我们在落地中会尝试的库之一

  1. 当然了,idp3外,像vla π0 我们也会考虑并行尝试,而π0此前已经封装进了lerobot,故π0官方库、lerobot库都是可以选择的 
  2. 顺带,我们非常缺人,如果有志于做具身智能或人形研发的,欢迎私我(目前在长沙、南京的具身团队之外,今25年q2正在上海再建立一支具身全职团队),共同让机器人干好活,推动新一轮工业革命的更快发展

综上,故写本文,而FFTAI/fourier-lerobot对huggingface/lerobot的主要扩展集中在以下几个方面:

  1. 数据集支持:添加了对Fourier ActionNet数据集的全面支持,包括转换工具和可视化工具
  2. 训练管道:扩展了IDP3训练管道,针对人形机器人优化
  3. 工具扩展:提供了特定于Fourier数据集的工具链

这些修改使FFTAI/fourier-lerobot成为一个专为人形机器人开发优化的版本,而保持了与原始huggingface/lerobot项目的核心框架兼容性

第一部分 fourier-lerobot新增的独立scripts

FFTAI/fourier-lerobot添加了一个完全独立的scripts目录(与lerobot/scripts不同),包含:

  • convert_to_lerobot_v2.py:这是一个专门为Fourier数据集开发的转换工具,用于将Fourier ActionNet数据集转换为LeRobotDatasetV2格式
    该脚本处理了:
    HDF5格式的机器人状态和动作数据
    摄像机RGB和深度视频数据
    点云数据生成和处理
    时间戳对齐
    特定于人形机器人的关节数据处理
  • fourier_viz.py:用于可视化Fourier数据集的专用工具
  • requirements.txt:Fourier数据集处理工具的依赖项

1.1 convert_to_lerobot_v2.py:机器人数据集转换脚本

这是一个用于将原始机器人数据转换为标准化 LeRobot 数据集格式的工具脚本。该脚本主要处理多模态数据,包括视频图像、深度图像、点云和机器人状态/动作数据,并将它们整合为时间同步的数据集

该脚本首先是导入模块,包括一些常用工具、外部依赖库以及内部项目的导入

其次是一系列工具函数

扫描二维码关注公众号,回复: 17551009 查看本文章
  1. 点云处理函数
    比如点云网格采样
    grid_sample_pcd(point_cloud, grid_size=0.005
    再比如从深度图创建点云
    create_colored_point_cloud_from_depth_oak(depth, far=1.0, near=0.1, num_points=10000)
  2. 时间同步函数
    比如匹配两个时间序列
    match_timestamps(candidate, ref)
  3. 辅助函数
    `get_cameras(hdf5_path)`:获取数据集中的相机列表
    `iso_to_datetime(filename: str) -> datetime`:ISO格式文件名转时间对象`filter_state_pose(state)`:从姿态状态筛选需要的部分

然后还包括一系列常量定义,比如机器人关节名称、收不关节名称、姿态参数名称等等

以及数据结构定义函数,比如下面这个定义数据集特征结构的函数

make_features(state_names, action_names, mode="video")

最后是主处理main函数,其工作流程是

  1. 参数与初始化
    处理输入参数,创建数据集对象,根据使用模式选择关节或姿态作为状态表示
  2. 数据读取
    遍历HDF5文件,读取机器人状态、动作数据,获取相机视频/图像数据和时间戳
  3. 数据同步
    匹配不同模态数据(状态、图像、点云)的时间戳,处理重复和冲突时间戳
  4. 点云生成
    包括从深度视频中读取每一帧、将深度图转换为3D点云、对点云进行降采样和归一化处理
  5. 数据整合与保存
    包括将各模态数据整合为一致的数据帧、按照定义的数据结构保存为标准化数据集、可选择推送到远程仓库

1.1.1 点云处理函数:grid_sample_pcd、create_colored_point_cloud_from_depth_oak

1.1.2 时间同步:match_timestamps

1.1.3 一系列辅助函数:get_cameras、iso_to_datetime、filter_state_pose

1.1.4 一系列常量定义和数据结构的定义

1.1.5 main函数:整体数据处理流程

`main`函数是一个复杂的数据处理流水线,用于将原始机器人数据转换为标准化的LeRobot数据集格式。该函数处理多模态数据,包括视频、图像、深度信息、点云和机器人状态/动作数据

首先,函数对输入参数进行验证,确保数据模式(`mode`)为"video"或"image",并清理任何现有的同名数据集
根据`use_pose`参数决定是使用关节数据还是姿态数据,这影响到状态和动作的表示方式

    # 如果使用姿态表示
    if use_pose:  
        # 使用过滤后的姿态状态名称和手部关节名称
        state_names = filter_state_pose(STATE_POSE_NAMES) + STATE_HAND6DOF_NAMES 

        # 使用姿态动作名称和手部关节名称 
        action_names = ACTION_POSE_NAMES + ACTION_HAND6DOF_NAMES 

通过`make_features`函数定义数据集的特征规格

    // # 创建特征字典,定义数据结构
    FOURIER_FEATURES = make_features(state_names, action_names, mode=mode)  

然后使用`LeRobotDataset.create`创建一个新的数据集实例

    // # 创建LeRobotDataset实例
    dataset = LeRobotDataset.create(  
        repo_id=repo_id,                # 设置仓库ID
        fps=30,                         # 设置帧率为30fps
        robot_type="gr1t1",             # 设置机器人类型
        features=FOURIER_FEATURES,      # 设置特征字典
        # tolerance_s=0.1,              # 注释掉的容差设置
        image_writer_threads=4,         # 设置图像写入线程数
        image_writer_processes=4,       # 设置图像写入进程数
    )

对于每个HDF5文件(代表一个机器人操作片段),函数提取状态和动作数据

    # 获取并排序所有的HDF5文件
    hdf5_files = sorted(raw_dir.glob("*.hdf5"))  
    num_episodes = len(hdf5_files)      # 计算片段数量
    episodes = range(num_episodes)      # 创建片段索引范围

然后便是通过循环处理每个摄像头,函数处理视频/图像数据和深度数据

以下是主要功能:

数据加载和预处理
检查是否有任务描述,如果没有则从`metadata.json`加载

   # 如果未指定任务 
   if task is None:  
        # 从元数据文件加载任务信息
        metadata = json.load(open(raw_dir / "metadata.json"))  
    
        # 将元数据转换为以ID为键的字典
        metadata = {m["id"]: m for m in metadata}

遍历每个数据集片段(episode),这些片段以HDF5文件形式存储

    # 遍历每个片段
    for ep_idx in episodes:  
        # 获取当前片段的HDF5文件路径
        ep_path = hdf5_files[ep_idx]

根据`use_pose`参数决定使用姿态数据还是关节数据:

接下来,便是从HDF5文件中提取和构建机器人状态和动作数据的关键过程

  1. 首先,代码使用`h5py.File`以只读模式打开机器人示范片段文件,然后基于`use_pose`参数选择不同的数据提取策略
    如果启用了姿态模式(`use_pose=True`)——姿态数据包括:姿态状态/state/pose、手部状态/state/hand
    具体而言,代码从`/state/pose`路径提取前18列和最后6列数据,并与手部状态数据(`/state/hand`)连接,形成完整的状态向量
            // # 以只读模式打开HDF5文件
            with h5py.File(ep_path, "r") as ep:  
                # 如果使用姿态表示
                if use_pose:  
    
                    # 将NumPy数组转换为PyTorch张量
                    state = torch.from_numpy(  
                        np.concatenate(  
    
                            # 提取姿态状态/state/pose的前18列和最后6列数据,然后再连接手部状态——从而组合姿态状态state/pose和手部状态state/hand
                            [ep["/state/pose"][:, :18], ep["/state/pose"][:, -6:], ep["/state/hand"][:]], axis=1
                        )
                    ) 
    相应地,动作数据通过连接姿态动作(`/action/pose`)和手部动作(`/action/hand`)构建,这种方式主要关注机器人的空间姿态表示
                    action = torch.from_numpy(  
                        # 连接姿态动作action/pose、和手部动作action/hand
                        np.concatenate([ep["/action/pose"][:], ep["/action/hand"][:]], axis=1)  
                    ) 
  2. 若不使用姿态模式,则采用另一策略:状态数据由机器人关节角度(`/state/robot`第12列开始)和手部状态构成——即组合`/state/robot`(从第12个关节开始)和`/state/hand`数据
                # 如果使用关节表示
                else:  
                    state = torch.from_numpy(  
                        # 连接机器人关节状态(从第12个关节开始)和手部状态
                        np.concatenate([ep["/state/robot"][:, 12:], ep["/state/hand"][:]], axis=1)  
                    )  
    动作数据同样由相应的关节动作和手部动作组合而成
                    # 连接末端执行器姿态状态和手部状态
                    action = torch.from_numpy(  
                        # 连接机器人关节动作(从第12个关节开始)和手部动作
                        np.concatenate([ep["/action/robot"][:, 12:], ep["/action/hand"][:]], axis=1)  
                    ) 
  3. 若不使用姿态模式的此模式下,代码还会额外构建
    `pose_state`
                    pose_state = torch.from_numpy(  
                        # 连接姿态状态的不同部分
                        np.concatenate(  
                            [ep["/state/pose"][:, :18], ep["/state/pose"][:, -6:], ep["/state/hand"][:]], axis=1
                        )
                    )
    和`pose_action`作为备选表示
                    # 将NumPy数组转换为PyTorch张量
                    # 连接姿态动作和手部动作
                    pose_action = torch.from_numpy(  
                        np.concatenate([ep["/action/pose"][:], ep["/action/hand"][:]], axis=1)  
                    ) 
    所有数据在提取后通过`torch.from_numpy()`转换为PyTorch张量,为后续深度学习模型训练做准备。这种数据结构设计确保了机器人的全身状态和动作能够被完整捕获,无论是以关节空间还是笛卡尔空间表示
为方便大家更好的理解,我再用下面表格总结一下
如果启用了姿态模式
状态数据 姿态状态/state/pose 手部状态/state/hand
动作数据 姿态动作(`/action/pose`) 手部动作(`/action/hand`)
若不使用姿态模式
状态数据 关节状态`/state/robot`(从第12个关节开始) 手部状态`/state/hand`
动作数据 关节动作/action/robot 手部动作/action/hand
                备选表示——相当于作为不使用姿态模式下的备选

pose_state

ep["/state/pose"][:, :18], ep["/state/pose"][:, -6:]

/state/pose /state/hand

pose_action

ep["/action/pose"][:], ep["/action/hand"

/action/pose /action/hand

接下来,是时间戳同步

  1. 获取每个摄像头的数据目录
                num_frames = None       # 初始化帧数为None
                matched = None          # 初始化匹配索引为None
    
                # 创建帧字典,用于存储所有模态数据
                frames = {}  
    
                # 遍历每个摄像头
                for camera in get_cameras(ep_path.parent / ep_path.stem):  
                    vid_key = f"observation.images.{camera}"               # 创建视频键名
                    depth_key = f"observation.depth.{camera}"              # 创建深度键名
                    pc_key = "observation.pointcloud"                      # 创建点云键名
    
                    raw_vid_dir = ep_path.parent / ep_path.stem / camera   # 获取原始视频目录
  2. 加载视频/图像时间戳
  3. 使用`match_timestamps`函数将机器人状态时间戳与图像时间戳对齐
                    # 匹配时间戳,获取对应索引
                    matched = match_timestamps(data_ts[non_duplicate], image_ts)  
    
                    # 裁剪匹配索引,去除开头和结尾的指定帧数
                    matched = matched[  
                        # 移除前后指定帧数
                        discard_frames : -discard_frames if discard_frames else None
                    ]  
    
                    # 计算帧数
                    num_frames = len(matched)  
  4. 可以选择丢弃开始和结束的几帧(通过`discard_frames`参数)

再其次,是多模态数据处理

  • 处理RGB视频或图像:
    \rightarrow  视频模式:直接复制视频文件,添加帧引用到数据集
    \rightarrow  图像模式:复制单独的PNG文件,构建帧序列
  • 处理深度数据:
    \rightarrow  复制原始深度视频
    \rightarrow  如果启用点云功能:
               - 从深度视频读取每一帧
               - 使用`create_colored_point_cloud_from_depth_oak`函数将每帧深度转换为3D点云
                            points = []          # 创建点云列表
                            # 注释掉的深度列表
    
                            while True:          # 循环读取视频帧
                                # 读取一帧
                                ret, frame = video_capture.read()  
                
                                # 如果读取成功
                                if ret:  
                                    # 注释掉的代码,用于保存深度数据
                                    points.append(      # 添加点云
                                        # 从深度图创建点云
                                        create_colored_point_cloud_from_depth_oak(frame * 1e-3, num_points=4096)  
                                    )
               - 生成标准化的点云数据(4096个点)

再其次,数据集构建和保存

  1. 设置最后一步为结束标志(`done`)
  2. 将状态和动作数据与时间戳匹配
  3. 按帧构建完整的多模态数据字典
  4. 逐帧添加到数据集
  5. 保存每个episode并进行垃圾回收
  6. 数据集整合和可选的上传到模型库

这个处理循环确保了视觉数据、状态数据和动作数据在时间上完美同步,生成可用于机器人学习的高质量多模态数据集

最后,函数保存每个片段,标记最后一帧为结束帧,并在所有片段处理完成后合并数据集。如果启用`push_to_hub`选项,数据集会被上传到模型库(可能是HuggingFace)供共享使用
整个过程构建了一个多模态、时间同步的机器人操作数据集,适用于机器学习和机器人学研究

1.2 fourier_viz.py

// 待更

第二部分 全新的idp3训练管道实现

原始huggingface/lerobot项目中不包含IDP3训练管道,而Fourier-Lerobot新增的lerobot/common/policies/idp3目录之下,包含以下文件:

  1. configuration_idp3.py
  2. modeling_idp3.py
  3. pointnet_extractor.py

2.1 configuration_idp3.py

2.2 modeling_idp3.py:从点云和状态数据生成机器人控制动作

该代码实现了一个名为"改进型3D扩散策略"(IDP3)的系统,用于从点云和状态数据生成机器人控制动作。以下是各组件间的调用关系分析:

整个系统由6个主要类构成,形成清晰的层次结构:

IDP3Policy (策略层)
    ↓ 包含
IDP3Model (模型层)
    ↓ 包含
├── IDP3Encoder (观测编码器)
├── DiffusionConditionalUnet1d (U-Net架构)
│   ↓ 包含
│   └── DiffusionConditionalResidualBlock1d (残差块)
└── DDPMScheduler/DDIMScheduler (噪声调度器)

以下是一些关键的调用路径

策略初始化路径
IDP3Policy.__init__
  ├── 创建数据归一化器(Normalize, Unnormalize)
  └── 创建IDP3Model实例(self.diffusion)
      ├── 创建IDP3Encoder(obs_encoder)
      ├── 创建DiffusionConditionalUnet1d(self.unet)
      │   └── 创建多个DiffusionConditionalResidualBlock1d
      └── 调用_make_noise_scheduler创建噪声调度器(self.noise_scheduler)
动作选择路径(推理时)
IDP3Policy.select_action
  ├── 调用normalize_inputs归一化输入
  ├── 调用populate_queues填充观测队列
  ├── [队列为空时]调用IDP3Model.generate_actions
  │   ├── 重组数据形状
  │   ├── 调用obs_encoder编码观测
  │   ├── 调用conditional_sample生成动作
  │   │   ├── 创建随机噪声作为起点
  │   │   └── 循环执行去噪步骤:
  │   │       ├── 调用unet进行预测
  │   │       └── 调用noise_scheduler.step执行去噪
  │   └── 提取相关时间段的动作
  └── 从动作队列中弹出并返回当前动作
训练路径
IDP3Policy.forward
  ├── 调用normalize_inputs归一化输入
  ├── 调用normalize_targets归一化目标
  └── 调用IDP3Model.compute_loss
      ├── 重组数据形状
      ├── 调用obs_encoder编码观测
      ├── 添加随机噪声到轨迹
      ├── 调用unet进行预测
      └── 计算预测与目标之间的损失
U-Net前向传播路径
DiffusionConditionalUnet1d.forward
  ├── 重排输入张量维度
  ├── 编码时间步
  ├── 合并条件信息
  ├── 执行编码器路径(每个下采样模块)
  │   └── 多次调用DiffusionConditionalResidualBlock1d.forward
  ├── 执行中间处理
  ├── 执行解码器路径(每个上采样模块)
  │   └── 多次调用DiffusionConditionalResidualBlock1d.forward
  └── 通过最终卷积层

整个系统中的数据流向遵循以下模式:

  1. 输入处理:点云和状态数据 → 归一化 → 观测队列
  2. 特征提取:观测数据 → IDP3Encoder → 全局条件特征
  3. 动作生成:噪声 → 迭代去噪 → 原始动作 → 逆归一化 → 最终动作
  4. 训练反馈:动作 → 添加噪声 → 预测去噪 → 计算损失 → 反向传播

2.2.1 IDP3Policy类:包含IDP3Model,涉及初始化、select_action、forward

`IDP3Policy`类:主策略类,负责:输入/输出数据归一化处理、管理观测和动作的历史队列、使用扩散模型选择动作

在初始化过程中,配置了三个关键组件:

  1. 输入/输出归一化器:确保所有数据按统一标准处理
  2. 观测/动作队列:维护历史状态和预生成的动作
  3. 扩散模型:核心动作生成引擎

特别值得注意的是,该策略支持多模态输入,可以处理状态数据、点云,以及配置允许时的图像数据和环境状态

`select_action`方法展示了该策略的精妙之处,它实现了一个优雅的滑动窗口机制:

  1. 维护最新的`n_obs_steps`个观测历史
    系统会使用过去的`n_obs_steps`个观测,生成长度为`horizon`的动作序列,但只实际执行其中的`n_action_steps`个动作
                (legend: o = n_obs_steps, h = horizon, a = n_action_steps)  # 图例:o = n_obs_steps, h = horizon, a = n_action_steps
                |timestep            | n-o+1 | n-o+2 | ..... | n     | ..... | n+a-1 | n+a   | ..... | n-o+h |  # 时间步
                |observation is used | YES   | YES   | YES   | YES   | NO    | NO    | NO    | NO    | NO    |  # 观测是否使用
                |action is generated | YES   | YES   | YES   | YES   | YES   | YES   | YES   | YES   | YES   |  # 动作是否生成
                |action is used      | NO    | NO    | NO    | YES   | YES   | YES   | NO    | NO    | NO    |  # 动作是否使用
  2. 当动作队列为空时,使用历史观测生成未来`horizon`步骤的动作轨迹
  3. 只使用其中的`n_action_steps`步骤作为实际执行的动作

该类同时支持推理(动作生成)和训练功能,通过`forward`方法计算训练或验证损失。在训练过程中:

  1. 对输入和目标进行归一化处理
        # 前向传播方法
        def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:  
            # 通过模型运行批次并计算训练或验证的损失
            batch = self.normalize_inputs(batch)       # 对输入批次进行归一化
            batch = self.normalize_targets(batch)      # 对目标批次进行归一化
  2. 将处理后的数据传递给扩散模型的`compute_loss`方法——此compute_loss下一节会详述
            # 使用扩散模型计算损失
            loss = self.diffusion.compute_loss(batch)  
  3. 返回平均损失用于反向传播
            return loss, None      # 返回损失和None

这种设计实现了训练和推理代码的统一,使模型能够无缝地从训练过渡到部署

此外,还涉及对噪声调度器的选择

# 创建噪声调度器函数
def _make_noise_scheduler(name: str, **kwargs: dict) -> DDPMScheduler | DDIMScheduler:  
    """
    # 请求类型的噪声调度器实例的工厂函数
    to the scheduler.       # 所有kwargs都传递给调度器
    """
    if name == "DDPM":      # 如果名称是DDPM
        return DDPMScheduler(**kwargs)      # 返回DDPM调度器实例
    elif name == "DDIM":    # 如果名称是DDIM
        return DDIMScheduler(**kwargs)      # 返回DDIM调度器实例
    else:  
        # 抛出不支持的噪声调度器类型错误
        raise ValueError(f"Unsupported noise scheduler type {name}")  

2.2.2 IDP3Model类:初始化、conditional_sample、generate_actions、compute_loss

`IDP3Model`类:核心扩散模型,实现:通过PointNet编码器处理点云和状态数据、使用扩散过程生成动作轨迹、同时支持训练(计算损失)和推理(生成动作)

2.2.2.1 初始化方法:obs_encoder、条件U-Net模型、_make_noise_scheduler

`IDP3Model`类的初始化方法构建了一个先进的3D扩散策略模型,该模型能够处理点云数据并生成连贯的机器人动作序列。这个初始化过程设计精巧,构建了三个核心组件,共同形成了一个强大的动作生成系统。

  1. 首先,方法创建了一个观测编码器(`obs_encoder`),这是一个`IDP3Encoder`实例,配置为使用多阶段PointNet架构处理点云数据
    # 定义IDP3模型类,继承自nn.Module
    class IDP3Model(nn.Module):  
        def __init__(self, config: IDP3Config):  # 初始化方法
            super().__init__()                   # 调用父类初始化方法
            self.config = config                 # 保存配置
    
            # 构建观测编码器(取决于提供哪些观测)
            obs_encoder = IDP3Encoder(                  # 创建观测编码器
                observation_space=config.obs_dict,      # 观测空间
                pointcloud_encoder_cfg=config.pointcloud_encoder_cfg,  # 点云编码器配置
                use_pc_color=False,                        # 不使用点云颜色
                pointnet_type="multi_stage_pointnet",      # 使用多阶段PointNet
                point_downsample=False,                    # 不对点进行下采样
            )
    这种设计选择非常关键,因为多阶段PointNet能够有效提取点云的层次化特征——将原始3D点转换为语义特征,而不需要依赖点云颜色信息(`use_pc_color=False`)或进行点云下采样(`point_downsample=False`)
    编码器的输出维度被用来计算全局条件向量的维度,这个向量将随后用于引导扩散过程
            obs_feature_dim = obs_encoder.output_shape()    # 获取观测特征维度
            global_cond_dim = obs_feature_dim * config.n_obs_steps      # 计算全局条件维度
            self.obs_encoder = obs_encoder                  # 保存观测编码器
  2. 接下来,代码实例化了一个条件U-Net模型(`self.unet`),这是一个一维卷积结构,专门设计用于处理时序动作数据
            # 创建条件U-Net
            self.unet = DiffusionConditionalUnet1d(config, global_cond_dim=global_cond_dim)  
    该U-Net接收通过观测编码器生成的全局条件,使用FiLM(特征线性调制)技术将条件信息整合到网络的各个层级中
  3. 最后,通过`_make_noise_scheduler`工厂函数创建了噪声调度器,这是扩散过程的核心控制器
            # 创建噪声调度器
            self.noise_scheduler = _make_noise_scheduler(  
                config.noise_scheduler_type,                     # 噪声调度器类型
                num_train_timesteps=config.num_train_timesteps,  # 训练时间步数
                beta_start=config.beta_start,      # beta起始值
                beta_end=config.beta_end,          # beta结束值
                beta_schedule=config.beta_schedule,      # beta调度
                clip_sample=config.clip_sample,          # 是否裁剪样本
                clip_sample_range=config.clip_sample_range,      # 样本裁剪范围
                prediction_type=config.prediction_type,          # 预测类型
            )
    代码支持两种扩散算法:DDPM(去噪扩散概率模型)和DDIM(去噪扩散隐式模型),并根据配置参数控制噪声添加和移除的速率。这个调度器在训练时控制向干净轨迹添加的噪声量,在推理时则控制去噪步骤的执行方式
2.2.2.2 conditional_sample:利用U-Net,从纯随机噪声开始 多步迭代去噪

`conditional_sample`方法实现了扩散模型的推理过程:从纯随机噪声开始,通过多步迭代去噪,逐步形成有意义的动作轨迹。每一步去噪过程都利用U-Net网络预测,并受全局条件(编码后的观测)引导

  1. 首先,方法获取模型的设备和数据类型信息,确保所有张量操作在一致的计算环境中进行。
        # inference 推理部分,下面为条件采样方法
        def conditional_sample( 
            # 批次大小,全局条件,随机数生成器
            self, batch_size: int, global_cond: Tensor | None = None, generator: torch.Generator | None = None  
        ) -> Tensor:                # 返回张量
            device = get_device_from_parameters(self)      # 从参数获取设备
            dtype = get_dtype_from_parameters(self)        # 从参数获取数据类型
    然后通过`torch.randn`函数创建初始噪声样本,这个随机噪声tensor的形状为(batch_size, horizon, action_dim),表示一批完全随机的动作轨迹
            # Sample prior —— 采样先验
            # 生成随机正态分布样本
            sample = torch.randn(  
                # 大小为(批次大小, 视野, 动作特征维度)
                size=(batch_size, self.config.horizon, self.config.action_feature.shape[0]),  
                dtype=dtype,              # 使用获取的数据类型
                device=device,            # 使用获取的设备
                generator=generator,      # 使用提供的随机数生成器
            )
    这一步是扩散模型生成过程的起点,类似于从一片混沌中逐渐雕刻出有意义的结构
  2. 接下来,方法配置噪声调度器的时间步数,这决定了去噪过程的精细程度
            # 设置噪声调度器的时间步数
            self.noise_scheduler.set_timesteps(self.num_inference_steps)
    调度器通常按照从高噪声到低噪声的顺序设置时间步,形成一个渐进去噪的路径
  3. 随后进入扩散模型的核心循环,对每个时间步:
    1. 调用U-Net模型预测去噪方向或目标,此处U-Net接收当前噪声样本、时间步信息(通过`torch.full`创建的均匀时间步张量)以及全局条件(观测编码)
            for t in self.noise_scheduler.timesteps:  # 对于每个时间步
                # 预测模型输出
                model_output = self.unet(      # 调用U-Net
                    sample,                    # 当前样本
    
                    # 创建全为t的张量作为时间步
                    torch.full(sample.shape[:1], t, dtype=torch.long, device=sample.device),  
                    # 全局条件
                    global_cond=global_cond,  
                )
    2. 使用噪声调度器的`step`方法执行一步去噪,将样本从状态 x_t 转换为略微清晰的状态x_{t-1}
                # 计算前一图像:x_t -> x_t-1,且使用调度器进行去噪步骤
                sample = self.noise_scheduler.step(model_output, t, sample, generator=generator).prev_sample  
    
            # 返回最终样本
            return sample  

经过所有时间步的迭代,初始的随机噪声被逐渐转换为有意义的动作轨迹

2.2.2.3 generate_actions:使用条件采样生成动作序列

当调用`generate_actions`时,模型处理机器人状态和点云数据,重组数据形状以适应编码器,然后使用条件采样生成动作序列。最后,从生成的序列中精确提取当前需要的动作部分

首先,是输入数据处理

  • 函数从批次数据中提取基本维度信息(批次大小和观测步数),并通过断言确保观测历史长度符合模型配置要求。这是保证时序数据完整性的重要前提
        # 生成动作方法
        def generate_actions(self, batch: dict[str, Tensor]) -> Tensor:  
            """
            # 该函数期望batch具有:
            {
                # (批次大小, 观测步数, 状态维度)
                "observation.state": (B, n_obs_steps, state_dim)  
    
                 # (批次大小, 观测步数, 相机数量, 通道数, 高度, 宽度)
                "observation.images": (B, n_obs_steps, num_cameras, C, H, W) 
                    AND/OR  # 和/或
    
                 # (批次大小, 环境维度)
                "observation.environment_state": (B, environment_dim) 
    
                # (批次大小, 观测步数, 点数量 * 3)
                "observation.pointcloud": (B, n_obs_steps, num_points * 3)  
                        }
            """
            # 获取批次大小和观测步数
            batch_size, n_obs_steps = batch["observation.state"].shape[:2]  
    
            # 断言观测步数等于配置的观测步数
            assert n_obs_steps == self.config.n_obs_steps

接下来,函数对点云数据进行两次关键的重塑操作:

  1. 第一次重塑将扁平的点云数据`(B, n_obs_steps, num_points*3)`转换为带有明确几何结构的形式`(B, n_obs_steps, num_points, 3)`,每个点有x、y、z三个坐标
            # 使用多阶段PointNet编码器编码点云特征并将它们与状态向量连接在一起
            # 重塑点云形状
            batch["observation.pointcloud"] = batch["observation.pointcloud"].reshape(  
                # (批次大小, 观测步数, 点数量, 3)
                batch["observation.pointcloud"].shape[0], batch["observation.pointcloud"].shape[1], -1, 3  
            )
  2. 第二次重塑将批次维度和时间步维度合并,形成`(-1, num_points, 3)`的形状,这使得多阶段PointNet编码器可以批量处理所有时间步的点云
            # 再次重塑点云形状
            batch["observation.pointcloud"] = batch["observation.pointcloud"].reshape( 
                # (批次大小 * 观测步数, 点数量, 3) 
                -1, *batch["observation.pointcloud"].shape[2:]  
            )

状态数据也进行了类似的维度合并处理,确保与点云数据格式一致,为后续特征提取做准备

         # 重塑状态形状
        batch["observation.state"] = batch["observation.state"].reshape( 
             # (批次大小 * 观测步数, 状态维度)
            -1, *batch["observation.state"].shape[2:] 
        )

其次,是特征编码与条件生成

  1. 经过重塑的观测数据被送入观测编码器(多阶段PointNet),生成能够捕捉环境几何特性的全局条件向量
            # 使用观测编码器获取全局条件
            global_cond = self.obs_encoder(batch)  # (B, global_cond_dim)  
    这个向量随后被重新整形回`(batch_size, -1)`,保留了批次维度但合并了所有特征维度,为扩散模型提供条件引导信息
            # 重塑全局条件形状为(批次大小, 特征维度)
            global_cond = global_cond.reshape(batch_size, -1)  
  2. 调用`conditional_sample`方法是整个过程的核心,它执行条件扩散采样,从随机噪声开始,在全局条件的引导下逐步去噪,最终生成完整的动作轨迹
            # 运行采样,使用条件采样生成动作
            actions = self.conditional_sample(batch_size, global_cond=global_cond) 
    这个过程实现了从纯噪声到有意义动作序列的结构化生成

最后,动作提取与返回

最后,函数从生成的完整轨迹中精确提取所需的时间片段:

  • 起始索引`start = n_obs_steps - 1`对应当前时刻(最新的观测)
            # 从当前观测提取n_action_steps步的动作
            start = n_obs_steps - 1          # 起始索引为观测步数减1
  • 结束索引`end = start + self.config.n_action_steps`确保提取配置指定数量的连续动作
            end = start + self.config.n_action_steps     # 结束索引为起始索引加动作步数
            actions = actions[:, start:end]              # 提取指定范围的动作
    最后返回动作
            return actions  # 返回动作

这种提取方式确保返回的动作序列紧接当前观测,为机器人提供即时且连贯的控制指令

2.2.2.4 compute_loss:使用U-Net预测噪声

`compute_loss`方法实现了模型的训练机制:

  1. 对干净的动作轨迹添加随机噪声,噪声程度由随机采样的时间步决定
  2. 使用U-Net预测噪声(epsilon模式)或原始轨迹(sample模式)
  3. 计算预测与目标之间的均方误差
  4. 可选地通过掩码机制忽略填充部分的误差 

具体而言,分为4个阶段

第一,数据准备与验证

首先,方法严格验证输入数据的完整性,确保包含必要的状态、动作、填充标记和点云数据。这种防御性编程确保了训练过程的稳定性。接着,将点云数据从扁平形式`(B, n_obs_steps, num_points*3)`重塑为结构化的点坐标表示`(B, n_obs_steps, num_points, 3)`,便于后续处理

    def compute_loss(self, batch: dict[str, Tensor]) -> Tensor:
        """
        This function expects `batch` to have (at least):
        {
            "observation.state": (B, n_obs_steps, state_dim)

            "observation.images": (B, n_obs_steps, num_cameras, C, H, W)
                AND/OR
            "observation.environment_state": (B, environment_dim)

            "observation.pointcloud": (B, n_obs_steps, num_points * 3)

            "action": (B, horizon, action_dim)
            "action_is_pad": (B, horizon)
        }
        """
        # 输入验证:确保批次数据包含所需的关键字段
        assert set(batch).issuperset({"observation.state", "action", "action_is_pad"})

        # 确保点云数据存在
        assert "observation.pointcloud" in batch

        # 重塑点云数据为标准格式:(batch_size, n_obs_steps, num_points, 3)
        batch["observation.pointcloud"] = batch["observation.pointcloud"].reshape(
            batch["observation.pointcloud"].shape[0], batch["observation.pointcloud"].shape[1], -1, 3
        )

方法还提取并验证关键维度信息(批次大小、观测步数、动作范围),确保与模型配置一致。这些检查对于防止形状不匹配错误至关重要

        # 获取观察步数、动作步数和批次大小
        n_obs_steps = batch["observation.state"].shape[1]
        horizon = batch["action"].shape[1]
        batch_size = batch["observation.state"].shape[0]

        # 确保动作步数和观察步数与配置一致
        assert horizon == self.config.horizon
        assert n_obs_steps == self.config.n_obs_steps

第二,观测编码与扩散过程

经过初始重塑后,方法进一步合并批次和时间维度,创建适合PointNet编码器处理的格式。编码器将点云和状态数据转换为全局条件向量,表示环境的关键特征

        # 重塑点云数据以适应编码器输入:将所有观察步骤展平
        batch["observation.pointcloud"] = batch["observation.pointcloud"].reshape(
            -1, *batch["observation.pointcloud"].shape[2:]
        )

        # 同样重塑状态数据
        batch["observation.state"] = batch["observation.state"].reshape(
            -1, *batch["observation.state"].shape[2:]
        )

        # 使用观察编码器获取全局条件特征
        global_cond = self.obs_encoder(batch)  # (B, global_cond_dim)

        # 重塑全局条件特征以匹配批次大小
        global_cond = global_cond.reshape(batch_size, -1)

扩散过程的核心步骤如下:

  1. 提取干净的动作轨迹
            # 前向扩散过程
            trajectory = batch["action"]  # 获取干净的轨迹数据
  2. 生成与轨迹同形状的随机噪声(`eps`)
            # 采样噪声以添加到轨迹中
            eps = torch.randn(trajectory.shape, device=trajectory.device)
  3. 为每个批次样本随机采样时间步,确定添加噪声的程度
            # 为批次中的每个样本随机采样一个噪声时间步
            timesteps = torch.randint(
                low=0,
                high=self.noise_scheduler.config.num_train_timesteps,
                size=(trajectory.shape[0],),
                device=trajectory.device,
            ).long()
  4. 使用噪声调度器添加噪声,创建`noisy_trajectory`
            # 根据时间步的噪声幅度将噪声添加到干净轨迹中
            noisy_trajectory = self.noise_scheduler.add_noise(trajectory, eps, timesteps)
  5. 将噪声轨迹输入U-Net模型,得到预测结果
            # 运行去噪网络(可以是预测去噪后的轨迹,也可以是预测噪声)
            pred = self.unet(noisy_trajectory, timesteps, global_cond=global_cond)

第三,损失计算与掩码处理

损失计算展现了扩散模型的灵活性,支持两种预测模式:

  1. epsilon模式:模型预测添加的噪声,损失基于预测值与原始噪声的差异
            # 计算损失
            # 目标可以是原始轨迹或噪声,取决于预测类型
            if self.config.prediction_type == "epsilon":
                target = eps  # 目标是预测噪声
  2. sample模式:模型直接预测原始轨迹,损失基于预测值与干净动作的差异
            elif self.config.prediction_type == "sample":
                target = batch["action"]      # 目标是预测原始轨迹
            else:
                raise ValueError(f"Unsupported prediction type {self.config.prediction_type}")
    
            # 计算预测值与目标值之间的均方误差
            loss = F.mse_loss(pred, target, reduction="none")

第四,方法通过掩码机制优化训练过程,排除填充动作对损失的影响。这确保了模型只从真实动作数据中学习,而不会被人工填充的数据误导

        # 掩盖那些被填充的动作(数据集轨迹边缘的复制)的损失
        if self.config.do_mask_loss_for_padding:
            # 确保batch中包含action_is_pad字段
            if "action_is_pad" not in batch:
                raise ValueError(
                    "You need to provide 'action_is_pad' in the batch when "
                    f"{self.config.do_mask_loss_for_padding=}."
                )

            # 创建掩码:只保留非填充动作的损失
            in_episode_bound = ~batch["action_is_pad"]

            # 应用掩码到损失
            loss = loss * in_episode_bound.unsqueeze(-1)

        # 返回平均损失
        return loss.mean()

这种设计使IDP3能够从混沌中学习秩序,掌握从噪声中恢复干净动作轨迹的能力,为推理时的动作生成奠定基础

2.2.3 条件U-Net架构:`DiffusionConditionalUnet1d`类

DiffusionConditionalUnet1d类构建了一个一维卷积U-Net网络,专门用于处理时序动作数据,其特点是:下采样/上采样路径的跳跃连接、FiLM(特征线性调制)条件机制、残差块设计增强梯度流

// 待更

最后,总计一下,对于iDP3

  • 训练阶段:向干净的动作轨迹添加噪声、训练网络预测原始噪声(epsilon预测)或原始动作(直接预测)、使用MSE损失优化网络参数
  • 推理阶段:从纯噪声开始、通过迭代去噪步骤(由`noise_scheduler`控制)、逐步将噪声转换为有效的动作轨迹、最终抽取适当时间段的动作用于执行

2.3 pointnet_extractor.py——PointNet提取器:3D扩散策略的空间感知基础

pointnet_extractor.py文件实现了改进型3D扩散策略(IDP3)的空间感知核心,它负责从原始点云数据中提取结构化特征,为机器人控制提供空间理解能力

2.3.1 点云处理工具集

文件开头定义了一系列点云操作函数,分别基于NumPy和PyTorch实现。这些函数处理三个关键任务:

  1. 点云随机打乱:通过`shuffle_point_numpy/torch`函数重新排列点的顺序,增强模型对点序的不变性
  2. 点云填充:当点数不足时,`pad_point_numpy/torch`函数添加零点并随机打乱,确保批处理一致性
  3. 均匀采样:`uniform_sampling_numpy/torch`函数实现点云下采样或上采样,控制计算复杂度

2.3.2 多阶段PointNet编码器

`MultiStagePointNetEncoder`类是该文件的核心,它实现了一种改进的PointNet架构:

  1. 首先将点云转置为形状`[B, 3, N]`的形式,便于卷积处理
  2. 通过初始卷积层提取每个点的局部特征
  3. 在多个阶段中,每个点的特征与全局特征(通过max pooling获得)连接,形成全局-局部上下文
  4. 各阶段特征通过跳跃连接汇总,通过输出卷积层生成最终特征向量

这种多阶段设计使编码器能够捕捉点云的层次化信息,相比原始PointNet更适合复杂的空间理解任务

2.3.3 IDP3编码器:多模态融合

`IDP3Encoder`类整合了点云和状态信息,形成对环境的完整理解:

  1. 可选地对点云进行下采样处理
  2. 通过多阶段PointNet提取点云特征
  3. 通过MLP处理机器人状态向量
  4. 将两种特征连接形成最终的环境表示

该架构的创新之处在于它无缝融合了离散的3D空间信息和连续的机器人状态,为扩散模型的条件生成提供了丰富的环境上下文,使机器人能够根据复杂的3D场景做出精确决策

// 待更

第三部分 Fourier数据集的可视化和处理与文档扩展

3.1 pyproject.toml中的修改:以支持Fourier数据集的可视化和处理

添加了`fourier_viz`可选依赖组,包含以下软件包:

fourier_viz = [
    "opencv-python>=4.10.0.84",
    "rerun-sdk==0.22.0",
    "h5py>=3.12.1",
    "tqdm>=4.67.1",
    "loguru",
    "numpy",
    "rich",
]

这些依赖项专门用于支持Fourier数据集的可视化和处理

3.2 文档扩展

Fourier-Lerobot相比原始的Lerobot,其

  1. 添加了完整的`DATASET.md`文件,详细说明了Fourier ActionNet数据集的结构和使用方法
  2. 修改了`README.md`,添加了关于Fourier-Lerobot的介绍段落和功能说明

// 待更

第四部分 傅利叶开源人形机器人数据集Fourier ActionNet

// 待更

第五部分 我司七月的ipd3落地实践(部分)

七月内部从24年q3开始,一直给各种工厂做场景落地和定制开发,25年q2开始,一个个过了保密期之后,可以逐一拿出部分 分享下

// 待更