README.md

---
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

| File | Method | Action treatment | Inference steps |
|------|--------|------------------|-----------------|
| `bfn/latest.ckpt`  | **BFN-Hybrid** (categorical + continuous Bayesian flow) | true hybrid | 20 |
| `ddpm/latest.ckpt` | **DDPM** | one-hot continuous (9D) | 100 |

## Quick start (BFN)

```bash
pip install -r requirements.txt
python inference.py --ckpt bfn/latest.ckpt --config bfn/policy_config.yaml
```

Programmatic use:

```python
from inference import load_bfn_policy, infer_step

policy = load_bfn_policy("bfn/latest.ckpt", "bfn/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 checkpoint

The DDPM policy 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/
  latest.ckpt
  policy_config.yaml
ddpm/
  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
```