Insert Mouse Battery ResNet18 Reward V3
This repository contains a binary ResNet18 reward model for the insert-mouse-battery task, trained from scratch on mixed successful and failed trajectories.
Input Format
- Image input: RGB three-view vertical stack.
- View order:
top_head + hand_left + hand_right. - Image size:
[3, 720, 320]as a CHW tensor. - Layout: each view is resized to
240x320; the stack is720x320. - Preprocessing: RGB conversion, float tensor in
[0, 1], then ImageNet mean/std normalization.
Output Format
The checkpoint outputs one binary trajectory-success logit. Apply sigmoid to obtain the reward score.
Files
| File | Description |
|---|---|
full_weights.pt | ResNet18 reward checkpoint trained from scratch. |
model_metadata.json | Input/output format and dataset-construction metadata. |
eval_summary.json | Episode-level validation metrics. |
Test Results
| split | samples | positives | negatives | accuracy | precision | recall | F1 | loss |
|---|---|---|---|---|---|---|---|---|
| validation | 384 | 192 | 192 | 0.8177 | 0.8081 | 0.8333 | 0.8205 | 0.6416 |
Training Data
The training data combines existing processed LeRobot success trajectories and rollout HDF5 trajectories from m1ku2/assemble_battery_rollout_pi05_5000.
- Processed LeRobot data: 126 train success episodes and 32 validation success episodes.
- Rollout data: 20 success episodes and 40 failure episodes, with a fixed episode-level train/validation split.
- Frames are sampled across early, middle, and late trajectory regions for both classes.
- Train and validation samples are class-balanced; no episode appears in both splits.
- No pretrained reward-model weights were used.