Flow-matching DiT policy (fm_dit) — R1 Lite "insert two flowers"
Baseline checkpoints for the ODE-policy comparison on the Galaxea R1 Lite right arm inserting flowers
(long-horizon task r1lite_insert_two_flowers). 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 the same data, encoder, batch, learning rate and
schedule as the ODE policy on this task
(at237299966/ODE_policy_euler_insert_two_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 (branch r1lite_exp;
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:
LeRobot v2.1, 98 episodes / 45,684 frames at 15 Hz, task "insert flower" (two flowers per episode).
| 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; 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 + 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; 470 steps/epoch, 141,000 steps in total |
| Wall clock | 17 h 52 min (214 s/epoch) |
| epoch | 0 | 10 | 50 | 100 | 150 | 200 | 250 | 300 |
|---|---|---|---|---|---|---|---|---|
| train loss | 0.55 | 0.044 | 0.018 | 0.012 | 0.0095 | 0.0078 | 0.0056 | 0.0055 |
| val loss | 0.22 | 0.046 | 0.061 | 0.079 | 0.061 | 0.112 | 0.130 | 0.089 (ep 295) |
The validation loss (two held-out episodes, raw weights) is lowest at epoch 30 (0.027) 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; the left arm is echoed, i.e. held). 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/50.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 # pre-flight
One prediction takes ~78 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.7745, 1.4887, 0.6998] 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/.