smolvla_so101_screwdriver_box_v2

SmolVLA (450 M parameter vision-language-action model) fully fine-tuned (vision encoder, language model and action expert all trained) from lerobot/smolvla_base on digas99/so101_screwdriver_box: a LeRobot SO-101 follower arm picking objects off a table and dropping them in a box. The instruction names the object, so one policy does all five tasks.

This is v2. v1 (digas99/smolvla_so101_screwdriver_box) trained only the action expert on 201 episodes and, on the arm, approached and transported well but often missed the final grasp. v2 differs in two ways: the dataset grew to 264 episodes, including 23 human-in-the-loop correction clips recorded from exactly the states where v1 failed (episodes 241-263), and the whole network is trained so the vision side can adapt to this scene and these cameras. Keep v1 around for A/B comparison on the arm.

task string (must be passed exactly)	training episodes
Pick up the screwdriver and put it in the box	151
Pick up all the objects and put them in the box	40
Pick up the yellow tape and put it in the box	33
Pick up the nut and put it in the box	20
Pick up the battery and put it in the box	20

Dataset: 264 episodes, 136,258 frames at 30 fps, two cameras (fixed top camera and a wrist camera on the gripper), LeRobot format v3.0, robot_type: so_follower. Box position is fixed in the first 50 episodes, then varied; from episode ~137 several objects lie on the table and the instruction selects the target. Episodes 201-240 are the multi-object task (top-camera colour corrected offline to match the rest); episodes 241-263 are short DAgger-style corrections (8-26 s) that start where v1 failed and end with the object in the box and the arm at home. Some episodes contain natural recoveries (missed grasp, re-approach) and were kept on purpose.

Training

Init	lerobot/smolvla_base, full fine-tune: train_expert_only=false, freeze_vision_encoder=false (lerobot keeps the last VLM layer, lm_head and the final norm frozen)
Precision	bf16 autocast; the VLM backbone is upcast to fp32 master weights for training (lerobot loads it in bf16, which would swallow most AdamW updates at this learning rate), see train_v2.py
Steps / batch size	30,000 / 16 (about 3.5 epochs over the 136,258 frames)
Learning rate	5e-05 (half of v1), 1,000 warm-up steps, cosine decay to 2.5e-6 at the last step
Chunk	chunk_size = 50, n_action_steps = 50, n_obs_steps = 1 (1.67 s of actions per inference at 30 fps)
Image augmentation	LeRobot default image_transforms (brightness, contrast, saturation, hue, sharpness, small affine)
Held-out data	none: all 264 episodes were used for training
Hardware	1x NVIDIA L4 (24 GB) on AWS g6.4xlarge, lerobot 0.6.1, see train_env_versions.txt
Loss	flow-matching MSE; the full training log is in train_log.txt, the exact CLI config in train_config.json; intermediate checkpoints are under checkpoints/<step>/ in this repo

Camera mapping during training: observation.images.top -> camera1, observation.images.wrist -> camera2 (smolvla_base's third camera slot is unused). This rename is stored inside the preprocessor of this checkpoint, so at inference you feed the dataset's own keys (observation.images.top, observation.images.wrist) and never rename anything yourself.

Input / output contract

observation.state           float32 (6,)  shoulder_pan.pos, shoulder_lift.pos, elbow_flex.pos, wrist_flex.pos, wrist_roll.pos, gripper.pos
observation.images.top      uint8 HWC (480, 640, 3) RGB     fixed camera looking down at the table
observation.images.wrist    uint8 HWC (480, 640, 3) RGB     camera on the gripper
task                        str   one of the five strings above, character for character
action                      float32 (6,)  same joint order and units as observation.state

• Joint values are whatever lerobot-record produced with the robot configuration used for recording (the SO-101 in degrees mode: state ranges roughly -115..100). Use the same robot config at inference; normalization statistics live inside the checkpoint (policy_preprocessor_* / policy_postprocessor_*), so you feed raw values and get raw values back.
• Joint order matters: the normalizer is per joint.
• Images must be RGB (not BGR) and ideally the same cameras, mounting and framing as the dataset. SmolVLA resizes with padding to 512x512 internally, so other resolutions load, but changing the cameras means a different distribution.
• Dataset action std per joint, for sanity-checking magnitudes: see train_config.json / the dataset's meta/stats.json.

Run it on the robot

Same environment as training: pip install "lerobot[smolvla]==0.6.1" (other torch versions are fine; the config schema is tied to the lerobot minor version).

lerobot-record \
  --robot.type=so_follower --robot.port=/dev/ttyACM0 --robot.id=<your_follower_id> \
  --robot.cameras='{"top": {"type": "opencv", "index_or_path": 0, "width": 640, "height": 480, "fps": 30},
                    "wrist": {"type": "opencv", "index_or_path": 1, "width": 640, "height": 480, "fps": 30}}' \
  --policy.path=digas99/smolvla_so101_screwdriver_box_v2 \
  --dataset.repo_id=<user>/eval_smolvla_screwdriver_box --dataset.push_to_hub=false \
  --dataset.single_task="Pick up the screwdriver and put it in the box"   # or any of the other four strings \
  --dataset.num_episodes=10 --dataset.episode_time_s=30 --dataset.reset_time_s=10

Use the same camera names (top, wrist) and the same --robot.id / calibration as when the dataset was recorded. Change --dataset.single_task to one of the other four strings to pick a different object.

Or from Python, which is what lerobot-record does internally:

import torch
from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
from lerobot.policies.factory import make_pre_post_processors

PATH = "digas99/smolvla_so101_screwdriver_box_v2"
policy = SmolVLAPolicy.from_pretrained(PATH).to("cuda").eval()
pre, post = make_pre_post_processors(policy.config, pretrained_path=PATH,
                                     preprocessor_overrides={"device_processor": {"device": "cuda"}})
policy.reset()                                   # at the start of every episode
# obs = {"observation.state": tensor(6,), "observation.images.top": uint8 HWC, "observation.images.wrist": uint8 HWC,
#        "task": "Pick up the nut and put it in the box"}
# action = post(policy.select_action(pre(obs)))   # call at 30 Hz; the policy re-infers every 50 calls

If the policy reaches the object but hesitates to close the gripper, try cfg.n_action_steps = 25 (re-observe every 0.8 s instead of 1.7 s) by loading the config with PreTrainedConfig.from_pretrained(PATH), editing it and passing config=cfg to from_pretrained.

This repo is private; hf auth login as a user with access before downloading.

Offline verification before touching the robot

The first from_pretrained needs the VLM backbone repo in the Hub cache (HuggingFaceTB/SmolVLM2-500M-Video-Instruct, 1.9 GB, downloaded automatically; hf download HuggingFaceTB/SmolVLM2-500M-Video-Instruct to prefetch for offline use).

hf download digas99/smolvla_so101_screwdriver_box_v2 --local-dir ./smolvla_so101
python ./smolvla_so101/verify_policy.py ./smolvla_so101 cuda

It loads the checkpoint the way an inference script does, runs 15 dataset frames (3 per task, each with its own instruction) and compares the predicted 50-step chunk to the recorded actions. The numbers it prints should match the Results section below; errors near 1.0 (in units of action std) mean a contract or version problem.

Results (auto-appended after training)

Training loss (flow-matching MSE, every ~1k steps):

step	loss	lr
100	0.509	2.6e-06
1K	0.253	5.0e-05
11K	0.119	3.5e-05
21K	0.075	1.2e-05
30K	0.061	2.5e-06

Offline verification on 15 dataset frames (3 per task). Must reproduce on the deployment machine to the 3rd decimal:

over 15 frames from 15 episodes (up to 3 per task), errors in raw joint units (degrees):
  first-step  MAE per joint: [2.2173, 2.3868, 6.4581, 1.688, 0.9081, 0.5792]
  50-step     MAE per joint: [3.3643, 4.0599, 3.9291, 2.2502, 2.5286, 1.2263]
  first-step MAE / std (mean over joints): 0.074   50-step: 0.102
  finite outputs: True   latency per chunk on cuda: median 234.5 ms
  fingerprint (sum of first-step preds): 77.73454