| """ONNX policy wrapper with the training-side metadata.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
| try: |
| import onnxruntime as ort |
| except Exception: |
| ort = None |
|
|
|
|
| def _csv_floats(value: str | None) -> list[float]: |
| return [float(x) for x in value.split(",") if x.strip()] if value else [] |
|
|
|
|
| def _csv_strings(value: str | None) -> list[str]: |
| return [x.strip() for x in value.split(",") if x.strip()] if value else [] |
|
|
|
|
| @dataclass |
| class PolicyMetadata: |
| joint_names: list[str] |
| action_scale: list[float] |
| default_joint_pos: list[float] |
| body_names: list[str] |
| observation_names: list[str] |
| raw: dict |
|
|
|
|
| class OnnxPolicy: |
| """``actions = policy(obs)``; the observation normaliser is baked into the graph.""" |
|
|
| def __init__(self, path: str, device: str = "cpu"): |
| if ort is None: |
| raise ImportError("onnxruntime is required to run an exported policy (pip install onnxruntime).") |
| available = set(ort.get_available_providers()) |
| if device == "cuda" and "CUDAExecutionProvider" in available: |
| providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] |
| elif device in ("cpu", "cuda"): |
| providers = ["CPUExecutionProvider"] |
| if device == "cuda": |
| print("[WARN] ONNX Runtime has no CUDA provider; the 182-D policy will run on CPU.") |
| else: |
| raise ValueError(f"unsupported policy device {device!r}; expected 'cpu' or 'cuda'") |
| self.session = ort.InferenceSession(path, providers=providers) |
| inputs = [i.name for i in self.session.get_inputs()] |
| outputs = [o.name for o in self.session.get_outputs()] |
| if inputs != ["obs"] or "actions" not in outputs: |
| raise RuntimeError(f"expected an action-only policy with input ['obs'] and output 'actions', got {inputs}/{outputs}") |
| self.obs_dim = int(self.session.get_inputs()[0].shape[-1]) |
| self.act_dim = int([o for o in self.session.get_outputs() if o.name == "actions"][0].shape[-1]) |
| meta = self.session.get_modelmeta().custom_metadata_map |
| self.metadata = PolicyMetadata( |
| joint_names=_csv_strings(meta.get("joint_names")), |
| action_scale=_csv_floats(meta.get("action_scale")), |
| default_joint_pos=_csv_floats(meta.get("default_joint_pos")), |
| body_names=_csv_strings(meta.get("body_names")), |
| observation_names=_csv_strings(meta.get("observation_names")), |
| raw=dict(meta), |
| ) |
| if not self.metadata.joint_names: |
| raise RuntimeError("policy.onnx carries no joint_names metadata; export with attach_onnx_metadata") |
|
|
| def __call__(self, obs: np.ndarray) -> np.ndarray: |
| obs = np.asarray(obs, dtype=np.float32).reshape(-1) |
| if obs.size != self.obs_dim or not np.isfinite(obs).all(): |
| raise ValueError(f"policy observation must be finite shape ({self.obs_dim},), got {obs.shape}") |
| out = self.session.run(["actions"], {"obs": obs.reshape(1, -1)})[0] |
| action = out.reshape(-1).astype(np.float32) |
| if action.size != self.act_dim or not np.isfinite(action).all(): |
| raise RuntimeError(f"policy returned invalid action shape/values: {action.shape}") |
| return action |
|
|
|
|
| class HoldPosePolicy: |
| """Stand-in policy (zero actions = default pose) for pipeline self-tests without ONNX.""" |
|
|
| def __init__(self, num_actions: int, obs_dim: int, joint_names: list[str], default_joint_pos: list[float], action_scale: list[float]): |
| self.obs_dim = obs_dim |
| self.act_dim = num_actions |
| self.metadata = PolicyMetadata(joint_names, action_scale, default_joint_pos, [], [], {}) |
|
|
| def __call__(self, obs: np.ndarray) -> np.ndarray: |
| return np.zeros(self.act_dim, dtype=np.float32) |
|
|
