Flow-matching DiT policy (fm_dit) — R1 Lite "insert flower"
Baseline checkpoints for the ODE-policy comparison on the Galaxea R1 Lite right arm. 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 two cameras and proprioception.
Trained with exactly the same data, encoder, batch, learning rate and schedule as the ODE policy
(at237299966/ODE_policy_insert_flowers)
and the Diffusion Policy baseline (at237299966/dp_insert_flowers).
Files
| Path | Content |
|---|---|
checkpoints/{50,100,150,200,250,300}.ckpt |
Full training checkpoints (741 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 (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_insert_flower_single:
LeRobot v2.1, 95 episodes / 23,034 frames at 15 Hz, single task "insert flower".
| Input / output | Definition |
|---|---|
| Cameras | head_rgb (720×1280) and right_wrist_rgb (360×640), padded to square then resized to 256×256 (same geometry as the robot client) |
agent_pos |
right-arm joint positions (6, rad) + right gripper opening (1, mm) |
| Action | right-arm absolute joint targets (6, rad) + gripper command (1, mm; 1 = 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 + proprioception, 8 frames → global condition; DiT head 512 wide × 4 layers × 8 heads (23.93M parameters outside the 22.35M 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), 48 per GPU = 96 global; 236 steps/epoch, 70,800 steps in total |
| Wall clock | 9 h 00 min (108 s/epoch) |
| epoch | 0 | 10 | 50 | 100 | 150 | 200 | 250 | 300 |
|---|---|---|---|---|---|---|---|---|
| train loss | 0.83 | 0.066 | 0.026 | 0.017 | 0.012 | 0.0095 | 0.0082 | 0.0079 |
| val loss | 0.35 | 0.069 | 0.042 | 0.068 | 0.067 | 0.082 | 0.123 | 0.142 (ep 295) |
The validation loss (one held-out episode, raw weights) is lowest at epoch 85 (0.042) 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). 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
One prediction takes ~78 ms on an A100 (20 DiT calls), longer than the 66.7 ms tick; if the host GPU is not
faster, reduce --num_inference_steps. Code: ODE_Policy repository, neural_ode/real/ and
roboverse_learn/il/deploy/.