poliwrap 0.0.6

A lightweight, framework-agnostic abstraction for reinforcement learning policy inference.

PoliWrap

A lightweight, framework-agnostic wrapper for reinforcement learning policy inference.

Installation

pip install poliwrap            # core only (numpy)
pip install poliwrap[onnx]      # + ONNX Runtime backend (CPU)
pip install poliwrap[onnx-gpu]  # + ONNX Runtime backend (CUDA)
pip install poliwrap[torch]     # + PyTorch backend
pip install poliwrap[all]       # all backends (CPU ONNX)

Each wrapper is importable only when its backend is installed.

Usage

ONNX

from poliwrap import ONNXPolicyWrapper

policy = ONNXPolicyWrapper("policy.onnx")
actions = policy({"actor_obs": obs})  # obs: dict[str, np.ndarray]

PyTorch

from poliwrap import PytorchPolicyWrapper

policy = PytorchPolicyWrapper("policy.pt")
actions = policy({"actor_obs": obs})  # obs: dict[str, torch.Tensor]

Extending

All wrappers subclass PolicyWrapper. To support a new backend, implement _load_model and __call__:

from poliwrap import PolicyWrapper


class MyPolicyWrapper(PolicyWrapper):
    def _load_model(self) -> None:
        self.model = ...  # load from self.model_path

    def __call__(self, observations):
        return self.model(observations)

PolicyWrapper also has overridable preprocess_observations, postprocess_actions, and reset hooks for input/output transforms and stateful (e.g. RNN) policies.

License

MIT