This repository provides training and evaluation scripts for a family of stage-aware reward models on both the LeRobot dataset and raw robot trajectories:

• SARM — single-task, stage-aware reward model (a stage estimator + a subtask progress head).
• SARM2 — multi-task stage-aware reward model that replaces SARM's per-task stage annotations with a task-agnostic action-primitive vocabulary and a multi-gate Mixture-of-Experts (MMoE) value head, producing dense per-step rewards across many manipulation tasks within a single model.
• ReWiND — our reproduction of the ReWiND baseline.

Each model also supports offline reward labeling of raw robot rollouts (used to relabel autonomous rollouts in a self-improvement loop).

All configs live under config/ (label configs under config/label/); encoders, model sizes, dataset paths, and checkpoint paths are set there.

SARM2 is a multi-task stage-aware reward model. Three camera views plus proprioceptive state are encoded by a shared frozen SigLIP-2 backbone, whose cached frame embeddings feed two separately trained causal Transformers:

1. Action-Primitive Stage Estimator (act_pri) — a task-agnostic 4-layer causal Transformer that classifies the current segment over K+1 = 22 candidates (K = 21 action primitives + a null/fallback class). Because the primitive vocabulary is shared across tasks, this stage estimator transfers to new tasks without per-task annotation. Implemented as ActionTransformer (models/action_estimator.py), trained with cross-entropy.

2. MMoE Value Decoder (sarm2) — a 6-layer causal Transformer whose head is a multi-gate Mixture-of-Experts. The K+1 primitives are clustered into M+1 = 8 semantic groups (Acquire, Release, Translate, Insert, Shape Clothes, Rotate, Force, Other); the predicted primitive selects the corresponding gate, which routes the fused token through top-k of a shared expert pool. The head predicts a normalized remaining-steps-to-completion value r* = -(T-t)/T ∈ [-1, 0]. Implemented as the multi-gate RewardTransformer (models/moe_deco_gate_reward_model.py), trained with MSE + per-gate load-balance / entropy auxiliary losses.

At inference the predicted primitive both supplies the value head's primitive embedding and selects its MoE gate, so the value estimate specializes to the action currently being performed.

We recommend using uv for dependency management.

git clone https://github.com/xdofai/opensarm

pip install uv

uv sync

source .venv/bin/activate

python train.py --config-name sarm

SARM2 is trained in two stages — first the action-primitive stage estimator, then the MMoE value decoder that consumes it:

# 1) Train the action-primitive stage estimator
python train.py --config-name act_pri

# 2) Point sarm2's `model.act_pri_model_path` at the act_pri checkpoint from step 1,
#    then train the MMoE value decoder (act_pri can be frozen or jointly finetuned
#    via `model.finetune_act_pri`).
python train.py --config-name sarm2

python train.py --config-name rewind

eval.py selects what to run with --mode (default eval):

# Validation set
python eval.py --config-name sarm2
python eval.py --config-name act_pri

# Raw robot trajectories
python eval.py --config-name sarm2 --mode raw_data

Any trained reward model can label a directory of raw robot rollouts with dense per-frame rewards. Each rollout writes reward.npy and progress.npy into its episode folder (per-frame, aligned to the episode's timestamp.npy). Set eval.label_data_dir (and eval.ckpt_path*) in the chosen config.

# SARM2 / SARM / ReWiND labelers (configs live under config/label/)
python eval.py --config-name label/sarm2_label  --mode label
python eval.py --config-name label/sarm_label   --mode label
python eval.py --config-name label/rewind_label --mode label

# Multi-GPU labeling, or clear existing labels
python eval.py --config-name label/sarm2_label  --mode label_mp
python eval.py --config-name label/sarm2_label  --mode clear

The code for labeling rollouts to finetune the reward model lives in label_data.

• Top-level configs are under config/; labeling configs are under config/label/.
• Swap --config-name among rewind, sarm, act_pri, sarm2 (and their label/* variants) to run a different model.
• SARM/ReWiND use a CLIP encoder; SARM2 (act_pri + sarm2) uses a frozen SigLIP-2 encoder — keep a config's encoder consistent with the checkpoint you load.

A typical LeRobotDataset looks like this from its root path:
        .
        ├── data
        │   ├── chunk-000
        │   │   ├── episode_000000.parquet
        │   │   ├── episode_000001.parquet
        │   │   ├── episode_000002.parquet
        │   │   └── ...
        │   ├── chunk-001
        │   │   ├── episode_001000.parquet
        │   │   ├── episode_001001.parquet
        │   │   ├── episode_001002.parquet
        │   │   └── ...
        │   └── ...
        ├── meta
        │   ├── episodes.jsonl
        │   ├── info.json
        │   ├── stats.json
        │   └── tasks.jsonl
        └── videos
            ├── chunk-000
            │   ├── left_camera-images-rgb
            │   │   ├── episode_000000.mp4
            │   │   ├── episode_000001.mp4
            │   │   ├── episode_000002.mp4
            │   │   └── ...
            |   ├── right_camera-images-rgb
            │   │   ├── episode_000000.mp4
            │   │   ├── episode_000001.mp4
            │   │   ├── episode_000002.mp4
            │   │   └── ...
            |   ├── top_camera-images-rgb
            │   │   ├── episode_000000.mp4
            │   │   ├── episode_000001.mp4
            │   │   ├── episode_000002.mp4
            │   │   └── ...
            ├── chunk-001
            └── ...

Each episode_XXXXXX.parquet stores a time-series trajectory.
Each row corresponds to one timestep and contains the following columns:

Note: Despite its name, reward stores the absolute task progress (value) rather than a reinforcement learning reward.

• The repository structure is adapted from diffusion_policy.
• The dataset format follows the LeRobotDataset specification.
• The ReWiND model included in this repository is our reproduction from the paper; please refer to the official ReWiND repository for the original implementation.

If you find our papers or code useful, please consider citing:

@inproceedings{
chen2026sarm,
title={{SARM}: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation},
author={Qianzhong Chen and Justin Yu and Mac Schwager and Pieter Abbeel and Fred Shentu and Philipp Wu},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
}

@article{chen2026sarm2,
  title={SARM2: Multi-Task Stage Aware Reward Modeling for Self Improving Robotic Manipulation},
  author={Chen, Qianzhong and Zheng, Hau and Yu, Justin and Huang, Suning and Sun, Jiankai and Goldberg, Ken and Wen, Chuan and Abbeel, Pieter and Shentu, Yide and Wu, Philipp and others},
  journal={arXiv preprint arXiv:2606.10305},
  year={2026}
}