inference.py

"""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()