feat(rollout): remote inference draft

tests/policy_server/test_codec.py

@@ -0,0 +1,412 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the MessagePack wire codecs (tensors, images, messages)."""
+
+import numpy as np
+import pytest
+
+msgpack = pytest.importorskip("msgpack")
+
+from lerobot.policy_server.codec import (  # noqa: E402
+    decode_action_chunk,
+    decode_image,
+    decode_observation,
+    decode_raw,
+    decode_reset,
+    decode_reset_ack,
+    decode_session_ack,
+    decode_session_close,
+    decode_session_open,
+    decode_status,
+    decode_tensor,
+    encode_action_chunk,
+    encode_image,
+    encode_observation,
+    encode_reset,
+    encode_reset_ack,
+    encode_session_ack,
+    encode_session_close,
+    encode_session_open,
+    encode_status,
+    encode_tensor,
+)
+from lerobot.policy_server.schema import (  # noqa: E402
+    IMAGE_CODEC_JPEG,
+    IMAGE_CODEC_RAW,
+    ActionChunkMsg,
+    ObservationMsg,
+    ResetAckMsg,
+    ResetMsg,
+    SessionAckMsg,
+    SessionCloseMsg,
+    SessionOpenMsg,
+    StatusMsg,
+)
+
+# ---------------------------------------------------------------------------
+# Tensor codec
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "arr",
+    [
+        np.array([1.5, -2.25, 3.0], dtype=np.float32),
+        np.arange(12, dtype=np.float64).reshape(3, 4),
+        np.array([[1, -2], [3, 4]], dtype=np.int64),
+        np.array([[True, False], [False, True]], dtype=np.bool_),
+        np.zeros((0,), dtype=np.float32),  # empty 1-d
+        np.zeros((0, 6), dtype=np.float64),  # empty 2-d
+        np.arange(24, dtype=np.int64).reshape(2, 3, 4),
+    ],
+    ids=["f32_1d", "f64_2d", "i64_2d", "bool_2d", "f32_empty", "f64_empty_2d", "i64_3d"],
+)
+def test_tensor_roundtrip(arr):
+    out = decode_tensor(encode_tensor(arr))
+    assert out.dtype == arr.dtype
+    assert out.shape == arr.shape
+    np.testing.assert_array_equal(out, arr)
+
+
+def test_tensor_roundtrip_0d_preserves_value():
+    # KNOWN QUIRK: np.ascontiguousarray inside encode_tensor promotes
+    # 0-d arrays to shape (1,), so the round-trip is value-preserving
+    # but not shape-preserving for scalars.
+    arr = np.array(3.5, dtype=np.float32)
+    out = decode_tensor(encode_tensor(arr))
+    assert out.dtype == arr.dtype
+    assert out.shape in ((), (1,))
+    assert float(np.squeeze(out)) == 3.5
+
+
+def test_tensor_none_passthrough():
+    assert encode_tensor(None) is None
+    assert decode_tensor(None) is None
+
+
+def test_tensor_big_endian_input_values_identical():
+    be = np.array([1.0, 2.5, -3.75], dtype=">f4")
+    enc = encode_tensor(be)
+    assert np.dtype(enc["dtype"]).byteorder != ">"
+    out = decode_tensor(enc)
+    np.testing.assert_array_equal(out, be.astype("<f4"))
+    np.testing.assert_array_equal(out, np.array([1.0, 2.5, -3.75], dtype=np.float32))
+
+
+def test_tensor_decoded_writable_and_contiguous():
+    arr = np.arange(6, dtype=np.float32).reshape(2, 3)
+    out = decode_tensor(encode_tensor(arr))
+    assert out.flags.writeable
+    assert out.flags.c_contiguous
+    out[0, 0] = 99.0  # must not raise
+    assert out[0, 0] == 99.0
+
+
+def test_tensor_decode_refuses_object_dtype():
+    with pytest.raises(ValueError, match="object dtype"):
+        decode_tensor({"dtype": "|O", "shape": [1], "data": b"\x00" * 8})
+
+
+def test_tensor_roundtrip_through_msgpack():
+    arr = np.arange(10, dtype=np.float32)
+    packed = msgpack.packb(encode_tensor(arr), use_bin_type=True)
+    out = decode_tensor(msgpack.unpackb(packed, raw=False))
+    np.testing.assert_array_equal(out, arr)
+
+
+# ---------------------------------------------------------------------------
+# Image codec
+# ---------------------------------------------------------------------------
+
+
+def _gradient_image(h: int = 32, w: int = 48) -> np.ndarray:
+    """Smooth RGB gradient: JPEG-friendly, deterministic."""
+    img = np.zeros((h, w, 3), dtype=np.uint8)
+    img[..., 0] = np.linspace(0, 255, w, dtype=np.uint8)[None, :]
+    img[..., 1] = np.linspace(0, 255, h, dtype=np.uint8)[:, None]
+    img[..., 2] = 128
+    return img
+
+
+def test_image_raw_roundtrip_byte_exact():
+    img = _gradient_image()
+    enc = encode_image(img, jpeg_quality=0)
+    assert enc["codec"] == IMAGE_CODEC_RAW
+    out = decode_image(enc)
+    assert out.dtype == np.uint8
+    assert out.shape == img.shape
+    np.testing.assert_array_equal(out, img)
+
+
+def test_image_jpeg_roundtrip_approximately_equal():
+    img = _gradient_image()
+    enc = encode_image(img, jpeg_quality=95)
+    assert enc["codec"] == IMAGE_CODEC_JPEG
+    out = decode_image(enc)
+    assert out.dtype == np.uint8
+    assert out.shape == img.shape
+    err = np.abs(out.astype(np.int32) - img.astype(np.int32)).mean()
+    assert err < 5.0, f"JPEG round-trip too lossy: mean abs error {err}"
+
+
+def test_image_jpeg_rgb_order_regression_pure_red_stays_red():
+    # A silent BGR swap would poison every VLA in a fleet: pure red must
+    # come back red-dominant, not blue-dominant.
+    img = np.zeros((32, 32, 3), dtype=np.uint8)
+    img[..., 0] = 255  # RGB: red channel
+    out = decode_image(encode_image(img, jpeg_quality=90))
+    red_mean = out[..., 0].astype(np.float64).mean()
+    blue_mean = out[..., 2].astype(np.float64).mean()
+    assert red_mean > 200, f"red channel lost: mean {red_mean}"
+    assert blue_mean < 50, f"blue channel gained: mean {blue_mean}"
+    assert red_mean > blue_mean
+
+
+def test_encode_image_rejects_float_arrays():
+    with pytest.raises(ValueError, match="uint8 HWC RGB"):
+        encode_image(np.zeros((8, 8, 3), dtype=np.float32))
+
+
+@pytest.mark.parametrize(
+    "shape", [(3, 16, 24), (16, 24), (16, 24, 4), (16, 24, 1)], ids=["chw", "hw", "hwc4", "hwc1"]
+)
+def test_encode_image_rejects_non_hwc(shape):
+    with pytest.raises(ValueError, match="uint8 HWC RGB"):
+        encode_image(np.zeros(shape, dtype=np.uint8))
+
+
+def test_decode_image_rejects_unknown_codec():
+    with pytest.raises(ValueError, match="Unknown image codec"):
+        decode_image({"codec": "webp", "data": b""})
+
+
+# ---------------------------------------------------------------------------
+# Data-plane messages
+# ---------------------------------------------------------------------------
+
+
+def test_observation_roundtrip_full():
+    rng = np.random.default_rng(0)
+    state = np.array([0.1, -0.2, 0.3, 0.4], dtype=np.float32)
+    front = rng.integers(0, 255, size=(16, 24, 3), dtype=np.uint8)
+    wrist = rng.integers(0, 255, size=(8, 12, 3), dtype=np.uint8)
+    prefix_model = rng.standard_normal((5, 4)).astype(np.float32)
+    prefix_robot = (prefix_model * 2.0).astype(np.float32)
+    msg = ObservationMsg(
+        state=state,
+        images={"front": front, "wrist": wrist},
+        task="fold the towel",
+        inference_delay_steps=3,
+        prefix_model=prefix_model,
+        prefix_robot=prefix_robot,
+        episode_start=True,
+        jpeg_quality=0,  # raw: byte-exact images
+    )
+    out = decode_observation(encode_observation(msg))
+    np.testing.assert_array_equal(out.state, state)
+    assert set(out.images) == {"front", "wrist"}
+    np.testing.assert_array_equal(out.images["front"], front)
+    np.testing.assert_array_equal(out.images["wrist"], wrist)
+    assert out.task == "fold the towel"
+    assert out.inference_delay_steps == 3
+    np.testing.assert_array_equal(out.prefix_model, prefix_model)
+    np.testing.assert_array_equal(out.prefix_robot, prefix_robot)
+    assert out.episode_start is True
+
+
+def test_observation_roundtrip_minimal_defaults():
+    out = decode_observation(encode_observation(ObservationMsg()))
+    assert out.state is None
+    assert out.images == {}
+    assert out.task == ""
+    assert out.inference_delay_steps == 0
+    assert out.prefix_model is None
+    assert out.prefix_robot is None
+    assert out.episode_start is False
+
+
+def test_action_chunk_roundtrip():
+    chunk_model = np.arange(12, dtype=np.float32).reshape(3, 4)
+    chunk_robot = chunk_model * 2.0
+    msg = ActionChunkMsg(
+        seq_id_echo=17,
+        client_mono_ns_echo=123456789,
+        episode_id_echo=2,
+        chunk_model=chunk_model,
+        chunk_robot=chunk_robot,
+        queue_wait_ms=1.5,
+        inference_ms=12.25,
+        superseded_seqs=4,
+        server_load=0.75,
+    )
+    out = decode_action_chunk(encode_action_chunk(msg))
+    assert out.seq_id_echo == 17
+    assert out.client_mono_ns_echo == 123456789
+    assert out.episode_id_echo == 2
+    np.testing.assert_array_equal(out.chunk_model, chunk_model)
+    np.testing.assert_array_equal(out.chunk_robot, chunk_robot)
+    assert out.queue_wait_ms == 1.5
+    assert out.inference_ms == 12.25
+    assert out.superseded_seqs == 4
+    assert out.server_load == 0.75
+
+
+# ---------------------------------------------------------------------------
+# Control-plane messages
+# ---------------------------------------------------------------------------
+
+
+def test_session_open_roundtrip():
+    msg = SessionOpenMsg(
+        client_uuid="uuid-1",
+        robot_type="so101_follower",
+        policy_type="pi0",
+        fps=30.0,
+        action_names=["j0.pos", "j1.pos"],
+        camera_names=["front", "wrist"],
+        state_dim=6,
+        rtc_enabled=True,
+        task="fold",
+        tags={"site": "lab-3"},
+    )
+    out = decode_session_open(encode_session_open(msg))
+    assert out == msg
+
+
+def test_session_ack_roundtrip():
+    msg = SessionAckMsg(
+        accepted=True,
+        warnings=["fps mismatch"],
+        session_id="sess-1",
+        model_repo="org/model",
+        model_revision="main",
+        policy_type="pi0",
+        action_names=["j0.pos"],
+        expected_cameras=["front"],
+        state_dim=6,
+        chunk_size=50,
+        trained_fps=30.0,
+        supports_rtc=True,
+        rtc_execution_horizon=25,
+        serving_mode="shared",
+        warmed_up=True,
+        server_load=0.5,
+    )
+    out = decode_session_ack(encode_session_ack(msg))
+    assert out == msg
+
+
+def test_status_roundtrip():
+    msg = StatusMsg(
+        model_repo="org/model",
+        model_revision="abc123",
+        policy_type="act",
+        action_names=["j0.pos", "j1.pos"],
+        expected_cameras=["front"],
+        state_dim=6,
+        chunk_size=100,
+        trained_fps=30.0,
+        supports_rtc=False,
+        rtc_execution_horizon=0,
+        serving_mode="exclusive",
+        warmed_up=False,
+        active_sessions=2,
+        max_sessions=4,
+    )
+    out = decode_status(encode_status(msg))
+    assert out == msg
+
+
+def test_reset_and_reset_ack_roundtrip():
+    out = decode_reset(encode_reset(ResetMsg(client_uuid="uuid-1", episode_id=5)))
+    assert out == ResetMsg(client_uuid="uuid-1", episode_id=5)
+    out_ack = decode_reset_ack(encode_reset_ack(ResetAckMsg(ok=False, error="busy")))
+    assert out_ack == ResetAckMsg(ok=False, error="busy")
+
+
+def test_session_close_roundtrip():
+    msg = SessionCloseMsg(client_uuid="uuid-1", session_id="sess-1")
+    out = decode_session_close(encode_session_close(msg))
+    assert out == msg
+
+
+# ---------------------------------------------------------------------------
+# Schema evolution (additive-only contract)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("encoded", "decoder", "expected"),
+    [
+        (
+            encode_session_ack(SessionAckMsg(accepted=True, session_id="s")),
+            decode_session_ack,
+            SessionAckMsg(accepted=True, session_id="s"),
+        ),
+        (
+            encode_reset(ResetMsg(client_uuid="u", episode_id=1)),
+            decode_reset,
+            ResetMsg(client_uuid="u", episode_id=1),
+        ),
+        (
+            encode_session_open(SessionOpenMsg(client_uuid="u")),
+            decode_session_open,
+            SessionOpenMsg(client_uuid="u"),
+        ),
+    ],
+    ids=["session_ack", "reset", "session_open"],
+)
+def test_unknown_keys_ignored(encoded, decoder, expected):
+    obj = msgpack.unpackb(encoded, raw=False)
+    obj["a_future_field"] = {"nested": [1, 2, 3]}
+    out = decoder(msgpack.packb(obj, use_bin_type=True))
+    assert out == expected
+
+
+def test_missing_optional_keys_take_defaults():
+    minimal = msgpack.packb({"accepted": True}, use_bin_type=True)
+    out = decode_session_ack(minimal)
+    assert out.accepted is True
+    assert out.error == ""
+    assert out.warnings == []
+    assert out.chunk_size == 0
+    assert out.server_load == 0.0
+
+    out_chunk = decode_action_chunk(msgpack.packb({"seq_id_echo": 9}, use_bin_type=True))
+    assert out_chunk.seq_id_echo == 9
+    assert out_chunk.chunk_model is None
+    assert out_chunk.chunk_robot is None
+    assert out_chunk.queue_wait_ms == 0.0
+
+    out_obs = decode_observation(msgpack.packb({"task": "t"}, use_bin_type=True))
+    assert out_obs.task == "t"
+    assert out_obs.state is None
+    assert out_obs.images == {}
+    assert out_obs.episode_start is False
+
+
+# ---------------------------------------------------------------------------
+# decode_raw
+# ---------------------------------------------------------------------------
+
+
+def test_decode_raw_returns_plain_dict_with_op():
+    open_obj = decode_raw(encode_session_open(SessionOpenMsg(client_uuid="u")))
+    assert isinstance(open_obj, dict)
+    assert open_obj["op"] == "open"
+
+    close_obj = decode_raw(encode_session_close(SessionCloseMsg(client_uuid="u")))
+    assert isinstance(close_obj, dict)
+    assert close_obj["op"] == "close"