protocol.py

"""Shared serialization for robot observations and actions.

Uses msgpack with numpy support for compact, fast binary encoding.
Both server.py and client.py import from here.
"""

import msgpack
import msgpack_numpy as m
import numpy as np

# Patch msgpack to handle numpy dtypes natively
m.patch()


def pack_obs(obs: dict[str, np.ndarray]) -> bytes:
    """Serialize observation dict (str -> numpy array) to msgpack bytes."""
    return msgpack.packb(obs, default=m.encode)


def unpack_obs(data: bytes) -> dict[str, np.ndarray]:
    """Deserialize msgpack bytes to observation dict."""
    return msgpack.unpackb(data, object_hook=m.decode, raw=False)


def pack_response(action: np.ndarray, inference_ms: float) -> bytes:
    """Serialize action + server timing metadata."""
    return msgpack.packb(
        {"action": action, "inference_ms": inference_ms},
        default=m.encode,
    )


def unpack_response(data: bytes) -> tuple[np.ndarray, float]:
    """Deserialize action response. Returns (action_array, inference_ms)."""
    resp = msgpack.unpackb(data, object_hook=m.decode, raw=False)
    return resp["action"], resp["inference_ms"]