add onnx support

examples/onnx/eval_act_onnx.py

@@ -0,0 +1,179 @@
+#!/usr/bin/env python
+"""Evaluate an ACT policy in sim with either the PyTorch or ONNX network.
+
+The ONNX backend swaps only ``policy.model`` (ResNet + transformer + action head)
+with an onnxruntime session. Everything else - the LeRobot processor pipeline
+(normalization), the action queue, and the gym env - is identical, so any
+difference in success rate is attributable to the network backend alone.
+
+Run both backends with the same seed to compare:
+
+    python examples/onnx/eval_act_onnx.py \
+        --policy-path=lerobot/act_aloha_sim_transfer_cube_human \
+        --task=AlohaTransferCube-v0 \
+        --backend=torch --n-episodes=50 --batch-size=10 --device=cuda
+
+    python examples/onnx/eval_act_onnx.py \
+        --policy-path=lerobot/act_aloha_sim_transfer_cube_human \
+        --task=AlohaTransferCube-v0 \
+        --onnx=outputs/onnx/act_transfer_cube.onnx \
+        --backend=onnx --n-episodes=50 --batch-size=10 --device=cuda
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+
+from lerobot.envs.factory import make_env, make_env_config, make_env_pre_post_processors
+from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.policies.factory import make_pre_post_processors
+from lerobot.scripts.lerobot_eval import eval_policy
+from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
+from lerobot.utils.random_utils import set_seed
+
+
+class ONNXACTModel(nn.Module):
+    """Drop-in replacement for ``ACTPolicy.model`` backed by onnxruntime."""
+
+    def __init__(self, onnx_path: str, image_keys: list[str], has_state: bool, has_env_state: bool, device: str):
+        super().__init__()
+        import onnxruntime as ort
+
+        providers = (
+            ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            if str(device).startswith("cuda")
+            else ["CPUExecutionProvider"]
+        )
+        so = ort.SessionOptions()
+        so.log_severity_level = 3
+        self.sess = ort.InferenceSession(onnx_path, sess_options=so, providers=providers)
+        self.image_keys = image_keys
+        self.has_state = has_state
+        self.has_env_state = has_env_state
+        print(f"[onnx] providers in use: {self.sess.get_providers()}")
+
+    def forward(self, batch: dict):
+        if self.has_state:
+            state = batch[OBS_STATE]
+        else:
+            state = batch[OBS_ENV_STATE]
+        ref = state
+        ort_inputs = {"state": state.detach().cpu().numpy().astype(np.float32)}
+        images = batch[OBS_IMAGES]
+        for i, img in enumerate(images):
+            ort_inputs[f"image_{i}"] = img.detach().cpu().numpy().astype(np.float32)
+        out = self.sess.run(None, ort_inputs)[0]
+        actions = torch.from_numpy(out).to(ref.device, dtype=ref.dtype)
+        return actions, None
+
+
+def load_stats_from_checkpoint(policy_path: str, input_features, output_features) -> dict:
+    """Recover MEAN_STD stats baked into a legacy ACT checkpoint's safetensors buffers.
+
+    Legacy checkpoints store normalization as buffers like
+    ``normalize_inputs.buffer_observation_state.{mean,std}``. We map those back to
+    feature names so we can rebuild the processor pipeline without the dataset.
+    """
+    from safetensors.torch import load_file
+
+    p = Path(policy_path)
+    if p.is_dir():
+        st_path = p / "model.safetensors"
+    else:
+        from huggingface_hub import hf_hub_download
+
+        st_path = Path(hf_hub_download(policy_path, "model.safetensors"))
+
+    sd = load_file(str(st_path))
+    stats: dict = {}
+    for feat in list(input_features) + list(output_features):
+        buf = "buffer_" + feat.replace(".", "_")
+        for prefix in ("normalize_inputs", "normalize_targets", "unnormalize_outputs"):
+            mkey, skey = f"{prefix}.{buf}.mean", f"{prefix}.{buf}.std"
+            if mkey in sd and skey in sd:
+                stats[feat] = {"mean": sd[mkey].numpy(), "std": sd[skey].numpy()}
+                break
+    return stats
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--policy-path", required=True)
+    parser.add_argument("--task", required=True, help="e.g. AlohaTransferCube-v0")
+    parser.add_argument("--env-type", default="aloha")
+    parser.add_argument("--backend", choices=["torch", "onnx"], default="torch")
+    parser.add_argument("--onnx", default=None, help="Path to .onnx (required for --backend=onnx)")
+    parser.add_argument("--n-episodes", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=10)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--seed", type=int, default=1000)
+    args = parser.parse_args()
+
+    if args.backend == "onnx" and not args.onnx:
+        raise SystemExit("--backend=onnx requires --onnx=<path>")
+
+    device = "cuda" if (args.device == "cuda" and torch.cuda.is_available()) else "cpu"
+    set_seed(args.seed)
+
+    print(f"[1/4] Loading ACT policy from '{args.policy_path}'...")
+    policy = ACTPolicy.from_pretrained(args.policy_path)
+    policy.config.device = device
+    policy.eval()
+    policy.to(device)
+    cfg = policy.config
+
+    if args.backend == "onnx":
+        image_keys = list(cfg.image_features)
+        has_state = cfg.robot_state_feature is not None
+        has_env_state = cfg.env_state_feature is not None
+        print(f"[2/4] Swapping policy.model with ONNX backend ({args.onnx})")
+        policy.model = ONNXACTModel(args.onnx, image_keys, has_state, has_env_state, device)
+        policy.to(device)
+    else:
+        print("[2/4] Using PyTorch backend")
+
+    print("[3/4] Building processors and environment...")
+    stats = load_stats_from_checkpoint(args.policy_path, cfg.input_features, cfg.output_features)
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=cfg,
+        dataset_stats=stats,
+        preprocessor_overrides={"device_processor": {"device": device}},
+    )
+
+    env_cfg = make_env_config(args.env_type, task=args.task)
+    env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=env_cfg, policy_cfg=cfg)
+    env_groups = make_env(env_cfg, n_envs=args.batch_size, use_async_envs=False)
+    # make_env returns {task_group: {idx: VectorEnv}}; grab the single env.
+    first_group = next(iter(env_groups.values()))
+    env = next(iter(first_group.values()))
+
+    print(f"[4/4] Evaluating backend='{args.backend}' for {args.n_episodes} episodes (seed={args.seed})...")
+    with torch.no_grad():
+        info = eval_policy(
+            env=env,
+            policy=policy,
+            env_preprocessor=env_preprocessor,
+            env_postprocessor=env_postprocessor,
+            preprocessor=preprocessor,
+            postprocessor=postprocessor,
+            n_episodes=args.n_episodes,
+            start_seed=args.seed,
+        )
+
+    agg = info["aggregated"]
+    print("\n==== RESULT ====")
+    print(f"backend       : {args.backend}")
+    print(f"task          : {args.task}")
+    print(f"n_episodes    : {args.n_episodes}")
+    print(f"pc_success    : {agg['pc_success']:.1f}%")
+    print(f"avg_max_reward: {agg['avg_max_reward']:.4f}")
+    print(f"eval_ep_s     : {agg['eval_ep_s']:.2f}s")
+
+    env.close()
+
+
+if __name__ == "__main__":
+    main()