| --- |
| tags: |
| - robotics |
| - diffusion |
| - policy |
| datasets: |
| - villekuosmanen/build_block_tower |
| - villekuosmanen/dAgger_build_block_tower_1.0.0 |
| - villekuosmanen/dAgger_build_block_tower_1.1.0 |
| - villekuosmanen/dAgger_build_block_tower_1.2.0 |
| - villekuosmanen/dAgger_build_block_tower_1.3.0 |
| - villekuosmanen/dAgger_build_block_tower_1.4.0 |
| --- |
| |
| # Multi-Task DiT Policy — Block Tower (Config Fix) |
|
|
| Diffusion Transformer (DiT) policy trained on [villekuosmanen/build_block_tower](https://huggingface.co/datasets/villekuosmanen/build_block_tower) + DAgger rounds 1.0.0–1.4.0 for robotic block stacking. Training config (batch size, learning rate, resize, horizon) recommended by the repo author based on what worked for him. |
|
|
| ## Training Details |
|
|
| | Parameter | Value | |
| |---|---| |
| | Architecture | DiT with CLIP ViT-B/16 vision encoder + CLIP text conditioning | |
| | Dataset | build_block_tower + 5 DAgger rounds | |
| | State/Action dim | 17D — joint_pos(7) + eef_xyz(3) + rot6d(6) + gripper(1) | |
| | Delta actions | All dims except 6D rotation (absolute) | |
| | Normalization | Ramen (q02/q98 percentile, per-timestep, per-dim, clipped [-1.5, 1.5]); 6D rotation exempt | |
| | Batch size | 80 per GPU, 320 global (4x GPUs) | |
| | Training steps | 40,000 / 50,000 (in progress) | |
| | Learning rate | 3e-4, cosine schedule, 500 warmup steps | |
| | Diffusion | DDIM, 100 train timesteps, 20 inference steps | |
| | Horizon | 32 | |
| | Action steps | 32 | |
| | Obs steps | 2 | |
| | Vision resize | 224x224, no crop | |
| | Mixed precision | AMP | |
| | Optimizer | Adam, grad clip 1.0 | |
| | Hardware | 1 node, 4x NVIDIA GH200 (Isambard-AI AIP2) | |
| | Training time | ~42h across 2 runs (24h + 18h resume) | |
| | Final loss | ~0.008–0.012 | |
|
|
| ## Checkpoints |
|
|
| | Checkpoint | Steps | sha256 (model.safetensors) | |
| |---|---|---| |
| | `checkpoint_40000` | 40k | `455f0f6f...1f7c29` | |
|
|
| Each checkpoint contains: |
| - `model.safetensors` — model weights (~1.3GB) |
| - `config.json` — model configuration |
| - `ramen_stats.pt` — normalization statistics (required for inference) |
|
|
| ## Task |
|
|
| Stack a coloured wooden block on top of an existing block tower. |
|
|
| ## W&B |
|
|
| - [Run 1 (steps 0–29k)](https://wandb.ai/pravsels/dit_block_tower_config_fix/runs/c00fb0ai) |
| - [Run 2 (steps 20k–40k, resume)](https://wandb.ai/pravsels/dit_block_tower_config_fix/runs/a8ql2mse) |
|
|
| ## Usage |
|
|
| ```python |
| from multitask_dit_policy.model import MultiTaskDiTPolicy |
| |
| policy = MultiTaskDiTPolicy.load("pravsels/dit_block_tower_config_fix/checkpoint_40000") |
| ``` |
|
|
