SARM-robosuite-can-mh-stages

SARM Robosuite PickPlaceCan (Multi-Human) 数据集是一个专为训练阶段感知奖励模型（SARM）设计的机器人操作演示数据集。该数据集基于robosuite仿真环境中的拾放罐子任务，采集了多人类演示数据，并按照LeRobot v2.1格式进行组织。数据集包含300个完整操作片段，总计62,756帧数据，帧率为20 FPS。数据模态包括机器人状态、动作指令、连续进度奖励标签以及来自两个视角的视觉观测：第三人称视角（agentview）和安装在机器人手腕上的视角（robot0_eye_in_hand），图像分辨率均为256x256。每个数据帧都被标注为四个任务阶段之一：机器人接近罐子、机器人抓取并抬起罐子、机器人将罐子移向放置箱和机器人将罐子放入箱中。奖励字段（reward）编码了一个连续的进度标签，其整数部分表示阶段类别，小数部分表示阶段内的相对进度，范围为[0.0, 4.0)。阶段标签是通过直接读取仿真中的MuJoCo状态（如末端执行器位置、物体位置、接触信息）自动检测生成的。该数据集主要用于机器人模仿学习、奖励模型训练，特别是SARM这类需要细粒度阶段标注和进度监督的算法。

The SARM Robosuite PickPlaceCan (Multi-Human) dataset is a robotic manipulation demonstration dataset specifically designed for training Stage-Aware Reward Models (SARM). It is based on the pick and place can task in the robosuite simulation environment, collecting multi-human demonstration data and organized according to the LeRobot v2.1 format. The dataset contains 300 complete operation episodes, totaling 62,756 frames of data at a frame rate of 20 FPS. Data modalities include robot state, action commands, continuous progress reward labels, and visual observations from two perspectives: a third-person view (agentview) and a view mounted on the robots wrist (robot0_eye_in_hand), both with an image resolution of 256x256. Each data frame is annotated as one of four task stages: robot approaches the can, robot grasps and lifts the can, robot moves the can towards the bin, and robot places the can into the bin. The reward field encodes a continuous progress label where the integer part indicates the stage category and the fractional part indicates the relative progress within the stage, ranging from [0.0, 4.0). Stage labels are automatically generated by directly reading MuJoCo state in the simulation (such as end-effector positions, object positions, and contact information). This dataset is primarily used for robot imitation learning and reward model training, especially for algorithms like SARM that require fine-grained stage annotation and progress supervision.

原始信息汇总

数据集概述

SARM Robosuite PickPlaceCan (Multi-Human) — Stage-Annotated Dataset 是一个为机器人操作任务设计的分阶段标注数据集，专门用于训练 SARM（Stage-Aware Reward Models）。该数据集基于 robosuite 环境中的 PickPlaceCan 任务，并由多人类演示生成。

---

基本信息

属性	值
任务	抓取并放置易拉罐（Pick and Place Can）
机器人	Panda（Franka Emika）
环境	robosuite PickPlaceCan
数据来源	robosuite 多人类（Multi-Human）演示
总片段数	300
总帧数	62,756
帧率	20 FPS
摄像头	agentview（第三人称）、robot0_eye_in_hand（腕部）
图像分辨率	256 x 256
视频编码	AV1（libsvtav1）
LeRobot 版本	v2.1

---

数据特征

特征	类型	形状	描述
state	float32	(9,)	末端执行器位置 (3) + 四元数 (4) + 夹爪位置 (2)
actions	float32	(7,)	末端执行器增量动作（6 自由度 + 夹爪）
reward	float32	(1,)	连续进度标签（详见奖励编码）
agentview-images-rgb	视频	(256, 256, 3)	第三人称摄像头视角
robot0-eye-in-hand-images-rgb	视频	(256, 256, 3)	腕部摄像头视角

---

阶段标注（共 4 个阶段）

每个阶段对应一个任务提示，帧被分配到其所属的阶段（不包含边界/过渡帧）：

阶段编号	提示	描述
0	“机器人接近易拉罐”	从片段开始到夹爪接触易拉罐
1	“机器人抓取并举起易拉罐”	从接触开始到将易拉罐抬离桌面
2	“机器人将易拉罐移向放置箱”	运输抓取的易拉罐至放置箱
3	“机器人将易拉罐放入放置箱”	降低/释放易拉罐至箱中，直到完成

---

奖励编码

reward 字段编码了一个用于 SARM 训练的连续进度标签：

reward = 阶段编号 + (帧在阶段中的位置 / 阶段总帧数)

例如，如果阶段 0 有 60 帧：

• 第 0 帧：reward = 0.0000
• 第 30 帧：reward = 0.5000
• 第 59 帧：reward = 0.9833

SARM 解码方式：

• floor(reward) = 阶段类别（0, 1, 2, 或 3）
• frac(reward) = 阶段内进度（0.0 到约 1.0）

奖励范围：[0.0, 4.0) —— 永远不会达到 4.0。

在 SARM 配置中应使用 num_classes_sparse: 4。

---

阶段检测方法

阶段通过直接读取 robosuite 仿真中的 MuJoCo 状态 来检测：

• 末端执行器位置：来自 sim.data.site_xpos[eef_site_id]
• 易拉罐位置：来自 sim.data.body_xpos[can_body_id]
• 夹爪接触：来自 sim.data.contact 的几何体对匹配
• 放置箱边界：来自静态放置几何体

HDF5 文件仅存储每帧的整数阶段标签。连续奖励在 LeRobot 转换期间计算。

---

在 SARM / LeRobot 中使用

python from lerobot.common.datasets.rm_lerobot_dataset import FrameGapLeRobotDataset

dataset = FrameGapLeRobotDataset( repo_id="poolvarine/SARM-robosuite-can-mh-stages", episodes=list(range(300)), n_obs_steps=4, frame_gap=1, max_rewind_steps=2, image_names=["agentview-images-rgb"], task_name="pick and place can", )

sample = dataset[0]

sample["state"] -> (n_obs_steps+max_rewind_steps+1, 9)

sample["targets"] -> 用于奖励模型训练的连续进度标签

sample["agentview-images-rgb"] -> 解码后的视频帧

SARM 配置（PickPlaceCan）

yaml general: camera_names: ["agentview-images-rgb"] state_dim: 9 model: num_classes_sparse: 4 annotation_list: ["approaching", "grasping_lifting", "transporting", "placing"]

搜集汇总

数据集介绍

构建方式

该数据集源自robosuite仿真环境中的PickPlaceCan任务，借助Panda（Franka Emika）机器人平台，采集了多人类演示数据。数据集包含300个情节，总计62,756帧，帧率为20 FPS。通过直接读取MuJoCo模拟状态，包括末端执行器位置、罐体位置及夹爪接触信息，对每帧进行阶段标注，将任务划分为四个连续阶段：接近罐体、抓取与提升、搬运至目标区域、放置入箱。在转换为LeRobot v2.1格式时，将离散的阶段标签编码为连续的进度奖励值，最终以parquet文件形式存储，支持双摄像头视角（第三视角与腕部视角）的RGB图像。

特点

数据集的核心特点在于其阶段感知奖励编码机制。奖励字段被设计为连续进度标签，数值范围为[0.0, 4.0)，通过整数部分表示阶段类别（0至3），小数部分反映阶段内进度。这种编码方式使得SARM模型能够同时学习阶段离散分类与阶段内连续进度，提升奖励模型的细粒度判别能力。此外，数据集提供了两种视角的高清视频（256×256分辨率，AV1编码）、9维状态向量（末端位置、四元数与夹爪位置）及7维动作信息，全面支持机器人操作任务的训练需求。

使用方法

数据集可通过LeRobot框架中的FrameGapLeRobotDataset类轻松加载。用户需指定仓库ID、情节范围、观测步数、帧间隔等参数，并选择所需摄像头视角。加载后的样本包含状态张量、连续进度标签及解码后的视频帧，可直接用于SARM奖励模型的训练。建议在SARM配置中将num_classes_sparse设置为4，以匹配四个阶段分类任务。数据集默认采用agentview视角图像作为主要视觉输入，并支持灵活调整状态维度与图像名称，适用于模仿学习与奖励建模等多种机器人学习场景。

背景与挑战

背景概述

在机器人操作领域，奖励函数的设计是强化学习与环境交互的核心瓶颈之一，传统稀疏奖励难以引导策略在长程任务中稳定收敛。SARM（Stage-Aware Reward Models）方法通过将复杂任务分解为若干有序阶段，并利用连续进度标签进行细粒度监督，为解决这一难题提供了新思路。该数据集由研究机构于2025年基于robosuite仿真环境创建，聚焦于PickPlaceCan（拾取与放置易拉罐）操作任务，采用Franka Emika Panda机器人采集了300条多人类演示轨迹，总计62,756帧图像与状态数据。数据集的独特之处在于为每一帧标注了四个阶段——接近、抓取与提升、搬运、放置——并编码为连续的进度奖励，为SARM奖励模型训练提供了高质量、结构化的基准。该数据集对机器人学习领域的影响在于，它填补了阶段性标注操作数据的空白，有效推动了基于进度感知的奖励模型在长程操作任务中的应用与发展。

当前挑战

该数据集所解决的领域问题在于，传统机器人操作任务中的奖励学习常受限于稀疏或二值化的成功信号，难以提供密集且语义明确的引导，而SARM-robosuite-can-mh-stages通过阶段划分与连续进度编码，实现了对操作过程的细粒度监督，提升了策略学习效率。在构建过程中面临的核心挑战包括：首先，阶段检测依赖直接读取MuJoCo仿真状态（如末端执行器位置、抓取器接触检测及物体坐标），这种精确的物理状态追踪在真实场景中难以复现，限制了数据集向真实世界的迁移。其次，多人类演示轨迹间的风格、速度与阶段性过渡存在显著差异，如何通过统一的编码规则（如机器人接近阶段定义）确保跨轨迹的标注一致性是一大难题。此外，数据压缩采用AV1视频编码，虽然节省了存储空间，但在解码与实时训练时可能引入额外的计算开销。

常用场景

经典使用场景

在机器人操作领域，精细化的任务分解与过程监督是提升学习效率的关键。该数据集专为机器人抓取与放置任务设计，通过将‘拿起放置罐头’这一复合操作精准划分为四个子阶段——接近、抓取提升、运输和放置，为阶段感知奖励模型（SARM）的训练提供了结构化监督信号。研究者可利用该数据集对机器人操控策略进行细粒度学习，例如训练模型在不同阶段输出差异化的动作指令，从而实现从粗放式端到端控制到分阶段精细决策的范式转变，显著提升复杂装配或拾放任务的执行成功率。

解决学术问题

传统机器人模仿学习常受困于稀疏奖励导致的探索效率低下问题，而密集奖励的人工设计又极为繁琐且易引入偏差。该数据集通过创新的连续进度标签编码机制，将阶段信息与帧内进度融合为单一奖励值，解决了如何在不依赖手工规则的前提下为长程任务提供渐进式监督的学术难题。这一设计使得奖励模型能够同时捕捉离散的阶段边界与连续的进程动态，为研究多阶段任务中的信用分配、时序对齐和子目标解耦提供了标准化数据基准，推动了机器人学习从‘结果导向’向‘过程可解释’的演进。

衍生相关工作

该数据集衍生出的核心工作是SARM（Stage-Aware Reward Model）框架，其通过在潜在空间中嵌入阶段先验，使得奖励模型能够在不依赖完整状态观测的情况下提供过程监督。后续研究在此基础上扩展了跨任务迁移能力，例如将PickPlaceCan数据集的阶段标注模式应用于其他robosuite操作任务（如螺母拧紧或抽屉开启），验证了阶段分解策略的泛化性。此外，部分工作探索了将连续进度标签与因果推理结合，用于提升奖励模型对分布外场景的鲁棒性，形成了从数据集构建到算法优化的完整研究链条，持续推动机器人奖励塑形理论的发展。