license: apache-2.0
library_name: lerobot
pipeline_tag: robotics
tags:
- lerobot
- reward-model
- rewind
- progress-reward
- dexmimicgen
- bimanual
- robotics
datasets:
- chomeed/dexmg_two_arm_threading
base_model: facebook/dinov2-base
ReWiND-state reward model -- DexMimicGen Two-Arm Threading
Progress reward model (ReWiND-state, LeRobot plugin type rewind_state) for the bimanual task Pick up the needle
and thread it through the hole in the tripod. It predicts task progress in [0, 1] per frame from 3 cameras, the
robot state and the instruction. It was trained on successful demos only, so it never sees a failure.
This exact file is the progress-only baseline of the FARM comparison, the starting point of the failure-aware
companion chomeed/farm_rewind_state_dexmg_two_arm_threading,
and the dense reward of the "ReWiND" RL runs summarised under Evaluation. On fresh rollouts, a failed attempt that
looks like progress keeps a high reward: failed episodes end at 0.92, successful ones at 0.94.
Training data
chomeed/dexmg_two_arm_threading: all 1,025
DexMimicGen-generated success demos (218,858 frames, 20 fps; the normalizer statistics count every frame). There is
no validation split and there are no failures.
Objective (ReWiND). Per-frame MSE between the sigmoid output and the progress label (i - clip start) / (episode end - clip start - 1). Clips are random contiguous segments of random demos. Half of them are "rewound": a forward segment from the first half of the episode into the second half, followed by a reversed tail whose labels decrease. Clips are uniformly subsampled or padded to 16 frames. No negatives, no language perturbation, no image augmentation.
Model and configuration
Inputs per frame: 3 RGB cameras at 84x84 (agentview, robot0_eye_in_hand, robot1_eye_in_hand), the 18-d
observation.state (mean/std normalised with the demo statistics) and the task string. DINOv2-base encodes each
camera at its native 84x84 (a 6x6 patch grid), pooled with the last layer's CLS attention; its last 2 blocks and final
LayerNorm are fine-tuned, the rest is frozen. The instruction goes through a frozen
sentence-transformers/all-MiniLM-L12-v2 inside the preprocessor (not stored in this repo). Per-modality MLPs, with
the state added to every token, feed a 4-layer Transformer (512-d, 8 heads) over a 16-frame window; the head is
LayerNorm -> MLP -> sigmoid. 102,945,537 parameters, 30,545,409 trainable.
| setting | value |
|---|---|
| checkpoint | final step 50,000 (checkpoints/last); model.safetensors sha256 8dad424e31cfb3e6a0b896cf1e00459137684a906505fb781019f390f47a3bdc |
| batch | 16 clips of 16 frames |
| optimiser | AdamW (0.9, 0.95), weight decay 1e-10, grad clip 10, bf16 mixed precision, seed 1000, 1 GPU |
| learning rate | 1e-4 for all non-DINOv2 parameters (MLPs, Transformer, head), 1e-5 for the fine-tuned DINOv2 blocks and final LayerNorm; 1,000-step warmup, cosine decay over 50k steps to 5 % of the base rate |
scheduler_peak_lr: 5e-5 in the config does not set the peak: LeRobot's cosine schedule scales the optimizer LR and
uses peak_lr only for the 5 % floor ratio.
Evaluation
Causal window (hybrid8). For t >= 16, the reward at frame t is the last output of a 16-frame window:
np.linspace(0, t-8, 8).astype(int) followed by frames t-7..t. Only frames <= t are used. This is also the window of
the RL reward shaper (RewindRewardShaper, trajectory_mode="hybrid", window_recent_n=8).
Protocol. Episode score = mean of the last 3 per-frame rewards, threshold 0.5. Failures caught = failure score < 0.5; successes recognised = success score >= 0.5; balanced accuracy = the mean of the two. AUROC = probability that a random success outscores a random failure.
Fresh held-out rollouts (headline). 100 new rollouts of the frozen base Diffusion Policy with env seeds 20000-20099, disjoint from every seed used for training, RL or earlier evaluation: 57 successes (8.8-13.85 s long) and 43 failures (all 25 s time-outs). No reward model saw them.
| model | success score | failure score (lower is better) | failures caught | successes recognised | balanced acc. | AUROC |
|---|---|---|---|---|---|---|
| ReWiND-state (this model) | 0.94 | 0.92 | 1/43 | 57/57 | 0.51 | 0.33 |
| FARM, step 5k (companion) | 0.98 | 0.003 | 43/43 | 56/57 | 0.99 | 1.00 |
- An AUROC of 0.33 means failures outrank successes more often than not. On a fresh failure the reward climbs to a per-episode peak of 0.97 on average and stays near it.
- In-sample reference, not out-of-sample: the last 50 demos (975-1024, in this model's training data) and failed rollouts 150-199 (never seen by this model, but mined by FARM): 0.98 / 0.92 / 1/50 / 50/50 / 0.51 / 0.38.
- As an RL reward (ResFit-TD3, same window): 76.6 % mean evaluation success over 730-780k steps, against 40.3 % with the sparse environment reward, 56.6 % with a ReWiND variant trained with failures labelled 0, and 71.3 % with FARM v6, an earlier failure-aware checkpoint (its value is digitised from a plot of its best replicates, so it is optimistic). At most 3 seeds; one 96-episode evaluation has an SD of about 5 points, and the evaluated sparse runs are the weaker replicate. In a separate 400k-step sweep, step 2,000 of the FARM v7 run was within seed noise of this model (training success 0.54 vs 0.51 at reward scale 0.5, 3 seeds each).
How the FARM companion lowers the failure reward
FARM fine-tunes this checkpoint for 5,000 steps with 200 failed rollouts. On each failure it finds the critical point
t*, the peak of this model's 7-frame-smoothed reward (the furthest the attempt got), scored with a 16-frame uniform
window (frames linspace(0, t, 16)) rather than the hybrid8 window above; the two windows give different t* in
46 of 50 checked failures. A hinge on the reward logits then penalises the reward unless it falls by at least 0.2
logits from each of 16 anchors to the next over frames t*-5..t*+60, each anchor scored in the causal hybrid8
window. Successes are protected by this model's progress loss plus a no-decrease hinge and a progress MSE on matched
success windows. Failures never get a "0" label, so the progress credit before t* is kept. Details are on the
FARM card.
Loading
A sketch of the evaluation path (reference implementation: load_rewind_reward_model +
compute_rewind_reward_for_episode); this exact snippet was not run on its own.
import lerobot_policy_rewind_state # registers `rewind_state` (local plugin, see caveats)
from huggingface_hub import snapshot_download
from lerobot.configs.policies import PreTrainedConfig
from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
from lerobot.policies.factory import make_policy, make_pre_post_processors
from lerobot_policy_rewind_state.common.eval_rewind_reward_visualization_dense import dense_predict_per_frame
path = snapshot_download("chomeed/rewind_state_dexmg_two_arm_threading")
cfg = PreTrainedConfig.from_pretrained(path)
cfg.pretrained_path = path # required: config.json has pretrained_path = null (make_policy would build a random head)
cfg.device = "cuda"
meta = LeRobotDatasetMetadata("chomeed/dexmg_two_arm_threading")
policy = make_policy(cfg=cfg, ds_meta=meta).eval()
pre, _ = make_pre_post_processors(policy_cfg=cfg, pretrained_path=None, dataset_stats=meta.stats)
# cams: {"observation.images.agentview", "...robot0_eye_in_hand", "...robot1_eye_in_hand"} -> float (1, T, 3, 84, 84) in [0, 1]
# state: raw observation.state, shape (1, T, 18)
batch = pre({**cams, "observation.state": state,
"task": ["Pick up the needle and thread it through the hole in the tripod"]})
reward = dense_predict_per_frame(policy, {k: batch[k] for k in cams}, batch["task.minilm"],
state=batch["observation.state"], max_window=16, sampling="hybrid", n_recent=8)
Caveats
Progress only. The model never saw a failure, so a failed attempt that looks like progress keeps a high reward; fresh failures end at 0.92.
Window layout differs between training and use. Training clips were contiguous segments, uniformly subsampled
or padded to 16 frames. The hybrid8 window used for evaluation and RL (8 sparse history frames + 8 consecutive
recent frames) was never seen in training.
Causal only through the window. The trunk lays tokens out stream-major, masks them with triu and regroups them
with h.view(B, T, -1), so early output positions inside a window see later frames. Read only the last output of a
window that ends at t. dense_predict_per_frame scores frames 0-15 from one pass over frames 0-15 (outputs 0-14 see
later frames, up to frame 15); for strictly causal online use, feed the prefix 0..t and read its last output, as the
RL shaper does. Episode scores (last 3 frames) are unaffected.
Episode length is a confound in every DexMimicGen evaluation: all failures are 25 s time-outs, while successes end after 8.8-13.85 s.
In-sample is not held out. There is no train/val split, so the evaluator's "validation" slice (the last 50 demos) is training data. In-sample successes are generated demos; fresh successes are policy rollouts.
Loading. Needs the LeRobot plugin package lerobot_policy_rewind_state, which has no public code location yet.
The results were produced with a locally modified LeRobot 0.4.4, and the snippet is untested on stock LeRobot:
config.json loads through the plugin, but train_config.json contains fork-only fields (dataset.cf_aug, val,
test) that stock LeRobot may not parse. Set cfg.pretrained_path before make_policy. config.json's repo_id
(chomeed/rewind_state_dexmg_threading_v1) is the internal run name, not this repo. The 14-d action output and the
postprocessor exist for LeRobot compatibility only; use predict_reward (select_action raises
NotImplementedError).
Companion artifacts
- Failure-aware fine-tune of this model:
chomeed/farm_rewind_state_dexmg_two_arm_threading - Same baseline on Three-Piece Assembly:
chomeed/rewind_state_dexmg_two_arm_three_piece_assembly - Datasets, base policies and the MimicGen models are in the MimicGen & DexMimicGen Reward Modeling collection.