insert-mouse-battery-resnet18-reward-v2

Insert Mouse Battery ResNet18 Reward V2

This repository contains a binary ResNet18 reward model for the insert-mouse-battery task, retrained from scratch on the local processed LeRobot data.

Input Format

• Image input: RGB three-view vertical stack from the processed LeRobot videos.
• View order: observation.images.top_head + observation.images.hand_left + observation.images.hand_right.
• Image size: [3, 1440, 640] as CHW tensor.
• Layout: each view is 480x640; the stacked image is 1440x640.
• Preprocessing: RGB conversion, CHW float tensor in [0, 1], ImageNet mean/std normalization.

Output Format

The checkpoint is a binary reward model. For a single input image, the model outputs one scalar logit. Applying sigmoid gives the estimated full-task success probability.

Files

Test Results

Training Data

The dataset uses three-view videos from the existing processed LeRobot splits.

• Positive samples: tail frames from successful episodes.
• Negative samples: head, early, and middle frames from the same successful episodes.
• Train/validation splitting follows the existing processed split directories and stays episode-level.
• The train and validation sets are class-balanced.
• No pretrained weights from qmddn/insert-mouse-battery-resnet18-reward-v1 were used.