Add model card (FARM reward modeling collection)

README.md

@@ -0,0 +1,168 @@
+---
+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.