at237299966/fm_dit_fold_towel

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

Baseline checkpoints for the ODE-policy comparison on the Galaxea R1 Lite folding a towel with both arms: the two grippers grasp the towel, fold it once, release, then grasp the folded towel again and drag it close to the robot. 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 the same data, encoder, batch, learning rate and schedule as the ODE policy with the multimodal side head (at237299966/ODE_policy_euler_sidehead_fold_towel).

Files

Path	Content
checkpoints/{50,100,150,200,250,300}.ckpt	Full training checkpoints: 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_fold_towel: LeRobot v2.1, 15 Hz, 35 episodes of 455–604 frames (30–40 s). 33 are used (17,280 frames): episodes 2 and 11 are dropped because in the drag phase the left gripper's first grasp did not hold and was repeated (selection in neural_ode/real/merges/r1lite_fold_towel_clean.txt).

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 or a few mm = closed on the towel, 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
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

Wall clock	9 h 54 min (119 s/epoch), 300 epochs, 208 steps/epoch, 62,400 steps in total

epoch	0	10	50	100	150	200	250	300
train loss	0.978	0.0888	0.0352	0.0214	0.0159	0.0122	0.00944	0.00856
val loss	0.377	0.0953	0.0595	0.0646	0.0747	0.108	0.12	0.143 (ep 295)

Lowest validation loss: 0.0489 at epoch 85 (validation every 5 epochs). Columns 50, 100, ... are the saved checkpoints (checkpoints/<epoch>.ckpt).

The validation loss (one held-out episode, raw weights) is a weak signal; compare early and late checkpoints on the robot. Loss values are not comparable across methods (the ODE and flow-matching 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). fm_dit outputs the gripper commands in mm directly, so no gripper flags are needed. 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. Start the robot from rest with the torso at the demonstrations' pose [-0.6944, 1.5135, 0.8066] rad (it is not part of the policy's state). Code: ODE_Policy repository, neural_ode/real/ and roboverse_learn/il/deploy/.