so101_chess

A vision-language-action policy that does the physical half of chess on an SO-101 arm. Told pick up the piece on e2 and place it on e4 or take the piece on d5 off the board, it carries it out from two camera images and six joint angles. A chess engine picks the move; this policy moves the piece. Trained and evaluated in MuJoCo.

Showcase

Five successes, then five failures, per family - captioned with the instruction and the outcome.

On the workstation

What the robot sees - its two cameras side by side: overhead on the left, wrist on the right. These are the only images the policy gets.

If the players do not appear: workstation moves, workstation captures, robot's view of moves, robot's view of captures.

Use this model

Every command below works with the Hub id as it stands.

1. Install LeRobot with MolmoAct2 support (its main branch, Python 3.12), and the simulator the policy was trained in:

pip install "lerobot[molmoact2] @ git+https://github.com/huggingface/lerobot"
git clone https://github.com/qalby-tech/embryo_lab_chess_so101
pip install -e embryo_lab_chess_so101

2. Watch it play. One command: random positions, the policy driving, and a video of exactly what its two cameras see.

cd embryo_lab_chess_so101
MUJOCO_GL=egl python examples/evaluate_policy.py --checkpoint XvKuoMing/so101_chess \
    --moves 4 --captures 2 --video so101_chess.mp4 --video-cameras top wrist

3. Drive it from Python.

import random
from chess_sim import (ChessSimEnv, EnvConfig, LeRobotPolicy, LeRobotPolicyConfig,
                       MoveSampler, RolloutConfig, run_episode)

policy = LeRobotPolicy.load(LeRobotPolicyConfig(checkpoint="XvKuoMing/so101_chess"))
env = ChessSimEnv(EnvConfig())
task = MoveSampler().sample(env, random.Random(0))       # a position, and a move the arm can make
result = run_episode(env, policy, task, RolloutConfig(max_steps=450))
print(task.instruction, "->", "success" if result.success else result.reason)

Or call LeRobot directly, outside the simulator:

from lerobot.policies import make_pre_post_processors
from lerobot.policies.molmoact2.modeling_molmoact2 import MolmoAct2Policy

policy = MolmoAct2Policy.from_pretrained("XvKuoMing/so101_chess").to("cuda").eval()
pre, post = make_pre_post_processors(policy_cfg=policy.config, pretrained_path="XvKuoMing/so101_chess")

batch = {"task": ["pick up the piece on e2 and place it on e4"],
         "observation.state": joints_deg,                              # (1, 6)
         "observation.images.top": top, "observation.images.wrist": wrist}  # (1, 3, 480, 640), [0, 1]
action = post(policy.select_action(pre(batch)))                        # (1, 6)

4. Fine-tune it on recordings of your own:

lerobot-train --policy.path=XvKuoMing/so101_chess \
    --dataset.repo_id=YOUR/DATASET --dataset.video_backend=pyav \
    --steps=20000 --policy.scheduler_decay_steps=20000 \
    --batch_size=8 --policy.gradient_checkpointing=true \
    --policy.push_to_hub=false --output_dir=outputs/so101_chess_finetune

Set scheduler_decay_steps to the run length: MolmoAct2 otherwise decays its learning rate over a fixed 24,000 steps, whatever --steps says. Batch 8 with gradient checkpointing needs ~26 GB.

5. On a real SO-101 - not yet run on hardware:

lerobot-rollout --strategy.type=base --policy.path=XvKuoMing/so101_chess \
    --fps=10 --interpolation_multiplier=3 \
    --robot.type=so101_follower --robot.port=/dev/ttyACM0 \
    --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} }" \
    --task="pick up the piece on e2 and place it on e4" --duration=30
  • --fps=10 --interpolation_multiplier=3: the policy was trained on 10 Hz targets, and ramping the arm between them at 30 Hz is what stops it knocking pieces over.
  • The cameras must be named top and wrist, and placed as the policy saw them: the build guide has the rig, the 254 mm board and where each camera goes.
  • Joint angles are in the degrees-with-offset convention of the public SO-100/101 datasets. Check your follower's calibration reports the same before letting it move.
  • A model call takes 0.8-3.3 s on an RTX 5090. If the arm pauses between chunks, try --inference.type=rtc, which overlaps inference with motion - this checkpoint's continuous action mode supports it, but it is untested here.

Inputs and outputs

key shape what
task string the instruction, in the two forms above
observation.state (6,) joint angles, SO-101 degrees, shoulder pan first, gripper last
observation.images.top (3, 480, 640) overhead camera, RGB in [0, 1]
observation.images.wrist (3, 480, 640) wrist camera, RGB in [0, 1]
action (6,) joint targets, SO-101 degrees, one every 0.1 s

Results

Scored on 300 positions in simulation, 450 control steps each, with the board shifted up to 10 mm and the arm starting up to 0.1 rad off its parked pose, as in training:

as shipped (30 actions per call) 5 actions per call
moves (200) 159/200 (80%) 156/200 (78%)
captures (100) 68/100 (68%) 85/100 (85%)

Median placement error is about 5 mm. An earlier 96-episode evaluation read 84% on moves and 78% on captures, inside the intervals of these larger numbers.

The open-loop horizon is a capture setting. The model predicts 30 actions per call and, as shipped, executes all 30 - three seconds of motion acting on what the board looked like before it started moving. Executing five and re-planning is worth 17 points on captures (paired McNemar p = 0.002) and nothing on moves: a capture is the longer trajectory, out past the board edge and into a tray slot, and that is where blind motion fails. It costs six times the model calls, which is why the checkpoint ships 30:

policy.config.n_action_steps = 5   # then policy.reset(): the action queue is rebuilt from it

More flow-matching steps do not help. 20 or 40 integration steps per chunk score the same as the default 10 (82%, 82% and 80% on 100 paired positions) for two to four times the inference.

Evaluation noise. The policy samples its actions, so the same weights on the same positions do not give the same score: about a quarter of identical positions flip between runs. Compare checkpoints on the same positions, paired, and on a hundred or more.

Training

Fine-tuned from allenai/MolmoAct2-SO100_101 with LoRA (rank 64) on 6,095 scripted demonstrations (4,903 moves, 1,192 captures; 497,565 frames at 10 Hz), batch 8, 70,000 steps = 1.13 epochs, bf16 with gradient checkpointing, on one RTX 5090. Two cameras: overhead and wrist.

One detail mattered more than any other: MolmoAct2 decays its learning rate over a fixed 24,000 steps whatever --steps says. An earlier 57,000-step run spent everything past 24,000 at the floor and stopped improving at 47% on moves. Sizing scheduler_decay_steps to the run produced the curve above.

Every episode varies piece colours, lighting and piece size, shifts the board up to 10 mm on each axis and starts the arm up to 0.1 rad from its parked pose, so the policy has to find squares in the image rather than memorise joint angles.

docs/EXPERIMENTS.md in this repository carries the full record: every run, what failed, and the measurements behind each setting.

Data

The demonstrations come from a scripted expert in MuJoCo with a calibrated SO-101 model (28 mm squares, chosen because the arm reaches ~330 mm). Every episode that was recorded is verified: the piece ended where the instruction said, upright, with nothing else displaced.

XvKuoMing/so101_chess holds exactly the data this policy was trained on: 6,095 episodes, 497,565 frames at 10 Hz, 1,291 distinct instructions, two camera views per frame.

Limitations

  • Simulation only so far; no real-robot results yet.
  • Two instruction families. Naming a piece by type ("the white knight") needs recognition the policy has never been trained for.
  • A piece lying on its side cannot be picked up: the grasp assumes a standing cylinder.
  • 28 mm squares and the specific camera geometry are baked into every demonstration.

Roadmap

  • A continuation on rebalanced data is training now: captures were a fifth of the set and are the weaker family, so 1,982 more capture demonstrations and the first DAgger corrections (expert recoveries recorded from where this policy went wrong) were added.
  • Real SO-101 recordings, published alongside the simulated ones.
  • A sim-to-real showcase: the same instruction executed in simulation and on hardware.

Licensing follows the base model, allenai/MolmoAct2-SO100_101.

Downloads last month
67
Safetensors
Model size
6B params
Tensor type
F32
·
BF16
·
Video Preview
loading

Model tree for XvKuoMing/so101_chess

Finetuned
(16)
this model

Dataset used to train XvKuoMing/so101_chess

Collection including XvKuoMing/so101_chess