FARM reward model -- DexMimicGen Two-Arm Threading

Failure-aware reward model (FARM, v7 recipe) for the bimanual task Pick up the needle and thread it through the hole in the tripod. It is the progress-only ReWiND-state model chomeed/rewind_state_dexmg_two_arm_threading fine-tuned for 5,000 steps on the same success demos plus 200 failed rollouts of the base policy. It is trained so that the reward keeps rising while an attempt makes progress and falls once the attempt starts to fail. The model outputs one reward in [0, 1] per frame. Use it only with the causal hybrid8 window described under Evaluation.

On 100 fresh rollouts that no model had seen, failed episodes end at a reward of 0.003 (baseline: 0.92) and successful ones at 0.98 (baseline: 0.94).

How the failure reward is lowered

FARM never labels a failure "0". It finds where each failed rollout goes wrong and trains the reward to fall over the ~3 s after that point, in the same causal windows the model sees at inference.

Find the critical point t*. The frozen baseline scores each of the 200 failed rollouts frame by frame with a 16-frame uniform window (frames linspace(0, t, 16), not the hybrid8 window used at inference). The scores are smoothed over 7 frames, and t* is their maximum: the furthest the attempt got before it went wrong. The median t* is frame 234 (11.7 s into a 25 s rollout).
Take the frames right after it. The failure clip covers frames t*-5 to t*+60 (about 3.3 s). It is sampled at 16 evenly spaced anchors about 4 frames apart; the first 2 anchors lie just before t*.
Make the reward fall along the clip. Each anchor k is scored in its inference window (8 frames spread over the history before it + the 8 frames ending at it), which gives a reward logit z_k. The hinge relu(z_{k+1} - z_k + 0.2) penalises any anchor that is not at least 0.2 logits below the previous one; when the hinge is fully satisfied, the logit falls by at least 3 (15 x 0.2) across the clip. It is a soft penalty (the mean over the 15 steps, averaged with the success hinge in 5), not a hard constraint.
Keep the credit earned before the failure. There is no absolute target, and no loss term acts on failure frames outside the clip. The reward before t*-5 is left to what the model learned from successes. On the mined failures it sits slightly above the baseline's: the mean over the 5 s before t* is 0.86 (baseline 0.78).
Protect successes. (a) The original ReWiND progress loss keeps training on all 1,025 success demos (MSE to linear progress, with half of the clips "rewound"). (b) Each failure clip comes with 3 matched success windows, ranked by a 50/50 blend of (i) cosine similarity on the normalised state trajectory (weight 0.4) and on the frozen baseline's per-camera Transformer embeddings (0.2 per camera), over 16 subsampled frames, and (ii) banded DTW on the state trajectory (band 3, early-warp penalty 0.5). One of the 3 is drawn per sample and scored at 16 anchors in the same causal window. A hinge without margin, relu(z_k - z_{k+1}), penalises any fall there, and an MSE pulls sigmoid(z_k) to linear progress (anchor frame / (episode length - 1)). On Threading these windows sit at the ends of successful demos (95 % of them end in the last tenth of their demo).

L = L_ReWiND + 1.0 * (L_fail_hinge + L_succ_hinge) / 2 + 1.0 * L_succ_MSE, with the hinge and MSE weights ramped in over the first 500 steps. A Bradley-Terry success-vs-failure term exists in the code but has weight 0 (logged only), and TV smoothing is off. The matching therefore only decides which success windows feed (b); no success-vs-failure comparison enters the loss.

Training data

chomeed/dexmg_two_arm_threading: all 1,025 DexMimicGen-generated success demos (218,858 frames, 20 fps) for the progress term. There is no validation split.
chomeed/dexmg_two_arm_threading_failures: 200 failed rollouts of the frozen base Diffusion Policy, each a 500-frame (25 s) time-out.
A pair manifest mined with the frozen baseline: one failure clip per failed rollout, each with 3 matched success windows (sampled at 16-64 frames; the chosen ones are 16-58 frames long, median 18), taken from 46 different demos.

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 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. Architecture and parameter names are identical to the baseline.

setting	value
initialisation	`chomeed/rewind_state_dexmg_two_arm_threading` (step 50,000; byte-identical file)
released checkpoint	step 5,000 of a run that continued to 100k; `model.safetensors` sha256 `bac26c33cc228344288c17ba61ee59bc370328f27f716e89608f6beb65bc6ee6`
batch	4 samples, each = 1 ReWiND clip + 16 failure-anchor + 16 success-anchor windows (132 windows of 16 frames per step)
optimiser	AdamW (0.9, 0.95), weight decay 1e-10, grad clip 10, bf16 mixed precision, seed 1000, 1x RTX A6000
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 over 100k steps, so still 9.94e-5 at step 5,000
failure hinge	margin 0.2 on logits, 16 anchors over t-5..t+60, `hybrid8` windows
success anchors	16 per window: no-decrease hinge (margin 0) + MSE to linear progress
loss weights	ReWiND 1, hinges 1 (averaged), success MSE 1, Bradley-Terry 0, TV 0; 500-step ramp

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. The hinge and anchor terms were trained on this layout, and the RL reward shaper uses the same window (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 (baseline)	0.94	0.92	1/43	57/57	0.51	0.33
FARM, step 5k (this model)	0.98	0.003	43/43	56/57	0.99	1.00

The one missed success scores 0.48. Every success still outscores every failure (highest failure score: 0.11).
FARM still credits a failed attempt early on: its per-episode peak averages 0.99 (baseline 0.97). Its reward then stays below 0.5 from a median 14.25 s, about 10.8 s before the time-out. The mean reward over all failure frames is 0.35 (baseline 0.72).
Against t* (the 50 mined failures 150-199, in-sample: their post-t* clips were trained): the mean reward over the 5 s before t* is 0.86 (baseline 0.78) and 0.97 around t* (t* +- 0.25 s; baseline also 0.97); only 1/50 is already below 0.5 at t*. The reward first drops below 0.5 a median 2.1 s after t* (84 % within 3 s), and that time correlates with t* at 0.92. The baseline drops below 0.5 on only 2 of these 50, so it has no comparable value.
In-sample reference, not evidence of generalisation (successes 975-1024 are training demos, failures 150-199 were mined for training): FARM 1.00 / 0.001 / 50/50 / 50/50 / 1.00 / 1.00; baseline 0.98 / 0.92 / 1/50 / 50/50 / 0.51 / 0.38.
Step 2,000 of the same run (not uploaded) scored 0.99 / 0.004 / 43/43 / 57/57 / 1.00 / 1.00 on the fresh rollouts.

Loading

A sketch of the evaluation path. The numbers above come from the reference implementation (load_rewind_reward_model + compute_rewind_reward_for_episode), which performs these steps; this exact snippet was not run on its own.

import lerobot_policy_rewind_state_ours_variant7  # registers `rewind_state_ours_variant7` (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/farm_rewind_state_dexmg_two_arm_threading")
cfg = PreTrainedConfig.from_pretrained(path)
# Required: config.json's pretrained_path is the *baseline's* local training checkpoint. Without this line
# make_policy loads the baseline (on the training cluster) or fails (anywhere else).
cfg.pretrained_path = path
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

Episode length is a confound. Every failure is a 25 s time-out, fresh successes end after 8.8-13.85 s, and the score is read at the episode end, so length alone would separate them. On the 50 mined failures, whose post-t* clips were trained, the drop follows t* (correlation 0.92, median lag 2.1 s). That checks the fit to training data, not generalisation: fresh failures have no t*, and failures that end early were not tested.

Slow successes may be under-rewarded. The one missed fresh success (13.45 s, score 0.48) is the second-longest; the longest (13.85 s) scores 0.82, against 0.99 under the baseline.

Only the clip after t* is trained. Low rewards late in a failure, and behaviour on failure modes unlike the 200 training failures (all time-outs of one base policy), are generalisation, not training targets.

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.

In-sample numbers overstate. Two fb variants of this recipe (v3, v5; not released) looked near-perfect in-sample (balanced accuracy 0.99-1.00) but scored fresh successes only 0.27 and 0.64. v8 already lowered successes in-sample (0.58) and scored fresh successes 0.44. Judge models on fresh rollouts.

Two window layouts in training. The ReWiND progress term still trains on uniformly subsampled contiguous clips; only the hinge and anchor terms use hybrid8.

v7, not v6. Our main FARM recipe in the RL experiments was v6 (margin 0.3, no success-anchor MSE); its Threading checkpoint no longer exists, so do not quote v6 numbers for this model. This checkpoint (step 5,000) was not used as an RL reward. Step 2,000 of the same run was, in a 400k-step ResFit-TD3 sweep at reward scale 0.5: training success 0.54 against 0.51 for the ReWiND baseline at the same scale (3 seeds each, within seed noise); its evaluations were not kept. FARM v6 did not beat ReWiND as an RL reward either: 71 % vs 77 % evaluation success on Threading (the FARM value is digitised from a plot of its best replicates, so it is optimistic) and 69 % vs 73 % on Three-Piece.

Loading needs the plugins. lerobot_policy_rewind_state_ours_variant7 and lerobot_policy_rewind_state are local LeRobot plugin packages with 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. Override cfg.pretrained_path: repo_id, pretrained_path (the baseline's local checkpoint) and preference_pairs_path in config.json are internal training values, and the manifest is only needed for training. The 14-d action output and the postprocessor exist for LeRobot compatibility only; use predict_reward.