longnav-mshab-pick-flow-native-speedmix-20hz

longnav-mshab-pick-flow-native-speedmix-20hz

A dense-fine-tuned Qwen3-VL-2B trunk with a grouped flow-matching action head that drives the Fetch robot in ManiSkill-HAB Pick (set_table / prepare_groceries / tidy_house objects) from two RGB cameras and proprioception, emitting native joint-rate commands at 20 Hz. Trained by behaviour cloning on MS-HAB expert demonstrations replayed at a mixture of speeds (constant 0.5x plus per-tick schedules in [0.2x, 0.5x]), 48,000 steps, EMA weights. Headline: on 252 held-out set_table val episodes at 2.5 Hz decisions (8 of 20 ticks executed), success-once 0.575 and success-at-end 0.313, level with a pi0.5 SFT baseline on the same benchmark (0.563 / 0.222) and with a per-tick validation flow loss of 0.099 at the first tick. This is the initialisation of our MS-HAB RL runs.

What this is

Not a discrete-action agent and not an end-effector policy. Each decision produces a chunk of 20 native Fetch commands at 0.05 s spacing (1.0 s of motion): 7 arm joint rates, torso rate, absolute gripper target, 2 head rates, base forward speed and base yaw rate, 13 dims, under a certified action contract (physical_rates_v1) that fixes their physical units and the controller they assume (pd_joint_delta_pos, controller bounds scaled to 0.5x: arm/body delta 0.05 rad per tick, base 0.5 m/s and 1.57 rad/s). The deployment loop executes the first execute_ticks rows of every chunk (we use 8, i.e. 2.5 Hz decisions) and re-plans.

Architecture: Qwen3-VL-2B trunk (all 2.13B parameters fine-tuned) reads the prompt, a 512x320 head image, a 224x224 wrist image and 25-dim proprioception; two readout tokens are projected (per-token LayerNorm -> Linear -> GELU -> Linear, 2048 wide) into 8 context tokens of width 512 that condition a 6-layer, 512-wide transformer flow decoder (28M parameters) with separate arm and base token groups, a continuous sinusoidal time embedding in seconds (periods 0.05-8 s) and a group-type embedding. Actions are sampled with a 10-step Euler flow integration. Proprioception enters the trunk only, never the head.

Results

Per-tick validation flow-matching loss (512 held-out rows, K=4 flow times, fixed seeds; the cross-run reference protocol; lower is better):

Closed loop, set_table/pick/val/all.json (apple and bowl plans), val spawns, 63 held-out scene configurations x 4 seeds = 252 episodes, 200 steps, exec-8 (2.5 Hz decisions), deterministic flow sampling (10 Euler steps), measured through longnav.scripts.train_native_rl run.mode=calibrate with the MS-HAB success definition (grasped and end-effector at rest and robot at rest and static and contact force under the episode limit). Success-once = any tick; success-at-end = tick 200.

Binomial standard error at n=252 is about 0.03. The pi0.5 rows are a matched protocol, not a paired comparison (seeds differ, and their decisions are 4 of 8 ticks at 5 Hz). Success-at-end is the hold phase; the training data ends 5 ticks after first success, so post-success behaviour is untrained. Stochastic (flow-SDE) sampling at noise amplitudes 0.2-0.7 keeps success-once in 0.56-0.64 on the same episodes.

Training mixture

Snapshot native_pick_speedmix_snap13_20260918, sampled object-uniform, then task, then episode, then tick. Expert demonstrations come from MS-HAB's released RL experts replayed with a patched controller; the two components share no episodes.

Realised split: 26,945 train episodes (2,266,973 frames), 2,055 validation episodes (177,133 frames) with validation held out by scene configuration. Objects: apple, bowl (set_table); master chef can, cracker box, sugar box, tomato soup can, tuna fish can, pudding box, gelatin box, potted meat can (prepare_groceries); tidy_house objects. Episodes are cut 5 ticks after the first success.

Optimisation: warm start of the trunk from our constant-speed wide-head run (native_qwen_h20_wide_cont48k_20260916/checkpoint-48000) with a fresh grouped head; head trained alone for 1,000 steps, trunk unfrozen over a 1,000-step ramp; AdamW lr 2.5e-5 with 1,000 warmup steps, effective batch 32 (4 GPUs x 8 x 4 accumulation), 8 stratified flow-time samples per trunk forward, EMA 0.99, 48,000 steps. Full curve in trainer_state.json (wandb vlm_rl/mshab-native-qwen/3vo3nu6d).

Files

Optimizer and RNG state are excluded. The local checkpoint stores fp32 master weights plus a separate EMA copy (16 GB); this export folds the EMA in and stores bf16. Verified: identical per-tick validation loss to four decimals (0.0988 / 0.1246 / 0.1434 on both), and the policy runs under bf16 autocast in training and evaluation anyway.

What it depends on

• Code: github.com/averyxi42/spatial_training, branch mshab, commit fe2381a or later (longnav.utils.native_mshab.NativeFlowPolicy, longnav.env.fetch_patches, longnav.env.mshab_vector, tools/eval_native_qwen.py). The loader is NativeFlowPolicy.load(path, device, ema=True); with no ema.pt present it loads these weights as they are.
• Simulator: ManiSkill fork haosulab/ManiSkill branch mshab at 17121e3 and arth-shukla/mshab at e9ff3d2 with the small recording-wrapper patch in setup/pins/mshab-record.patch; MS-HAB ReplicaCAD scene data, set_table task plans and spawn data (standard MS-HAB downloads) under MS_ASSET_DIR.
• Two conda environments (simulator and policy cannot share an interpreter); pins in setup/pins/*.requirements.txt. torch 2.8.0+cu128, transformers 4.57.6.
• The half-speed controller patch and the two VLA cameras (512x320 head at 110 deg hfov, 224x224 wrist) are applied in-process by longnav.env.fetch_patches and must match what the data was recorded with; the contract certificate in native_config.json is checked at load.
• Not needed for evaluation: the expert demonstration generator (colleague's baseline repo), which only the data collectors call.

Running / evaluating it

Closed-loop evaluation with videos and per-episode traces (runs in the simulator env; the policy is loaded in-process, see the script header for the sys.path it needs):

env -u DISPLAY CUDA_VISIBLE_DEVICES=0 MS_ASSET_DIR=$PWD/data_mshab PYTHONPATH=$PWD/src HF_HUB_OFFLINE=1 \
python tools/eval_native_qwen.py --checkpoint <this repo, downloaded> --output-dir dump/eval_speedmix \
  --split val --task set_table --obj all --max-build-configs 63 --build-config-stride 1 --num-rollouts 4 \
  --max-steps 200 --speed-scale 0.5 --predict-ticks 20 --execute-ticks 8 --flow-steps 10 --seed 10000

--execute-ticks changes what the numbers mean (2 = 10 Hz ceiling, 4 = 5 Hz, 8 = 2.5 Hz; the table above is 8). --speed-scale 0.5 must match the contract's deployment_controller_scale.

The 252-episode table above came from the batched vector environment:

env -u DISPLAY HF_HUB_OFFLINE=1 PYTHONPATH=src:. python -m longnav.scripts.train_native_rl \
  run.mode=calibrate run.init_checkpoint=<this repo, downloaded> rollout.execute_ticks=8 run.eval_episodes=63 \
  sde.gap_noise_grid=[0.2,0.35,0.5,0.7] run.policy_gpus=[0,1,2,3] rollout.sim_gpus=[4,5,6,7] run.wandb_project=""

which writes per-episode arrays to <run>/calibration/*.npz (the val_ode condition is the row reported here).