README.md

metadata

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

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

Programmatic use:

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