A tiny vision-language-action model for one motor#

A one-joint motor must move to the object an instruction names on a one-dimensional table (“go to the red”, “go to the leftmost”, “go to the block”). TinyVLA grounds the instruction in the image with one attention step and a small action head turns the expected target into a motor command. When the instruction is true of two objects, a split-conformal gate makes the robot ask, and the answer is fused with the first belief. The data are simulated; this is a prototype, not a real robot.

import numpy as np

from quantum_mind.embodied import TinyVLA, bc_plan, calibrate_gate, make_dataset, oracle_plan, run_episode, vla_plan, BCBaseline

train, calib, test = make_dataset(2000, 0), make_dataset(800, 1), make_dataset(200, 2)
vla = TinyVLA(seed=0).fit(train, steps=2000)
print('TinyVLA parameters:', vla.n_params)
TinyVLA parameters: 817

Look at one ambiguous instruction#

s = next(s for s in test if s.ambiguous)
g = vla.ground(s.image, s.text)
print('instruction:', ' '.join(s.words), '| objects at cells', s.cells, '| intended cell', s.target)
print('belief over cells (top 3):', {int(i): round(float(g[i]), 2) for i in np.argsort(g)[::-1][:3]})
instruction: go to the block | objects at cells [1, 5] | intended cell 1
belief over cells (top 3): {5: 0.67, 1: 0.33, 8: 0.0}

Closed loop: act blindly, ask through a conformal gate, or use the end-to-end baseline#

gate = calibrate_gate(vla, calib, alpha=0.02)
bc = BCBaseline(hidden=128, seed=0).fit(train, steps=4000, lr=0.003)
print('%-22s %-9s %-12s %s' % ('policy', 'success', 'questions', '(ambiguous scenes only: success)'))
for name, plan in (('oracle', oracle_plan), ('BC baseline', bc_plan(bc)), ('TinyVLA, act', vla_plan(vla)),
                   ('TinyVLA, ask gate', vla_plan(vla, gate))):
    rng = np.random.default_rng(3)
    rows = [run_episode(sc, rng, plan) for sc in test]
    amb = [r['success'] for r, sc in zip(rows, test) if sc.ambiguous]
    print('%-22s %-9.3f %-12.2f %.3f' % (name, np.mean([r['success'] for r in rows]), np.mean([r['asked'] for r in rows]), np.mean(amb)))
policy                 success   questions    (ambiguous scenes only: success)
oracle                 1.000     0.00         1.000
BC baseline            0.200     0.00         0.029
TinyVLA, act           0.660     0.00         0.015
TinyVLA, ask gate      0.965     0.33         0.897

Total running time of the script: (0 minutes 9.338 seconds)