File size: 3,689 Bytes
752c314 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Standalone real-robot inference for PushT-xarm BFN-hybrid policy.
Inputs each step:
cam0: HxWx3 uint8 RGB image from the top camera (any size; will be resized to 224x224)
cam0_prev: same, one step earlier (n_obs_steps=2)
Output per step (predicts horizon=16, returns next n_action_steps=8):
direction: int in {0..7}
distance: float in [0, 50]
"""
import argparse
import sys
import time
from pathlib import Path
import numpy as np
import torch
import yaml
from PIL import Image
THIS_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(THIS_DIR))
from bfn_hybrid_image_policy import BFNHybridImagePolicy # noqa: E402
def load_bfn_policy(ckpt_path: str, config_path: str, device: str = "cuda"):
with open(config_path) as f:
cfg = yaml.safe_load(f)
pcfg = cfg["policy"]
policy = BFNHybridImagePolicy(
shape_meta=cfg["shape_meta"],
horizon=cfg["horizon"],
n_action_steps=cfg["n_action_steps"],
n_obs_steps=cfg["n_obs_steps"],
num_discrete_actions=pcfg.get("num_discrete_actions", 8),
continuous_param_dim=pcfg.get("continuous_param_dim", 1),
sigma_1=pcfg.get("sigma_1", 0.001),
beta_1=pcfg.get("beta_1", 0.2),
n_timesteps=pcfg.get("n_timesteps", 20),
crop_shape=tuple(pcfg.get("crop_shape", [216, 216])),
obs_encoder_group_norm=pcfg.get("obs_encoder_group_norm", True),
eval_fixed_crop=pcfg.get("eval_fixed_crop", True),
diffusion_step_embed_dim=pcfg.get("diffusion_step_embed_dim", 128),
down_dims=tuple(pcfg.get("down_dims", [256, 512, 1024])),
kernel_size=pcfg.get("kernel_size", 5),
n_groups=pcfg.get("n_groups", 8),
cond_predict_scale=pcfg.get("cond_predict_scale", True),
)
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
state = ckpt["state_dicts"]["model"] if "state_dicts" in ckpt else ckpt
policy.load_state_dict(state)
policy.to(device).eval()
return policy
def preprocess_image(img: np.ndarray) -> np.ndarray:
"""Resize HxWx3 uint8 -> 3x224x224 float32 in [0,1]."""
if img.shape[:2] != (224, 224):
img = np.array(Image.fromarray(img).resize((224, 224), Image.BILINEAR))
img = img.astype(np.float32) / 255.0
return img.transpose(2, 0, 1)
def infer_step(policy, cam0_now: np.ndarray, cam0_prev: np.ndarray, device: str = "cuda"):
"""One inference call. Returns a list of dicts [{direction, distance}, ...] of length n_action_steps."""
a = preprocess_image(cam0_prev)
b = preprocess_image(cam0_now)
obs = torch.from_numpy(np.stack([a, b])).unsqueeze(0).to(device) # [1, 2, 3, 224, 224]
with torch.no_grad():
out = policy.predict_action({"camera_0": obs})
actions = out["action"][0].cpu().numpy() # [n_action_steps, 2] = [direction, distance]
return [
{"direction": int(round(a[0])) % 8, "distance": float(np.clip(a[1], 0, 50))}
for a in actions
]
def main():
p = argparse.ArgumentParser()
p.add_argument("--ckpt", required=True)
p.add_argument("--config", required=True)
p.add_argument("--device", default="cuda")
args = p.parse_args()
print(f"Loading policy from {args.ckpt}...")
policy = load_bfn_policy(args.ckpt, args.config, args.device)
print("Policy loaded.")
# Dummy roundtrip test
dummy = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
t0 = time.time()
actions = infer_step(policy, dummy, dummy, args.device)
dt = (time.time() - t0) * 1000
print(f"Smoke test: {len(actions)} actions in {dt:.1f} ms")
print(f"First action: {actions[0]}")
if __name__ == "__main__":
main()
|