add onnx support

examples/onnx/convert_legacy_checkpoint.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+"""Convert a legacy LeRobot checkpoint to the current processor-pipeline format.
+
+Older hub checkpoints (e.g. ``lerobot/act_aloha_sim_insertion_human``) bake
+normalization stats into the model weights and do not ship
+``policy_preprocessor.json`` / ``policy_postprocessor.json``. Current ``main``
+loads those processor configs from the checkpoint, so eval/rollout fail with
+``FileNotFoundError: Could not find 'policy_preprocessor.json'``.
+
+This script rebuilds the processors from the training dataset's stats and saves
+a pipeline-format checkpoint locally that ``lerobot-eval`` can consume directly.
+
+Usage:
+    python examples/onnx/convert_legacy_checkpoint.py \
+        --policy-path=lerobot/act_aloha_sim_insertion_human \
+        --dataset-repo-id=lerobot/aloha_sim_insertion_human \
+        --output-dir=outputs/converted/act_aloha_sim_insertion_human
+
+Then:
+    lerobot-eval \
+        --policy.path=outputs/converted/act_aloha_sim_insertion_human \
+        --env.type=aloha --env.task=AlohaInsertion-v0 \
+        --eval.batch_size=10 --eval.n_episodes=50 \
+        --eval.use_async_envs=false --policy.device=cuda
+"""
+
+import argparse
+from pathlib import Path
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+from lerobot.policies.factory import make_policy, make_pre_post_processors
+from lerobot.utils.constants import (
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--policy-path", required=True, help="Legacy checkpoint repo id or local dir")
+    parser.add_argument(
+        "--dataset-repo-id",
+        required=True,
+        help="Training dataset repo id, used only for normalization stats",
+    )
+    parser.add_argument("--output-dir", required=True, help="Where to save the converted checkpoint")
+    parser.add_argument("--device", default="cpu", help="Device for building the policy (cpu is fine)")
+    args = parser.parse_args()
+
+    out = Path(args.output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    print(f"[1/4] Loading dataset stats from '{args.dataset_repo_id}' (metadata only)...")
+    ds_meta = LeRobotDatasetMetadata(args.dataset_repo_id)
+
+    print(f"[2/4] Loading policy weights from '{args.policy_path}'...")
+    cfg = PreTrainedConfig.from_pretrained(args.policy_path)
+    cfg.pretrained_path = args.policy_path
+    cfg.device = args.device
+    policy = make_policy(cfg, ds_meta=ds_meta)
+
+    print("[3/4] Building processors from dataset stats...")
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=policy.config,
+        dataset_stats=ds_meta.stats,
+    )
+
+    print(f"[4/4] Saving pipeline-format checkpoint to '{out}'...")
+    policy.save_pretrained(out)
+    preprocessor.save_pretrained(out, config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json")
+    postprocessor.save_pretrained(out, config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json")
+
+    print(f"\nDone. Converted checkpoint at: {out}")
+    print("Eval it with --policy.path=" + str(out))
+
+
+if __name__ == "__main__":
+    main()

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

examples/onnx/export_act.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python
+"""Export an ACT policy's network to ONNX and verify numerical parity.
+
+Only the inference network is exported (ResNet backbone + transformer enc/dec +
+action head). The VAE encoder is training-only and the inference latent is zeros,
+so the exported graph is a pure function of (state, images) -> action_chunk.
+Normalization stays in the LeRobot processor pipeline (outside ONNX).
+
+Usage:
+    python examples/onnx/export_act.py \
+        --policy-path=outputs/converted/act_aloha_sim_transfer_cube_human \
+        --output=outputs/onnx/act_transfer_cube.onnx
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+
+from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
+
+
+class ACTExportWrapper(nn.Module):
+    """Tensor-in/tensor-out wrapper around ACT's inference network."""
+
+    def __init__(self, model: nn.Module, image_keys: list[str], has_state: bool, has_env_state: bool):
+        super().__init__()
+        self.model = model
+        self.image_keys = image_keys
+        self.has_state = has_state
+        self.has_env_state = has_env_state
+
+    def forward(self, state: torch.Tensor, *images: torch.Tensor) -> torch.Tensor:
+        batch: dict = {}
+        if self.has_state:
+            batch[OBS_STATE] = state
+        if self.has_env_state:
+            # Convention: when env_state is used it is passed as `state`.
+            batch[OBS_ENV_STATE] = state
+        batch[OBS_IMAGES] = list(images)
+        actions, _ = self.model(batch)
+        return actions
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--policy-path", required=True, help="Converted ACT checkpoint dir or repo id")
+    parser.add_argument("--output", required=True, help="Output .onnx path")
+    parser.add_argument("--opset", type=int, default=17)
+    parser.add_argument("--atol", type=float, default=1e-3)
+    parser.add_argument("--device", default="cpu")
+    args = parser.parse_args()
+
+    out = Path(args.output)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    print(f"[1/4] Loading ACT policy from '{args.policy_path}'...")
+    policy = ACTPolicy.from_pretrained(args.policy_path)
+    policy.eval()
+    policy.to(args.device)
+    cfg = policy.config
+
+    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
+    state_dim = (cfg.robot_state_feature or cfg.env_state_feature).shape[0]
+
+    print(f"      image_keys={image_keys} state_dim={state_dim} "
+          f"chunk_size={cfg.chunk_size} action_dim={cfg.action_feature.shape[0]}")
+
+    wrapper = ACTExportWrapper(policy.model, image_keys, has_state, has_env_state).eval().to(args.device)
+
+    # Build example inputs (batch size 1) from the config feature shapes.
+    state_example = torch.randn(1, state_dim, device=args.device)
+    image_examples = [
+        torch.rand(1, *cfg.image_features[k].shape, device=args.device) for k in image_keys
+    ]
+    example_inputs = (state_example, *image_examples)
+
+    input_names = ["state"] + [f"image_{i}" for i in range(len(image_keys))]
+    output_names = ["action_chunk"]
+    dynamic_axes = {name: {0: "batch"} for name in input_names + output_names}
+
+    print(f"[2/4] Exporting to ONNX (opset {args.opset}) -> {out}")
+    torch.onnx.export(
+        wrapper,
+        example_inputs,
+        str(out),
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=dynamic_axes,
+        opset_version=args.opset,
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    print("[3/4] Running parity check (torch vs onnxruntime)...")
+    import onnxruntime as ort
+
+    providers = ["CPUExecutionProvider"]
+    so = ort.SessionOptions()
+    so.log_severity_level = 3
+    sess = ort.InferenceSession(str(out), sess_options=so, providers=providers)
+
+    # Fresh random inputs for the check.
+    state_check = torch.randn(2, state_dim, device=args.device)
+    image_check = [torch.rand(2, *cfg.image_features[k].shape, device=args.device) for k in image_keys]
+
+    with torch.no_grad():
+        torch_out = wrapper(state_check, *image_check).cpu().numpy()
+
+    ort_inputs = {"state": state_check.cpu().numpy()}
+    for i, img in enumerate(image_check):
+        ort_inputs[f"image_{i}"] = img.cpu().numpy()
+    ort_out = sess.run(None, ort_inputs)[0]
+
+    max_abs = float(np.max(np.abs(torch_out - ort_out)))
+    mean_abs = float(np.mean(np.abs(torch_out - ort_out)))
+    print(f"      shapes: torch={torch_out.shape} onnx={ort_out.shape}")
+    print(f"      max_abs_diff={max_abs:.3e} mean_abs_diff={mean_abs:.3e} (atol={args.atol:.0e})")
+
+    ok = max_abs <= args.atol
+    print(f"[4/4] Parity: {'PASS' if ok else 'FAIL'}")
+    if not ok:
+        raise SystemExit(f"Parity check failed: max_abs_diff {max_abs:.3e} > atol {args.atol:.0e}")
+    print(f"\nDone. ONNX model at: {out}")
+
+
+if __name__ == "__main__":
+    main()