Insert Mouse Battery VideoMAE Reward
Binary VideoMAE reward model for detecting the completed state of the insert-mouse-battery task from an approximately one-second, three-view clip. It was newly fine-tuned from MCG-NJU/videomae-base-finetuned-kinetics and does not use weights from an earlier reward model.
Input Format
- Eight RGB frames sampled every four source frames at 30 FPS; covered duration is 29 frames (about 0.97 seconds).
- View order per time step:
top_head,hand_left,hand_right. - Each view is resized to
240x320, then vertically stacked to720x320. - Tensor shape:
[batch, 8, 3, 720, 320]. - Normalize with mean
[0.45, 0.45, 0.45]and standard deviation[0.225, 0.225, 0.225]. - VideoMAE tubelet/patch size:
[2, 24, 32].
Output Format
The model returns two logits: not_success and success. The reward is softmax(logits)[..., 1]. The selected validation threshold is 0.92.
Training Data
- Training clips: 2255 (735 dense failure clips added).
- Validation clips: 568 (184 dense failure clips added).
- Positive clips come only from the final two-second stable window of successful episodes, with the complete clip inside that window.
- Negatives include early, middle, hard-preterminal successful clips and dense clips from unsuccessful trajectories.
- Train and validation episodes are disjoint.
- Sources include the processed LeRobot success demonstrations and the 20-success/40-failure rollout dataset.
Validation Results
| metric | value |
|---|---|
| accuracy @ 0.5 | 0.9824 |
| precision @ 0.5 | 0.9505 |
| recall @ 0.5 | 1.0000 |
| F1 @ 0.5 | 0.9746 |
| ROC AUC | 0.9939 |
Full Held-Out Rollouts
| metric | VideoMAE | ResNet18 v4 baseline |
|---|---|---|
| successful episodes detected | 4/4 | 4/4 |
| failed episodes with sustained false trigger | 3/8 | 2/8 |
| success terminal positive rate | 0.9750 | 0.8824 |
| failure frame false-positive rate | 0.0406 | 0.0364 |
A sustained trigger is at least three consecutive sampled predictions above the selected threshold.
Files
| file | description |
|---|---|
model.safetensors | Best VideoMAE reward weights. |
config.json | Transformers model configuration and label mapping. |
preprocessor_config.json | Clip, view, resize, normalization, and threshold metadata. |
evaluation.json | Aggregate and per-episode validation results. |
training_history.json | Per-epoch metrics. |
inference.py | Minimal preprocessing and reward helper. |
Limitations
Completion time is represented by the final stable window because the source datasets provide episode-level success rather than an annotated insertion timestamp. Validate the threshold in the target policy and camera setup before using the score for RL optimization.