andreu-collabs/hybrid-flow-planning-3blocks

Hybrid flow planner, 3-blocks task: production checkpoints (2026-09-04)

One transformer jointly generates a discrete plan (masked discrete flow) and a continuous end-effector trajectory (continuous flow), conditioned on one scene observation at an action boundary. Task: a Dingo base with a Kinova Lite arm stacks three cubes on the square table, from 121 real demonstrations. Code: the hybrid-flow-planning repository (examples_hybrid/hybrid_planning_3_blocks_ee_chunks.py, src/hybrid_flow_planning/); the options behind every checkpoint are documented in the repository's ablations.md.

Three policies, one folder each. Every number is on the 12 held-out demonstrations (plan lengths 4 to 18, 142 chunk starts, 50 samples each), with the deployment decoding options (--decode-grammar --freeze-after-eos); the three-seed means are in brackets.

folder	policy	training flags	plan valid	plan valid + proxy
baseline/	fixed tower, no extras, seed 42	none	0.70 [0.68]	0.35 [0.34]
baseline-recipe/	fixed tower, best ablations, seed 43	--pad-loss-weight 0.1 --loss-weighting uncertainty	0.77 [0.76]	0.37 [0.35]
perm-aug-recipe/	any tower order (goal-conditioned), best ablations, seed 42	--perm-aug --pad-loss-weight 0.1 --loss-weighting uncertainty	0.75 [0.76]	0.26 [0.24]

"Plan valid" = the predicted plan is the demonstration's plan or an equivalent plan under the symbolic world model; "+ proxy" = the kinematic MuJoCo proxy also executes the plan with the predicted waypoints to the finished tower. With the original sampler (no decoding options) the baseline scores 0.56 / 0.27; the decoding options and the recipe are described below. The goal-conditioned policy scores 0.76 / 0.25 over all six tower orders.

Deployment checkpoints trained on all 121 demonstrations

The three folders above hold out 12 demonstrations, which is what makes their numbers meaningful. For deployment, the same three configurations were also trained on every demonstration (--val-episodes "", seed 42): baseline-alltrain/, baseline-recipe-alltrain/, perm-aug-recipe-alltrain/. They cannot be scored on unseen demonstrations; on the twelve demonstrations they now train on (eval_grammar_freeze_twelve/) they score 1.00 / 0.94, 1.00 / 0.86 and 1.00 / 0.83 (plan valid / plan valid + proxy), the in-distribution behaviour. The held-out size study in the repository found no measurable difference between training on 106, 113 and 118 demonstrations, so the expected gain of the extra 12 on new scenes is small; the extra coverage of scenes is the reason to prefer them on the robot. Same files, same load commands, same deployment_decoding.

The recipe baseline on the twelve held-out demonstrations: demonstration left, policy right, the current plan step highlighted (green = executed to success, red = failed).

What the recipe is

• Pad supervision (--pad-loss-weight 0.1): the plan positions after EOS are masked and supervised like every other position, so the generated canvas has no out-of-distribution tail.
• Learned loss weighting (--loss-weighting uncertainty): the continuous and the discrete flow-matching losses are combined with learned log-variances (Kendall et al., 2018) instead of a hand-tuned weight.
• Decoding at inference (--decode-grammar --freeze-after-eos, also stored in manifest.json as deployment_decoding): every revealed token is the most likely token among those the plan grammar allows at its position, an illegal visible token is re-masked and re-predicted, and no position after the first visible EOS is revealed. Every sampled plan is grammatical.
• Permutation policy (--perm-aug): trained on six colour-permuted copies of every demonstration with a per-cube goal field (support one-hot + rank in the tower) and one memory token per cube; it builds any tower order given in the goal fields.

Files per folder

• config.json: the ExperimentConfig of the run and its provenance.
• step_20000_ema.pt: the EMA weights at step 20000.
• manifest.json: everything needed to run the policy without the dataset: observation and continuous layouts, per-field normalization, vocabulary (pad=0, EOS=20, mask=21), max_t/max_d, model_kwargs and trainer_kwargs (exact constructor arguments), and deployment_decoding (the prediction keyword options to use).
• eval/ and eval_grammar_freeze/: the 50-seed evaluations with the original sampler and with the deployment decoding (summary.md, summary.json, plots, the stacked video where rendered). perm-aug-recipe/eval_grammar_freeze_allperms/: the same evaluation over all six tower orders.

Load a checkpoint with the repository tooling

A downloaded folder is a valid run folder (config.json + step_<n>_ema.pt). From the repository root:

hf download andreu-collabs/hybrid-flow-planning-3blocks --local-dir runs/hf
uv run python examples_hybrid/evaluate_hybrid_policy.py --run-dir runs/hf/baseline-recipe --num-seeds 10 --decode-grammar --freeze-after-eos
uv run python examples_hybrid/evaluate_hybrid_policy.py --run-dir runs/hf/perm-aug-recipe --num-seeds 10 --decode-grammar --freeze-after-eos --perms all

The first evaluation regenerates the packed HDF5 from the raw demos in the repository if it is missing (3_blocks_ee_chunks.hdf5; 3_blocks_ee_chunks_perm_slotgoal.hdf5 for the permutation policy). In Python:

from hybrid_flow_planning.evaluation.runner import EvalConfig, load_policy

policy = load_policy(EvalConfig(run_dir="runs/hf/baseline-recipe"))
sampler = policy.sampler
obs = sampler.all_observations[0:1].unsqueeze(1).to(policy.device)          # (B, 1, d_obs), normalized
traj, tokens = policy.trainer.prediction(obs=obs, batch_size=1, num_steps=100,
                                         decode_grammar=True, freeze_after_eos=True)
traj = sampler.denormalize_cont(traj)   # (B, 1, max_t, 14): base xy, yaw sin/cos, ee pos, ee rot6d, gripper

Run the policy on the robot without the dataset

import json, torch
from hybrid_flow_planning.data_samplers import HybridTransform
from hybrid_flow_planning.models import HybridConditionalVectorField
from hybrid_flow_planning.trainers import HybridFlowTrainer

folder = "runs/hf/baseline-recipe"
m = json.load(open(f"{folder}/manifest.json"))
transform = HybridTransform.from_manifest(m)                 # normalize_obs / denormalize_cont, layouts, tokens
kwargs = dict(m["model_kwargs"], img_size=tuple(m["model_kwargs"]["img_size"]),
              patch_size=tuple(m["model_kwargs"]["patch_size"]))   # JSON lists -> tuples
model = HybridConditionalVectorField(**kwargs)
model.load_state_dict(torch.load(f"{folder}/step_20000_ema.pt", map_location="cpu"))
model.eval()
trainer = HybridFlowTrainer(sampler=transform, **m["trainer_kwargs"])
trainer.model = model

obs_phys = torch.zeros(1, 1, m["meta"]["d_obs"])             # one observation in the layout of m["meta"]["obs_layout"]
obs = transform.normalize_obs(obs_phys)
traj, tokens = trainer.prediction(obs=obs, batch_size=1, num_steps=100, **m["deployment_decoding"])
traj = transform.denormalize_cont(traj)

The observation is the concatenation of the fields in m["meta"]["obs_layout"], in that order, in physical units (metres, yaw as sin/cos, rot6d, gripper bit 1 = closed). For perm-aug-recipe the three <colour>_goal fields hold the goal: for each cube a one-hot over its support in the target tower (square table, red, yellow, green) and a one-hot over its rank (bottom, middle, top); the standard tower is red [0,0,1,0, 0,0,1], yellow [0,0,0,1, 0,1,0], green [1,0,0,0, 1,0,0]. The predicted trajectory has max_t = 69 waypoints, about 0.8 s apart, and repeats the last state after the end; the plan tokens are Move_to RGk | Pick X | Place X Y triples ending with EOS (m["meta"]["vocabulary"] gives the names).