ReWiND-state reward model -- DexMimicGen Two-Arm Three-Piece Assembly

Progress reward model (ReWiND-state, LeRobot plugin type rewind_state) for the bimanual task Assemble the three pieces: place the two pieces onto the base. 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. Architecture, objective and configuration are identical to the Threading baseline chomeed/rewind_state_dexmg_two_arm_threading; only the dataset differs.

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_three_piece_assembly, and the dense reward of the "ReWiND" RL runs summarised under Evaluation. Unlike on Threading, its failure reward already falls on most Three-Piece failures.

Training data

chomeed/dexmg_two_arm_three_piece_assembly: all 1,006 DexMimicGen-generated success demos (239,827 frames, 20 fps; median 238 frames, longest 272 frames = 13.6 s; 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 `cf7faf95a672fcaeb3710e6f37d95ebc02744429ed780bc923830f2f075b0af4`
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: 63 successes (10.25-15.1 s long) and 37 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.91	0.19	31/37	61/63	0.90	0.99
FARM, step 5k (companion)	0.93	0.003	37/37	61/63	0.98	0.998

The two successes it misses are both longer than any training demo: 15.1 s (the longest fresh success, score 0.02) and 14.2 s (score 0.27). Both reached 1.0 before dropping, at about frame 238, the median demo length. The three other fresh successes at least as long as the longest demo (two longer, one as long) score >= 0.998.
In-sample reference, not out-of-sample: the last 50 demos (956-1005, in this model's training data) and failed rollouts 150-199 (never seen by this model, but mined by FARM): 1.00 / 0.19 / 43/50 / 50/50 / 0.93 / 1.00; every success (lowest 0.996) outscores every failure (highest 0.980).
As an RL reward (ResFit-TD3, same window): 73.1 % mean evaluation success over 370-420k steps, against 74.5 % with the sparse environment reward (2 seeds) and 69.1 % with FARM v6, an earlier failure-aware checkpoint. No reward separated on Three-Piece.

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_three_piece_assembly")
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_three_piece_assembly")
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": ["Assemble the three pieces: place the two pieces onto the base"]})
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. On Three-Piece it already lowers most failure rewards, but 6/37 fresh failures still end at >= 0.5.

Long successes can score low. Both missed fresh successes are longer than the longest training demo (13.6 s) and drop after reaching 1.0 near the median demo length, which suggests the reward partly depends on elapsed time or position in the episode.

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 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: all failures are 25 s time-outs, while successes end after 10.25-15.1 s.

In-sample is not held out. There is no train/val split, so the in-sample successes (956-1005) are training demos.

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