Upload 100000 checkpoint

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+media/rollouts.gif filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,86 @@
+---
+license: apache-2.0
+library_name: lerobot
+pipeline_tag: robotics
+tags:
+  - lerobot
+  - diffusion-policy
+  - robotics
+  - bimanual
+  - mujoco
+  - imitation-learning
+---
+
+# Bimanual puzzle pushing: Diffusion Policy
+
+This is the policy I trained for my [puzzleBimanual](https://github.com/HarunSMetin/puzzleBimanual)
+project. Two Franka Panda arms with stick end-effectors stand on either side of a table in
+MuJoCo. Each arm pushes one half of a randomly cut puzzle, and the goal is to push the two halves
+together until they fit into one rectangle. There is no target marker, so the policy has to work
+out from the shapes alone how the pieces go together.
+
+<p align="center">
+  <img src="media/rollouts.gif" width="640" alt="Two rollouts on unseen puzzles">
+  <br><em>Two rollouts on test puzzles. Each row shows the top camera and both wrist cameras.</em>
+</p>
+
+## Results
+
+I tested the policy in the simulator on puzzles that were never used in training. An episode
+counts as a success when the pieces are within 5 mm and 4° of the mated pose and stay there for
+half a second, which is the same rule I used when recording the demonstrations.
+
+| Split | Puzzles | Success | 95% CI |
+|---|---|---|---|
+| dev (used to pick this checkpoint) | 50 | 88% | 76–94% |
+| test (evaluated once) | 200 | **82%** | 76–87% |
+
+When the policy fails, it's mostly because the pieces never got close: in 30 of the 36 failed
+test episodes they never came within 2 cm of the mated pose. Getting stuck on the final
+alignment is rare (3 episodes). Once the pieces are
+close, it usually finishes: the median final error in successful episodes is about 2 mm.
+
+## What the policy sees and does
+
+- **Inputs**: three 240×320 RGB cameras (a top view and one camera on each wrist) and an
+  18-number robot state (tip x, y and seven joint angles per arm), with the last 2 frames as
+  context. The piece poses are not given to the policy; it has to see the pieces.
+- **Output**: absolute x, y targets for the two stick tips, 4 numbers per step at 20 Hz. The
+  policy predicts 16 steps and executes 8 (0.4 s) before planning again.
+- **Model**: LeRobot's Diffusion Policy with a separate ResNet18 encoder per camera and a 1D
+  U-Net with channels [256, 512, 1024]. 103M parameters in total. DDIM with 10 denoising
+  steps at inference.
+
+## Training data and setup
+
+I recorded 217 demonstrations myself with a DualShock 4 (one stick per arm), about 50 minutes
+in total. A processing step removed the pauses in each arm's motion and added mirrored and
+rotated copies of every demonstration, which gave 797 episodes. Every episode has a different
+random puzzle.
+
+Training ran for 100k steps with batch size 32 and mixed precision on a laptop RTX 4050 (6 GB),
+which took about 12 hours. I compared the 50k, 70k and 100k checkpoints on the same 50 dev
+puzzles (64%, 70% and 88%) and kept the last one.
+
+## How to use it
+
+The simulator and evaluation code live in the GitHub repository. You don't need the dataset;
+the files here include everything the evaluation needs (`eval_profile.json` holds the camera
+setup and the training seeds, so evaluation puzzles never overlap with training).
+
+```bash
+git clone https://github.com/HarunSMetin/puzzleBimanual.git
+cd puzzleBimanual
+uv sync
+uv run python scripts/eval_policy.py --checkpoint HSM1/puzzle-bimanual-diffusion --episodes 20
+```
+
+It runs on a GPU if there is one and on the CPU otherwise (much slower). Results and videos
+go to `outputs/eval/`.
+
+## Limitations
+
+This policy only knows this simulated scene: the table, the arm placement, the camera positions
+and the puzzle generator's shapes. It has not been tried on a real robot. Rollouts are also not
+exactly repeatable, because tiny rendering differences between runs grow into different pushes,
+so compare success rates over many puzzles rather than single episodes.

config.json

@@ -0,0 +1,105 @@
+{
+    "type": "diffusion",
+    "n_obs_steps": 2,
+    "input_features": {
+        "observation.state": {
+            "type": "STATE",
+            "shape": [
+                18
+            ]
+        },
+        "observation.images.top": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                240,
+                320
+            ]
+        },
+        "observation.images.left_wrist": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                240,
+                320
+            ]
+        },
+        "observation.images.right_wrist": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                240,
+                320
+            ]
+        }
+    },
+    "output_features": {
+        "action": {
+            "type": "ACTION",
+            "shape": [
+                4
+            ]
+        }
+    },
+    "device": "cuda",
+    "use_amp": true,
+    "use_peft": false,
+    "push_to_hub": false,
+    "repo_id": null,
+    "private": null,
+    "tags": null,
+    "license": null,
+    "pretrained_path": null,
+    "pretrained_revision": null,
+    "horizon": 16,
+    "n_action_steps": 8,
+    "normalization_mapping": {
+        "VISUAL": "MEAN_STD",
+        "STATE": "MIN_MAX",
+        "ACTION": "MIN_MAX"
+    },
+    "drop_n_last_frames": 7,
+    "vision_backbone": "resnet18",
+    "resize_shape": null,
+    "crop_ratio": 1.0,
+    "crop_shape": [
+        216,
+        288
+    ],
+    "crop_is_random": true,
+    "pretrained_backbone_weights": "ResNet18_Weights.IMAGENET1K_V1",
+    "use_group_norm": false,
+    "spatial_softmax_num_keypoints": 32,
+    "use_separate_rgb_encoder_per_camera": true,
+    "down_dims": [
+        256,
+        512,
+        1024
+    ],
+    "kernel_size": 5,
+    "n_groups": 8,
+    "diffusion_step_embed_dim": 128,
+    "use_film_scale_modulation": true,
+    "gradient_checkpointing": false,
+    "noise_scheduler_type": "DDIM",
+    "num_train_timesteps": 100,
+    "beta_schedule": "squaredcos_cap_v2",
+    "beta_start": 0.0001,
+    "beta_end": 0.02,
+    "prediction_type": "epsilon",
+    "clip_sample": true,
+    "clip_sample_range": 1.0,
+    "num_inference_steps": 10,
+    "compile_model": false,
+    "compile_mode": "reduce-overhead",
+    "do_mask_loss_for_padding": false,
+    "optimizer_lr": 0.0001,
+    "optimizer_betas": [
+        0.95,
+        0.999
+    ],
+    "optimizer_eps": 1e-08,
+    "optimizer_weight_decay": 1e-06,
+    "scheduler_name": "cosine",
+    "scheduler_warmup_steps": 500
+}

eval_profile.json

@@ -0,0 +1,286 @@
+{
+ "profile": {
+  "task": "Push the two puzzle pieces together so they fit into one rectangle.",
+  "image_size": [
+   240,
+   320
+  ],
+  "cameras": {
+   "top": {
+    "pos": [
+     0.0,
+     0.0,
+     0.85
+    ],
+    "lookat": [
+     0.0,
+     0.0,
+     0.0
+    ],
+    "up": [
+     0.0,
+     1.0,
+     0.0
+    ],
+    "fovy": 45
+   },
+   "left_wrist": {
+    "body": "left_hand",
+    "pos": [
+     -0.08,
+     0.0,
+     0.08
+    ],
+    "lookat": [
+     0.1,
+     0.0,
+     0.2
+    ],
+    "up": [
+     0.0,
+     0.0,
+     -1.0
+    ],
+    "fovy": 90
+   },
+   "right_wrist": {
+    "body": "right_hand",
+    "pos": [
+     -0.08,
+     0.0,
+     0.08
+    ],
+    "lookat": [
+     0.1,
+     0.0,
+     0.2
+    ],
+    "up": [
+     0.0,
+     0.0,
+     -1.0
+    ],
+    "fovy": 90
+   }
+  }
+ },
+ "training_seeds": [
+  12028509,
+  20601358,
+  23801677,
+  25683142,
+  25692361,
+  43779169,
+  66429433,
+  108900320,
+  124137339,
+  126500499,
+  131099074,
+  132355821,
+  145108040,
+  146520148,
+  160963281,
+  162782454,
+  168038232,
+  179966057,
+  188119339,
+  227450912,
+  241231647,
+  243663250,
+  253013083,
+  258127594,
+  263002108,
+  283931877,
+  289978712,
+  308477138,
+  322224297,
+  328507075,
+  364618648,
+  373686606,
+  395149279,
+  406299962,
+  410021292,
+  422382219,
+  425134518,
+  436648182,
+  470213788,
+  490995676,
+  499350383,
+  500388152,
+  501539276,
+  506723995,
+  508030765,
+  543183270,
+  569638916,
+  578220624,
+  598463260,
+  601084775,
+  620675224,
+  624373224,
+  625812800,
+  632450625,
+  632771102,
+  636599108,
+  659982646,
+  691671442,
+  708391640,
+  710108672,
+  710332091,
+  713252927,
+  714116022,
+  714866032,
+  726749891,
+  765250386,
+  813436240,
+  815033438,
+  815924988,
+  839812273,
+  847498665,
+  858974136,
+  879424295,
+  909378324,
+  915802576,
+  941675750,
+  964172386,
+  964543550,
+  965086662,
+  970019947,
+  979221494,
+  1004759742,
+  1025236119,
+  1028979965,
+  1034477576,
+  1043759268,
+  1045740737,
+  1056971393,
+  1067787926,
+  1070837443,
+  1073289052,
+  1079791751,
+  1091655952,
+  1094839739,
+  1098999966,
+  1107072218,
+  1108731791,
+  1109767066,
+  1114447069,
+  1122268946,
+  1127848354,
+  1131812424,
+  1137151720,
+  1146141383,
+  1149343864,
+  1151645310,
+  1157473353,
+  1161513287,
+  1174215047,
+  1194203245,
+  1195894527,
+  1205912179,
+  1206838930,
+  1208487248,
+  1220156506,
+  1247921214,
+  1250401226,
+  1251154183,
+  1255069804,
+  1280208446,
+  1282509212,
+  1305098747,
+  1330137458,
+  1332697375,
+  1350000524,
+  1362940859,
+  1366581455,
+  1367897481,
+  1369672754,
+  1370830013,
+  1387444039,
+  1391136667,
+  1402239636,
+  1404889105,
+  1411408806,
+  1411569236,
+  1415229497,
+  1416993533,
+  1438156149,
+  1444890860,
+  1446097750,
+  1451178558,
+  1453933600,
+  1454293819,
+  1456275504,
+  1457156555,
+  1460740608,
+  1471581390,
+  1477781620,
+  1478884149,
+  1481559368,
+  1490495337,
+  1509650080,
+  1514522256,
+  1537489219,
+  1552091886,
+  1563744178,
+  1568231659,
+  1573438767,
+  1578294930,
+  1584534039,
+  1590425740,
+  1625457664,
+  1625984261,
+  1626146277,
+  1634076724,
+  1653484029,
+  1657090042,
+  1671439209,
+  1679564793,
+  1693046931,
+  1707298455,
+  1707578919,
+  1715149367,
+  1719931647,
+  1730832604,
+  1734572355,
+  1742990303,
+  1747236156,
+  1774832701,
+  1777842383,
+  1784658559,
+  1787538022,
+  1790246388,
+  1797078623,
+  1800201016,
+  1803759826,
+  1808910250,
+  1833013532,
+  1844796784,
+  1851223911,
+  1861085964,
+  1874451705,
+  1876145536,
+  1891028584,
+  1898577441,
+  1906197421,
+  1907063253,
+  1911760782,
+  1918871774,
+  1922540413,
+  1926264395,
+  1943839865,
+  1946496894,
+  1963587247,
+  1976819248,
+  1982099574,
+  2015986824,
+  2019737030,
+  2045507156,
+  2046666581,
+  2060815008,
+  2085704634,
+  2092512430,
+  2105038603,
+  2116958770,
+  2133164780
+ ]
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7845e4c01a0e299bbffbbf8c78146f1cb2779e65c1f9d4adcf583ed2a871949a
+size 410567112

policy_postprocessor.json

@@ -0,0 +1,32 @@
+{
+  "name": "policy_postprocessor",
+  "steps": [
+    {
+      "registry_name": "unnormalizer_processor",
+      "config": {
+        "eps": 1e-08,
+        "features": {
+          "action": {
+            "type": "ACTION",
+            "shape": [
+              4
+            ]
+          }
+        },
+        "norm_map": {
+          "VISUAL": "MEAN_STD",
+          "STATE": "MIN_MAX",
+          "ACTION": "MIN_MAX"
+        }
+      },
+      "state_file": "policy_postprocessor_step_0_unnormalizer_processor.safetensors"
+    },
+    {
+      "registry_name": "device_processor",
+      "config": {
+        "device": "cpu",
+        "float_dtype": null
+      }
+    }
+  ]
+}

policy_postprocessor_step_0_unnormalizer_processor.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9524edb58d656c0fed00a942aa2770f7923be41e6427d95b85d963bc2abdc5f9
+size 10568

policy_preprocessor.json

@@ -0,0 +1,72 @@
+{
+  "name": "policy_preprocessor",
+  "steps": [
+    {
+      "registry_name": "rename_observations_processor",
+      "config": {
+        "rename_map": {}
+      }
+    },
+    {
+      "registry_name": "to_batch_processor",
+      "config": {}
+    },
+    {
+      "registry_name": "device_processor",
+      "config": {
+        "device": "cuda",
+        "float_dtype": null
+      }
+    },
+    {
+      "registry_name": "normalizer_processor",
+      "config": {
+        "eps": 1e-08,
+        "features": {
+          "observation.state": {
+            "type": "STATE",
+            "shape": [
+              18
+            ]
+          },
+          "observation.images.top": {
+            "type": "VISUAL",
+            "shape": [
+              3,
+              240,
+              320
+            ]
+          },
+          "observation.images.left_wrist": {
+            "type": "VISUAL",
+            "shape": [
+              3,
+              240,
+              320
+            ]
+          },
+          "observation.images.right_wrist": {
+            "type": "VISUAL",
+            "shape": [
+              3,
+              240,
+              320
+            ]
+          },
+          "action": {
+            "type": "ACTION",
+            "shape": [
+              4
+            ]
+          }
+        },
+        "norm_map": {
+          "VISUAL": "MEAN_STD",
+          "STATE": "MIN_MAX",
+          "ACTION": "MIN_MAX"
+        }
+      },
+      "state_file": "policy_preprocessor_step_3_normalizer_processor.safetensors"
+    }
+  ]
+}

policy_preprocessor_step_3_normalizer_processor.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:200ed1d12687223b3f489d68dfd3ae9a9c4520e6346ec0025e99b87558949d37
+size 10576

train_config.json

@@ -0,0 +1,258 @@
+{
+    "dataset": {
+        "repo_id": "local/puzzle_bimanual_top_wrists",
+        "repo_type": "dataset",
+        "root": "datasets/lerobot/puzzle_bimanual_top_wrists",
+        "episodes": null,
+        "image_transforms": {
+            "enable": false,
+            "max_num_transforms": 3,
+            "random_order": false,
+            "tfs": {
+                "brightness": {
+                    "weight": 1.0,
+                    "type": "ColorJitter",
+                    "kwargs": {
+                        "brightness": [
+                            0.8,
+                            1.2
+                        ]
+                    }
+                },
+                "contrast": {
+                    "weight": 1.0,
+                    "type": "ColorJitter",
+                    "kwargs": {
+                        "contrast": [
+                            0.8,
+                            1.2
+                        ]
+                    }
+                },
+                "saturation": {
+                    "weight": 1.0,
+                    "type": "ColorJitter",
+                    "kwargs": {
+                        "saturation": [
+                            0.5,
+                            1.5
+                        ]
+                    }
+                },
+                "hue": {
+                    "weight": 1.0,
+                    "type": "ColorJitter",
+                    "kwargs": {
+                        "hue": [
+                            -0.05,
+                            0.05
+                        ]
+                    }
+                },
+                "sharpness": {
+                    "weight": 1.0,
+                    "type": "SharpnessJitter",
+                    "kwargs": {
+                        "sharpness": [
+                            0.5,
+                            1.5
+                        ]
+                    }
+                },
+                "affine": {
+                    "weight": 1.0,
+                    "type": "RandomAffine",
+                    "kwargs": {
+                        "degrees": [
+                            -5.0,
+                            5.0
+                        ],
+                        "translate": [
+                            0.05,
+                            0.05
+                        ]
+                    }
+                }
+            }
+        },
+        "revision": null,
+        "use_imagenet_stats": true,
+        "video_backend": "torchcodec",
+        "return_uint8": false,
+        "depth_output_unit": "mm",
+        "streaming": false,
+        "eval_split": 0.0
+    },
+    "env": null,
+    "policy": {
+        "type": "diffusion",
+        "n_obs_steps": 2,
+        "input_features": {
+            "observation.state": {
+                "type": "STATE",
+                "shape": [
+                    18
+                ]
+            },
+            "observation.images.top": {
+                "type": "VISUAL",
+                "shape": [
+                    3,
+                    240,
+                    320
+                ]
+            },
+            "observation.images.left_wrist": {
+                "type": "VISUAL",
+                "shape": [
+                    3,
+                    240,
+                    320
+                ]
+            },
+            "observation.images.right_wrist": {
+                "type": "VISUAL",
+                "shape": [
+                    3,
+                    240,
+                    320
+                ]
+            }
+        },
+        "output_features": {
+            "action": {
+                "type": "ACTION",
+                "shape": [
+                    4
+                ]
+            }
+        },
+        "device": "cuda",
+        "use_amp": true,
+        "use_peft": false,
+        "push_to_hub": false,
+        "repo_id": null,
+        "private": null,
+        "tags": null,
+        "license": null,
+        "pretrained_path": null,
+        "pretrained_revision": null,
+        "horizon": 16,
+        "n_action_steps": 8,
+        "normalization_mapping": {
+            "VISUAL": "MEAN_STD",
+            "STATE": "MIN_MAX",
+            "ACTION": "MIN_MAX"
+        },
+        "drop_n_last_frames": 7,
+        "vision_backbone": "resnet18",
+        "resize_shape": null,
+        "crop_ratio": 1.0,
+        "crop_shape": [
+            216,
+            288
+        ],
+        "crop_is_random": true,
+        "pretrained_backbone_weights": "ResNet18_Weights.IMAGENET1K_V1",
+        "use_group_norm": false,
+        "spatial_softmax_num_keypoints": 32,
+        "use_separate_rgb_encoder_per_camera": true,
+        "down_dims": [
+            256,
+            512,
+            1024
+        ],
+        "kernel_size": 5,
+        "n_groups": 8,
+        "diffusion_step_embed_dim": 128,
+        "use_film_scale_modulation": true,
+        "gradient_checkpointing": false,
+        "noise_scheduler_type": "DDIM",
+        "num_train_timesteps": 100,
+        "beta_schedule": "squaredcos_cap_v2",
+        "beta_start": 0.0001,
+        "beta_end": 0.02,
+        "prediction_type": "epsilon",
+        "clip_sample": true,
+        "clip_sample_range": 1.0,
+        "num_inference_steps": 10,
+        "compile_model": false,
+        "compile_mode": "reduce-overhead",
+        "do_mask_loss_for_padding": false,
+        "optimizer_lr": 0.0001,
+        "optimizer_betas": [
+            0.95,
+            0.999
+        ],
+        "optimizer_eps": 1e-08,
+        "optimizer_weight_decay": 1e-06,
+        "scheduler_name": "cosine",
+        "scheduler_warmup_steps": 500
+    },
+    "reward_model": null,
+    "output_dir": "outputs/train/diffusion",
+    "job_name": "diffusion",
+    "resume": false,
+    "seed": 1000,
+    "cudnn_deterministic": false,
+    "num_workers": 4,
+    "batch_size": 32,
+    "prefetch_factor": 2,
+    "persistent_workers": true,
+    "dataloader_multiprocessing_context": "spawn",
+    "steps": 100000,
+    "env_eval_freq": 20000,
+    "log_freq": 200,
+    "eval_steps": 0,
+    "max_eval_samples": 0,
+    "tolerance_s": 0.0001,
+    "save_checkpoint": true,
+    "save_freq": 10000,
+    "use_policy_training_preset": true,
+    "optimizer": {
+        "type": "adam",
+        "lr": 0.0001,
+        "weight_decay": 1e-06,
+        "grad_clip_norm": 10.0,
+        "betas": [
+            0.95,
+            0.999
+        ],
+        "eps": 1e-08
+    },
+    "scheduler": {
+        "type": "diffuser",
+        "num_warmup_steps": 500,
+        "name": "cosine"
+    },
+    "eval": {
+        "n_episodes": 50,
+        "batch_size": 11,
+        "use_async_envs": true,
+        "recording": false,
+        "recording_repo_id": null,
+        "recording_private": false
+    },
+    "wandb": {
+        "enable": false,
+        "disable_artifact": false,
+        "project": "lerobot",
+        "entity": null,
+        "notes": null,
+        "run_id": null,
+        "mode": null,
+        "add_tags": true
+    },
+    "peft": null,
+    "job": {
+        "target": null,
+        "image": "huggingface/lerobot-gpu:latest",
+        "timeout": "2d",
+        "detach": false,
+        "tags": []
+    },
+    "save_checkpoint_to_hub": false,
+    "sample_weighting": null,
+    "rename_map": {},
+    "checkpoint_path": null
+}