Piece / piece-with-pin classifier (two cameras, EfficientNet-B0)
Classifies the piece held by the SO-101 as nopin or pin (perno / tornillo) from the front and the wrist camera. About 85 % accurate on recording sessions it was not trained on, in the same room, with the same operator and lighting as the training sessions. Not tested on another day, lighting or batch of pieces.
Model
One ImageNet-pretrained EfficientNet-B0 (torchvision weights, 4.0 M parameters) shared by both cameras, fine-tuned (backbone at a tenth of the head learning rate). Each camera's feature map is pooled (mean + max), the two are concatenated and a linear layer gives the two logits. Input: two 448x336 RGB images (centre-cropped to 4:3, never stretched). Training: 25 epochs, class-weighted cross-entropy, AdamW + OneCycle, strong augmentation (rotation, zoom, shift, flip, saturation / white balance, blur, noise).
Data
1004 front+wrist pairs from the MikelRubi/so101-pick-cube_* recordings (2026-10-23 and 2026-10-24):
- 654 from the last 3 frames of 218 pick-cube episodes (112
pin, 106nopin), labelled by eye from the last frames (pin_labels.json); 16 of thepinand 72 of thenopinlabels are "moderate" (blurred or partial view). 39 further episodes were too ambiguous and 5 had no piece in the gripper or no view of the tray, so they were left out. - 350 snapshot pairs (186
pin, 164nopin) taken on purpose with17b_record_real_so101_live.py(keysp/s), stored in each repo'spin_snapshots/folder.
Evaluation
5-fold cross-validation where every fold holds out whole recording sessions (frames of one episode are near-duplicates, so a random split would look better than it is):
| subset | n | accuracy | majority-class baseline |
|---|---|---|---|
| all pairs | 1004 | 0.85 | 0.52 |
| episode last frames | 654 | 0.87 | 0.51 |
| snapshots | 350 | 0.83 | 0.53 |
| per episode | 568 | 0.84 | 0.52 |
| per episode, clear labels only | 480 | 0.84 | 0.59 |
Per fold: 0.76, 0.85, 1.00, 0.88, 0.88. The final model in this repo is trained on all the data and was not itself evaluated; the numbers above are from the five cross-validation models. Labels are a visual judgement, not ground truth.
Same cross-validation, other runs: ResNet18 (11.2 M parameters) at 448x336 gave 0.85 on all pairs, 0.87 on episode frames, 0.79 on snapshots, 0.82 per episode; ResNet18 trained and tested on the episode frames only gave 0.84 (0.87 when the snapshots were added to the training set). A small CNN from scratch stayed at chance level (0.55). Earlier ResNet18 runs on fewer pairs: 0.66 (627 pairs), 0.81 (905 pairs). More data, pretrained weights and resolution all mattered.
Files
model.pt is a dict with state_dict, channels, height, width, classes, backbone (efficientnet_b0), cameras (both). Load it with load_model / predict from 22_train_pin_classifier.py in demos/06_imitation_learning/ of the LeRobot_ workspace (needs torch and torchvision). metrics.json is the training-run report of the final model; the cross-validation report is in the table above.