---
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`](https://huggingface.co/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`](https://huggingface.co/datasets/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](https://huggingface.co/chomeed/farm_rewind_state_dexmg_two_arm_threading).

## 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.

```python
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`](https://huggingface.co/chomeed/farm_rewind_state_dexmg_two_arm_threading)
- Same baseline on Three-Piece Assembly:
  [`chomeed/rewind_state_dexmg_two_arm_three_piece_assembly`](https://huggingface.co/chomeed/rewind_state_dexmg_two_arm_three_piece_assembly)
- Datasets, base policies and the MimicGen models are in the
  [MimicGen & DexMimicGen Reward Modeling](https://huggingface.co/collections/chomeed/mimicgen-and-dexmimicgen-reward-modeling-6ad45df2fafcb692ae6f37b1)
  collection.