Migrate policy to PolicyProcessorPipeline system

README.md

@@ -1,70 +1,136 @@
 ---
+datasets: unknown
+library_name: lerobot
 license: apache-2.0
-datasets:
-- lerobot/pusht
+model_name: diffusion
+pipeline_tag: robotics
 tags:
-- diffusion-policy
-- model_hub_mixin
-- pytorch_model_hub_mixin
+- lerobot
 - robotics
-pipeline_tag: robotics
+- diffusion
+---
+
+# Model Card for diffusion
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+[Diffusion Policy](https://huggingface.co/papers/2303.04137) treats visuomotor control as a generative diffusion process, producing smooth, multi-step action trajectories that excel at contact-rich manipulation.
+
+
+
+
+<!-- A short demo is worth more than any description! Record a GIF/video of the policy
+running on your robot, upload it to this repo, and embed it here:
+<p align="center">
+  <img src="https://huggingface.co/<hf_user>/<policy_repo_id>/resolve/main/demo.gif" width="60%"/>
+</p>
+-->
+
+This policy has been trained and pushed to the Hub using [LeRobot](https://github.com/huggingface/lerobot).
+
+See the [full LeRobot documentation](https://huggingface.co/docs/lerobot/index).
+
+
 ---
-# Model Card for Diffusion Policy / PushT
 
-Diffusion Policy (as per [Diffusion Policy: Visuomotor Policy
-Learning via Action Diffusion](https://arxiv.org/abs/2303.04137)) trained for the `PushT` environment from [gym-pusht](https://github.com/huggingface/gym-pusht).
+## Model Details
+
+- **License:** apache-2.0
+
+
+## Inputs & Outputs
+
+The policy consumes these observation features and produces these action features.
+
+**Inputs**
+
+| Feature | Type | Shape |
+| --- | --- | --- |
+| `observation.image` | VISUAL | `(3, 96, 96)` |
+| `observation.state` | STATE | `(2,)` |
+
+**Outputs**
+
+| Feature | Type | Shape |
+| --- | --- | --- |
+| `action` | ACTION | `(2,)` |
+
+
+
+---
 
 ## How to Get Started with the Model
 
-See the [LeRobot library](https://github.com/huggingface/lerobot) (particularly the [evaluation script](https://github.com/huggingface/lerobot/blob/main/lerobot/scripts/eval.py)) for instructions on how to load and evaluate this model.
+New to LeRobot? These guides cover the full workflow:
 
-## Training Details
+- **[Install LeRobot](https://huggingface.co/docs/lerobot/main/en/installation)** — set up the `lerobot` package.
+- **[Hardware setup](https://huggingface.co/docs/lerobot/main/en/hardware_guide)** — assemble, wire, and calibrate your robot and cameras.
+- **[Record data & train a policy](https://huggingface.co/docs/lerobot/en/il_robots)** — the end-to-end imitation-learning walkthrough.
+- **[CLI cheat-sheet](https://huggingface.co/docs/lerobot/main/en/cheat-sheet)** — quick reference for the `lerobot-*` commands.
 
-Trained with [LeRobot@3c0a209](https://github.com/huggingface/lerobot/tree/3c0a209f9fac4d2a57617e686a7f2a2309144ba2).
+The short version to run and train this policy:
 
-The model was trained using [LeRobot's training script](https://github.com/huggingface/lerobot/blob/main/lerobot/scripts/train.py) and with the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using this command:
+### Run the policy on your robot
 
 ```bash
-python lerobot/scripts/train.py \
-    --output_dir=outputs/train/diffusion_pusht \
-    --policy.type=diffusion \
-    --dataset.repo_id=lerobot/pusht \
-    --seed=100000 \
-    --env.type=pusht \
-    --batch_size=64 \
-    --steps=200000 \
-    --eval_freq=25000 \
-    --save_freq=25000 \
-    --wandb.enable=true
+lerobot-rollout \
+  --strategy.type=base \
+  --robot.type=<your_robot_type> \
+  --robot.port=<your_robot_port> \
+  --robot.cameras="{ <camera_1>: {type: opencv, index_or_path: <index_or_path>, width: 640, height: 480, fps: 30}, <camera_2>: {type: opencv, index_or_path: <index_or_path>, width: 640, height: 480, fps: 30}}" \
+  --policy.path=<hf_user>/<policy_repo_id> \
+  --task="<your_task_description>" \
+  --duration=60
 ```
 
+Replace the remaining `<...>` placeholders with your own values: `--robot.port` and the camera names/indices are specific to your machine, and the camera names must match the observation keys this policy was trained on.
+
+When `--strategy.type=base` is used the script doesn't record the episodes. Skipping duration will make the policy run indefinitely. For more information look at [rollout documentation](https://huggingface.co/docs/lerobot/main/en/inference).
 
-The training curves may be found at https://wandb.ai/aliberts/lerobot/runs/s7elvf4r.
-The current model corresponds to the checkpoint at 175k steps.
+### Train your own policy
+
+```bash
+lerobot-train \
+  --dataset.repo_id=${HF_USER}/<dataset> \
+  --policy.type=diffusion \
+  --output_dir=outputs/train/<policy_repo_id> \
+  --job_name=lerobot_training \
+  --policy.device=cuda \
+  --policy.repo_id=${HF_USER}/<policy_repo_id> \
+  --wandb.enable=true
+```
+
+_Writes checkpoints to `outputs/train/<policy_repo_id>/checkpoints/`._
+
+---
 
 ## Evaluation
 
-The model was evaluated on the `PushT` environment from [gym-pusht](https://github.com/huggingface/gym-pusht) and compared to a similar model trained with the original [Diffusion Policy code](https://github.com/real-stanford/diffusion_policy). There are two evaluation metrics on a per-episode basis:
+<!-- Report real-robot results here: run the policy several times per task and count the
+successes. Delete the "No evaluation results" line and fill in this table instead:
 
-- Maximum overlap with target (seen as `eval/avg_max_reward` in the charts above). This ranges in [0, 1].
-- Success: whether or not the maximum overlap is at least 95%.
+| Task | Trials | Successes | Success rate |
+| ---- | ------ | --------- | ------------ |
+| pick the lego brick | 10 | 8 | 80% |
 
-Here are the metrics for 500 episodes worth of evaluation. The "Theirs" column is for an equivalent model trained on the original Diffusion Policy repository and evaluated on LeRobot (the model weights may be found in the [`original_dp_repo`](https://huggingface.co/lerobot/diffusion_pusht/tree/original_dp_repo) branch of this respository).
+Also worth noting: anything that affects difficulty (new object positions, lighting,
+distractors, a different robot of the same type, ...).
+-->
 
-<blank>|Ours|Theirs
--|-|-
-Average max. overlap ratio | 0.955 | 0.957
-Success rate for 500 episodes (%) | 65.4 | 64.2
+_No evaluation results have been provided for this policy yet._
 
-The results of each of the individual rollouts may be found in [eval_info.json](eval_info.json).
-It was produced after training with this command:
-```bash
-python lerobot/scripts/eval.py \
-    --policy.path=outputs/train/diffusion_pusht/checkpoints/175000/pretrained_model \
-    --output_dir=outputs/eval/diffusion_pusht/175000 \
-    --env.type=pusht \
-    --eval.n_episodes=500 \
-    --eval.batch_size=50 \
-    --device=cuda \
-    --use_amp=false
+---
+
+## Citation
+
+If you use this policy, please cite the method linked in the description above, along with LeRobot:
+
+```bibtex
+@misc{cadene2024lerobot,
+    author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascal, Caroline and Choghari, Jade and Moss, Jess and Wolf, Thomas},
+    title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
+    howpublished = "\url{https://github.com/huggingface/lerobot}",
+    year = {2024}
+}
 ```

config.json

@@ -1,76 +1,88 @@
 {
-  "beta_end": 0.02,
-  "beta_schedule": "squaredcos_cap_v2",
-  "beta_start": 0.0001,
-  "clip_sample": true,
-  "clip_sample_range": 1.0,
-  "crop_is_random": true,
-  "crop_shape": [
-    84,
-    84
-  ],
-  "diffusion_step_embed_dim": 128,
-  "do_mask_loss_for_padding": false,
-  "down_dims": [
-    512,
-    1024,
-    2048
-  ],
-  "drop_n_last_frames": 7,
-  "horizon": 16,
-  "input_features": {
-    "observation.image": {
-      "shape": [
-        3,
-        96,
-        96
-      ],
-      "type": "VISUAL"
+    "type": "diffusion",
+    "n_obs_steps": 2,
+    "input_features": {
+        "observation.image": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                96,
+                96
+            ]
+        },
+        "observation.state": {
+            "type": "STATE",
+            "shape": [
+                2
+            ]
+        }
     },
-    "observation.state": {
-      "shape": [
-        2
-      ],
-      "type": "STATE"
-    }
-  },
-  "kernel_size": 5,
-  "n_action_steps": 8,
-  "n_groups": 8,
-  "n_obs_steps": 2,
-  "device": "cuda",
-  "use_amp": false,
-  "noise_scheduler_type": "DDPM",
-  "normalization_mapping": {
-    "ACTION": "MIN_MAX",
-    "STATE": "MIN_MAX",
-    "VISUAL": "MEAN_STD"
-  },
-  "num_inference_steps": null,
-  "num_train_timesteps": 100,
-  "optimizer_betas": [
-    0.95,
-    0.999
-  ],
-  "optimizer_eps": 1e-08,
-  "optimizer_lr": 0.0001,
-  "optimizer_weight_decay": 1e-06,
-  "output_features": {
-    "action": {
-      "shape": [
-        2
-      ],
-      "type": "ACTION"
-    }
-  },
-  "prediction_type": "epsilon",
-  "pretrained_backbone_weights": null,
-  "scheduler_name": "cosine",
-  "scheduler_warmup_steps": 500,
-  "spatial_softmax_num_keypoints": 32,
-  "type": "diffusion",
-  "use_film_scale_modulation": true,
-  "use_group_norm": true,
-  "use_separate_rgb_encoder_per_camera": false,
-  "vision_backbone": "resnet18"
+    "output_features": {
+        "action": {
+            "type": "ACTION",
+            "shape": [
+                2
+            ]
+        }
+    },
+    "device": "cuda",
+    "use_amp": false,
+    "use_peft": false,
+    "push_to_hub": true,
+    "repo_id": null,
+    "private": null,
+    "tags": null,
+    "license": null,
+    "pretrained_path": null,
+    "pretrained_revision": null,
+    "horizon": 16,
+    "n_action_steps": 8,
+    "normalization_mapping": {
+        "ACTION": "MIN_MAX",
+        "STATE": "MIN_MAX",
+        "VISUAL": "MEAN_STD"
+    },
+    "drop_n_last_frames": 7,
+    "vision_backbone": "resnet18",
+    "resize_shape": null,
+    "crop_ratio": 1.0,
+    "crop_shape": [
+        84,
+        84
+    ],
+    "crop_is_random": true,
+    "pretrained_backbone_weights": null,
+    "use_group_norm": true,
+    "spatial_softmax_num_keypoints": 32,
+    "use_separate_rgb_encoder_per_camera": false,
+    "down_dims": [
+        512,
+        1024,
+        2048
+    ],
+    "kernel_size": 5,
+    "n_groups": 8,
+    "diffusion_step_embed_dim": 128,
+    "use_film_scale_modulation": true,
+    "noise_scheduler_type": "DDPM",
+    "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": null,
+    "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
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:995d14d35db57d95c35ad9704c3d79c8612b7bc45f3877e5c46c2cdc516856a8
-size 1050862408
+oid sha256:69ecc36b62286bd5472a0e92b64e53ee6e0cb8b96fc05c247c4b7d1be583c2b2
+size 1050861448

policy_postprocessor.json

@@ -0,0 +1,32 @@
+{
+  "name": "policy_postprocessor",
+  "steps": [
+    {
+      "registry_name": "unnormalizer_processor",
+      "config": {
+        "eps": 1e-08,
+        "features": {
+          "action": {
+            "type": "ACTION",
+            "shape": [
+              2
+            ]
+          }
+        },
+        "norm_map": {
+          "ACTION": "MIN_MAX",
+          "STATE": "MIN_MAX",
+          "VISUAL": "MEAN_STD"
+        }
+      },
+      "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:392af8778e650720aa230904998d2c11184583d956f0b7eada100afd09bf2378
+size 504

policy_preprocessor.json

@@ -0,0 +1,56 @@
+{
+  "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.image": {
+            "type": "VISUAL",
+            "shape": [
+              3,
+              96,
+              96
+            ]
+          },
+          "observation.state": {
+            "type": "STATE",
+            "shape": [
+              2
+            ]
+          },
+          "action": {
+            "type": "ACTION",
+            "shape": [
+              2
+            ]
+          }
+        },
+        "norm_map": {
+          "ACTION": "MIN_MAX",
+          "STATE": "MIN_MAX",
+          "VISUAL": "MEAN_STD"
+        }
+      },
+      "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:392af8778e650720aa230904998d2c11184583d956f0b7eada100afd09bf2378
+size 504