captainjaseel/smolvla-cube-8dim-v2

SmolVLA — ElRobot Cube Pick-and-Place (8-DoF, v2)

Fine-tuned SmolVLA policy for an 8-DoF ElRobot arm doing the task "put the black object inside the green box".

This is v2 of the model, trained with a ~4× faster pipeline (BF16 autocast + SDPA attention + torch.compile) — full training in ~2 h on a single RTX A6000. The training math is bit-identical to a stock FP32 run; only the wall-clock time changed.

Files

• config.json — SmolVLAConfig with state_dim=8, action_dim=8.
• model.safetensors — 450 M-parameter policy (99 M trainable). FP32 master weights.
• stats.safetensors — per-joint state/action mean & std for normalization at inference time. Required — without it actions are denormalized incorrectly.

Training

• Base model: lerobot/smolvla_base
• Trainable: 99.88 M / 450.05 M params (22.2 %)
• Dataset: 17 episodes, 6809 frames of cube pick-and-place demonstrations on the ElRobot.
• Optimizer: Fused AdamW, lr 1e-4, cosine decay over 15000 steps, warmup 1000 steps, betas (0.9, 0.95), wd 1e-10, grad clip 10.
• Batch: 48 (RTX A6000, 48 GB).
• Steps: 15 000 (~106 epochs).
• Precision: BF16 autocast in the forward pass; FP32 master weights / optimizer state / loss.
• Attention: PyTorch SDPA (scaled_dot_product_attention) replacing the upstream eager FP32 attention path.
• Final loss: 0.0141 (flow-matching MSE).

Speed / quality

Stack	step/s	Wall time @ 15000 steps
FP32 baseline (upstream)	0.52	7 h 51 m
+ BF16 autocast + TF32 + fused AdamW	1.42	2 h 56 m
+ SDPA attention	1.60	2 h 36 m
+ torch.compile(mode="default") (this run)	2.13	~2 h

Loss curve is bit-identical to the FP32 baseline at every step we measured — the speed wins do not change the trained model's behavior. The relevant PR upstreaming these changes (without the robot-specific 8-DoF hard-code) is at norma-core/norma-core.

How to use

from smolvla import SmolVLAPolicy
from safetensors.torch import load_file as load_safetensors

policy = SmolVLAPolicy.from_pretrained("captainjaseel/smolvla-cube-8dim-v2", strict=False).cuda().eval()
stats  = load_safetensors("stats.safetensors")  # download alongside the model

# build a batch (see Robotics_berlin_hack/sim/mujoco_eval.py for the full pattern)
# normalize state with stats["state_mean"], stats["state_std"]
# call policy.predict_action_chunk(batch) -> (B, chunk_size, 8) normalized actions
# un-normalize with stats["action_mean"], stats["action_std"]

The smolvla Python package lives in norma-core/software/ai/smolvla_py.

Closed-loop MuJoCo demo

A self-contained MuJoCo inference script that loads this checkpoint, builds an ElRobot + table + cube scene, and runs the policy in a closed loop lives at Robotics_berlin_hack/sim/.

License

Apache 2.0, matching the upstream SmolVLA license.