File size: 3,946 Bytes
d8290eb | 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 | """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: # pragma: no cover - optional at import time
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] # Isaac articulation order, with "_joint" suffix
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)
|