at237299966/fm_dit_microwave

Flow-matching DiT policy (fm_dit) — R1 Lite "microwave" (two arms)

Baseline checkpoints for the ODE-policy comparison on the Galaxea R1 Lite using both arms (the right arm puts a bowl into the microwave, the left arm works the door). fm_dit is a standard conditional flow-matching policy: a DiT action head denoises a 16-step action window (8 past + 8 future) from Gaussian noise, conditioned on 8 frames of three cameras and proprioception. Trained with exactly the same data, encoder, batch, learning rate and schedule as the ODE policy (at237299966/ODE_policy_microwave) and the Diffusion Policy baseline (at237299966/dp_microwave).

Files

Path	Content
checkpoints/{50,100,150,200,250,300}.ckpt	Full training checkpoints (954 MB each): model, ema_model, optimizer state dicts + the pickled Hydra config. Serve the EMA weights (the loader's default).
config/config.yaml, config/overrides.yaml	Resolved training config and the command-line overrides of this run
logs/logs.json.txt	Per-step / per-epoch training log (loss, lr, epoch time, val_loss every 5 epochs)
SHA256SUMS	Checksums of every uploaded file

Checkpoints are torch.save(..., pickle_module=dill) payloads written by the ODE_Policy training runner; loading needs that codebase at a version with two-arm support (branch r1lite_exp, commit e9cc9c0 or later: roboverse_learn.il.policies.fm.fm_dit_image_policy.FlowMatchingDiTImagePolicy) and torch.load(..., pickle_module=dill, weights_only=False) — only load checkpoints you trust.

Data

at237299966/r1lite_microwave: LeRobot v2.1, 34 episodes / 14,432 frames at 15 Hz, single task "microwave".

Input / output	Definition
Cameras	head_rgb (720×1280), right_wrist_rgb and left_wrist_rgb (360×640), padded to square then resized to 256×256 (same geometry as the robot client)
agent_pos	14-D, the robot client's order: [L1..L6 (rad), L_grip (mm), R1..R6 (rad), R_grip (mm)]
Action	14-D in the same order: absolute joint targets (rad) + gripper commands (mm; 0 = closed, 100 = open), as recorded
Window	8 observation frames; the policy predicts 16 steps and rows 7..14 (the current step and the next 7) are the action chunk

Training

Model	ResNet18 (GroupNorm, random crop) per camera (3 cameras) + proprioception, 8 frames → global condition; DiT head 512 wide × 4 layers × 8 heads (26.06M parameters outside the 33.53M encoder)
Objective	conditional flow matching (straight path from N(0, I)), 10 midpoint steps at inference
Optimisation	AdamW, lr 1.5e-4, cosine schedule over 300 epochs, warm-up 500 steps, EMA
Batch	2 × A100-40GB (DDP), 40 per GPU = 80 global; 175 steps/epoch, 52,500 steps in total
Wall clock	8 h 39 min (103 s/epoch)

epoch	0	10	50	100	150	200	250	300
train loss	1.14	0.086	0.036	0.021	0.015	0.012	0.0082	0.0077
val loss	0.50	0.086	0.037	0.055	0.058	0.078	0.078	0.110 (ep 295)

The validation loss (one held-out episode, raw weights) is lowest at epoch 50 (0.037) and rises afterwards, so compare early and late checkpoints on the robot. Loss values are not comparable across methods (the ODE and diffusion policies optimise different quantities).

Deployment

Same server and wire protocol as the ODE policy; the robot side is the unmodified openpi-r1lite closed-loop client (15 Hz, executes row 0 of each (8, 14) reply, must send all three cameras). For a like-for-like comparison with the ODE policy, predict a chunk every 8 ticks and walk through it:

python -m roboverse_learn.il.deploy.serve_policy \
    --checkpoint checkpoints/100.ckpt --host 0.0.0.0 --port 8000 --device cuda:0 --replan_every 8
python -m roboverse_learn.il.deploy.test_client --host 127.0.0.1 --port 8000 -n 40 --bimanual   # pre-flight

One prediction takes ~94 ms on an A100 (10 midpoint steps), longer than the 66.7 ms tick; with --replan_every 8 only one tick in eight is slow. If the host GPU is slower, reduce --num_inference_steps. Set the torso to the demonstrations' pose [-0.6845, 1.4986, 0.8031] rad before a run (it is not part of the policy's state). Code: ODE_Policy repository, neural_ode/real/ and roboverse_learn/il/deploy/.