LegoCount Diffusion Policy — original and big blocks
Public EMA policies for placing exactly 1, 2, or 3 Lego blocks into a bin and returning home. Inputs are front/wrist RGB, measured robot state, and requested count. Policies predict 50 actions and execute 20 at 20 Hz. Object poses and deposit labels are not policy inputs.
Models and results
Best checkpoints verified on 10 October 2026. Scores are full-task success on the repeated evaluation set used to select checkpoints, not held-out generalization or official challenge scores. A later update is retained only when success improves.
| Dataset | Policy | Update | Selection success | Weights |
|---|---|---|---|---|
| Original · 2,900 demos | DP · long action chunk | 160,000 | 7/12 (58.3%) | best.pt |
| Original · 2,900 demos | DP + LSTM | 120,000 | 11/12 (91.7%) | best.pt |
| Original · 2,900 demos | DP + kinematic traces | 280,000 | 10/12 (83.3%) | best.pt |
| Simplified · 1,000 demos | DP · long action chunk | 140,000 | 9/12 (75.0%) | best.pt |
| Simplified · 1,000 demos | DP + LSTM | 80,000 | 6/12 (50.0%) | best.pt |
| Simplified · 1,000 demos | DP + kinematic traces | 100,000 | 11/12 (91.7%) | best.pt |
| Original · 900 demos | DP · long action chunk | 29,304 | 2/6 (33.3%) | best.pt |
| Original · 900 demos | DP + LSTM | 40,348 | 3/6 (50.0%) | best.pt |
| Original · 900 demos | DP + kinematic traces | 27,431 | 0/6 (0.0%) | best.pt |
The original 900-demo dataset has 90 separate validation demonstrations. Original 2,900 adds 2,000 diverse demonstrations from the standard planner; it excludes the slower-expert and simplified datasets. Simplified 1,000 uses larger blocks, a wider opening, fixed initial yaw/robot pose, and 10 cm block spacing in LegoCountEasy-v2. Results from different tasks or evaluation sets are not directly comparable. This repository contains DP, DP LSTM, and DP with kinematic traces.
Loading
Each best.pt contains ema, normalization stats, public config, and step. These inference exports omit optimizer/RNG states and private training paths; they cannot resume training. EMA tensors and normalization are verified against the selected training checkpoints. Files can be read with torch.load(path, map_location="cpu", weights_only=True).
In the solution checkout, run python -m other.download_checkpoints --group original2900 --variant trace. No login is required. Then use python -m policies.dp.evaluate as described in its README. Groups: original900, original2900, easy1000; supported variants: default, lstm, trace.
The LSTM visual-history encoder is assets/resnet18-f37072fd.pth; put it under $TORCH_HOME/hub/checkpoints/. Each policy directory contains public model settings and per-episode selection results in metadata.json. checkpoint_manifest.json provides sizes and SHA256 digests.