| --- |
| license: mit |
| tags: |
| - robotics |
| - bfn |
| - bayesian-flow-networks |
| - diffusion-policy |
| - pusht |
| - xarm |
| - hybrid-action |
| --- |
| |
| # PushT-xarm Real-Robot Policies (BFN-Hybrid vs DDPM-OneHot) |
|
|
| Real-robot push-T policies trained on [borueihuang/pusht_xarm_merged](https://huggingface.co/datasets/borueihuang/pusht_xarm_merged). |
| Both policies were trained on 144 episodes / 7835 frames at 30 Hz, top camera only, 200 epochs. |
|
|
| ## Action space |
|
|
| - **Discrete:** 8 push directions (`action.direction` in {0..7}) |
| - **Continuous:** push distance (`action.distance` in `[0, 50]`) |
|
|
| ## Observation space |
|
|
| - `camera_0`: top-view RGB image, 3x224x224, two-step history (`n_obs_steps=2`) |
|
|
| ## Policies (4 checkpoints) |
|
|
| | Folder | Method | Dataset | Action treatment | Inference steps | |
| |--------|--------|---------|------------------|-----------------| |
| | `bfn_merged/` | **BFN-Hybrid** | merged (7835 frames) | true hybrid | 20 | |
| | `bfn_orig/` | **BFN-Hybrid** | original (9229 frames) | true hybrid | 20 | |
| | `ddpm_merged/` | **DDPM** | merged | one-hot continuous (9D) | 100 | |
| | `ddpm_orig/` | **DDPM** | original | one-hot continuous (9D) | 100 | |
|
|
| The **merged** dataset combines consecutive same-direction segments into single actions. |
| The **original** dataset keeps the operator's raw per-step input. |
|
|
| ## Quick start (BFN, merged) |
|
|
| ```bash |
| pip install -r requirements.txt |
| python inference.py --ckpt bfn_merged/latest.ckpt --config bfn_merged/policy_config.yaml |
| ``` |
|
|
| Programmatic use: |
|
|
| ```python |
| from inference import load_bfn_policy, infer_step |
| |
| policy = load_bfn_policy("bfn_merged/latest.ckpt", "bfn_merged/policy_config.yaml", "cuda") |
| actions = infer_step(policy, cam0_now, cam0_prev, "cuda") |
| # actions: List[{"direction": int 0..7, "distance": float 0..50}], len = n_action_steps (8) |
| ``` |
|
|
| ## DDPM checkpoints |
|
|
| DDPM uses `diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy` |
| from the `diffusion-policy` library. Action is a 9D continuous vector: `[one_hot(8), distance]`. |
| At inference time, take `argmax` of the first 8 dims for the direction, and the 9th dim for distance. |
|
|
| ## Files |
|
|
| ``` |
| bfn_merged/ |
| latest.ckpt |
| policy_config.yaml |
| bfn_orig/ |
| latest.ckpt |
| policy_config.yaml |
| ddpm_merged/ |
| latest.ckpt |
| policy_config.yaml |
| ddpm_orig/ |
| latest.ckpt |
| policy_config.yaml |
| bfn_hybrid_image_policy.py # standalone BFN policy class |
| policies/base.py # BasePolicy abstract class |
| networks/base.py # BFNetwork wrapper |
| inference.py # example loader + inference |
| requirements.txt |
| ``` |
|
|
