smolvla_so101_tray_v2box

Co-training experiment: the tray policy v2's data plus the single-object episodes of the earlier box project, to measure whether data from a related task on the same arm (different camera pose, a box instead of the tray, prompts naming the object, episodes from rest to rest) helps or hurts the tray task. Compare with digas99/smolvla_so101_tray_v2, trained on the tray data alone with the same recipe.

SmolVLA (450 M parameters) fully fine-tuned from lerobot/smolvla_base on digas99/so101_tray_train_v2box: a LeRobot SO-101 arm moving objects between a table and a transparent tray, in both directions. Two task strings, one per direction, the same for every object:

Pick up the object and put it in the tray (135 episodes, snapshot digas99/so101_tray_b1_pickA) Pick up the object from the tray and put it on the table (90 episodes, snapshot digas99/so101_tray_b2_placeB) plus the four box-project strings Pick up the <screwdriver|nut|battery|yellow tape> and put it in the box (370 episodes from digas99/so101_screwdriver_box, its multi-object episodes left out)

The policy is class-agnostic by design: it never hears an object name. Every episode starts with the gripper already hovering about 10 cm above the target with the jaws turned across it (a supervisor with a class-agnostic segmenter puts it there at run time), and ends right after the release with the gripper lifted a little above the tray. Nine household objects; table-to-tray episodes are placed over a 4 x 3 grid of the working area, tray-to-table episodes start from the reachable band of the tray, with touching neighbours, distractors and pre-filled trays in known proportions. Objects not in the training set are kept aside to measure generalisation.

Dataset: 595 episodes, 194,105 frames at 30 fps, two cameras (fixed top camera and a wrist camera on the gripper), LeRobot format v3.0, robot_type: so_follower. Every episode carries a provenance record (object, grid cell, adjacency, distractors, tray contents, hover position) in the dataset's meta/provenance.jsonl.

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)
Steps / batch size	40,000 / 16 (about 3.3 passes over the 194,105 frames)
Learning rate	5e-05, 1,000 warm-up steps, cosine decay to 1/20 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; the evaluation is on the robot, on ten fixed layouts, with training objects and with unseen ones
Checkpoints	every 2,500 steps under checkpoints/<step>/, so the step count can be chosen by that evaluation
Hardware	1x NVIDIA L4 (24 GB) on AWS g6.4xlarge, lerobot 0.6.1, see train_env_versions.txt

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

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 two strings above
action                      float32 (6,)  same joint order and units as observation.state

Joint values are raw SO-101 readings in degrees mode; normalisation statistics live inside the checkpoint. The policy expects to start from the hover described above; from the rest pose it has never seen the approach.

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

PATH = "digas99/smolvla_so101_tray_v2box"
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 object and put it in the tray"}   # or the tray-to-table string
# action = post(policy.select_action(pre(obs)))   # call at 30 Hz; the policy re-infers every 50 calls

Same environment as training: pip install "lerobot[smolvla]==0.6.1". verify_policy.py in this repo runs the checkpoint on dataset frames and prints the numbers of the Results section below; they should reproduce on the deployment machine.

Results (auto-appended after training)

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

step	loss	lr
100	0.508	2.6e-06
1K	0.254	5.0e-05
11K	0.132	4.1e-05
21K	0.087	2.4e-05
31K	0.064	7.8e-06
40K	0.055	2.5e-06

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

over 21 frames from 21 episodes (up to 3 per task), errors in raw joint units (degrees):
  first-step  MAE per joint: [1.4231, 2.2838, 5.3296, 1.533, 1.4722, 1.1155]
  50-step     MAE per joint: [2.4583, 3.1585, 3.4413, 2.3304, 1.76, 1.2634]
  first-step MAE / std (mean over joints): 0.075   50-step: 0.090
  finite outputs: True   latency per chunk on cuda: median 250.8 ms
  fingerprint (sum of first-step preds): 593.66479