52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
|
"""Common inference adapter abstraction.
|
||
|
|
|
||
|
|
The adapter owns all architecture-specific model setup and exposes a single
|
||
|
|
``predict_batch`` entry point. Environment drivers pass payloads through without
|
||
|
|
interpreting the env-adapter schema.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from abc import ABC, abstractmethod
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from wall_x._vendor.harrix.eval_config import EvalConfig
|
||
|
|
|
||
|
|
|
||
|
|
class BaseInferAdapter(ABC):
|
||
|
|
@abstractmethod
|
||
|
|
def __init__(self, cfg: EvalConfig) -> None:
|
||
|
|
"""Load processors, model config, checkpoints, and normalizers.
|
||
|
|
|
||
|
|
Subclasses should reject unsupported ``cfg.model.action_mode`` values
|
||
|
|
during construction.
|
||
|
|
"""
|
||
|
|
|
||
|
|
@property
|
||
|
|
@abstractmethod
|
||
|
|
def chunk_horizon(self) -> int:
|
||
|
|
"""Number of action steps returned by each ``predict_batch`` call."""
|
||
|
|
|
||
|
|
@property
|
||
|
|
@abstractmethod
|
||
|
|
def action_mode(self) -> str:
|
||
|
|
"""Configured inference algorithm, fixed at construction time."""
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def predict_batch(self, payloads: list[dict]) -> list[np.ndarray]:
|
||
|
|
"""Run one batched inference call and return one chunk per payload.
|
||
|
|
|
||
|
|
``payloads[i]`` schema is defined by the env-adapter pair:
|
||
|
|
{
|
||
|
|
"observation": dict[str, np.ndarray],
|
||
|
|
# Env-defined ndarray bundle, for example LIBERO:
|
||
|
|
# {"eef_pos":(3,), "eef_axisangle":(3,), "gripper":(1,),
|
||
|
|
# "face_view":(H,W,3), "wrist_view":(H,W,3)}
|
||
|
|
"instruction": str,
|
||
|
|
"noise": np.ndarray | None,
|
||
|
|
# Flow may pass (chunk_horizon, action_dim); other modes
|
||
|
|
# may leave this as None.
|
||
|
|
}
|
||
|
|
"""
|