poolvarine/SARM-robosuite-can-mh-stages

SARM Robosuite PickPlaceCan (Multi-Human) — Stage-Annotated Dataset

Stage-annotated demonstration dataset for the PickPlaceCan task from robosuite, formatted in LeRobot v2.1 format for training SARM (Stage-Aware Reward Models).

Dataset Summary

Features

Stage Annotations (4 Stages)

Each stage corresponds to a task prompt. Frames are assigned to the stage they belong to — not edges/transitions.

Reward Encoding

The reward field encodes a continuous progress label for SARM training:

reward = stage + (frame_position_in_stage / total_frames_in_stage)

For example, if stage 0 has 60 frames:

• —Frame 0: reward = 0.0000
• —Frame 30: reward = 0.5000
• —Frame 59: reward = 0.9833

SARM decodes this as:

• —floor(reward) = stage class (0, 1, 2, or 3)
• —frac(reward) = intra-stage progress (0.0 to ~1.0)

Reward range: [0.0, 4.0) — never reaches 4.0.

Use num_classes_sparse: 4 in SARM config.

Stage Detection Method

Stages are detected via direct MuJoCo state reading from robosuite simulation:

• —End-effector position from sim.data.site_xpos[eef_site_id]
• —Can position from sim.data.body_xpos[can_body_id]
• —Gripper contact from sim.data.contact geom-pair matching
• —Bin bounds from static placement geometry

The HDF5 stores only integer stage labels per frame. The continuous reward is computed during LeRobot conversion.

Usage with SARM / LeRobot

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"] -> continuous progress labels for reward model training
# sample["agentview-images-rgb"] -> decoded video frames

SARM Config (PickPlaceCan)

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

Citation

If you use this dataset, please cite the SARM paper and robosuite:

@inproceedings{robosuite2020,
  title={robosuite: A Modular Simulation Framework and Benchmark for Robot Learning},
  author={Zhu, Yuke and Wong, Josiah and Mandlekar, Ajay and others},
  booktitle={arXiv preprint arXiv:2009.12293},
  year={2020}
}