admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 01:41:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(rollout): remote inference draft

Browse Source
This commit is contained in:
Steven Palma
2026-06-12 02:01:41 +02:00
parent 87242cfced
commit fc019d3902
54 changed files with 7883 additions and 3392 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/async_inference/test_e2e.py
-187
View File
@@ -1,187 +0,0 @@
# Copyright 2025 The HuggingFace Inc. team.
#
# 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.
"""End-to-end test of the asynchronous inference stack (client ↔ server).
This test spins up a lightweight gRPC `PolicyServer` instance with a stubbed
policy network and launches a `RobotClient` that uses a `MockRobot`. The goal
is to exercise the full communication loop:
1. Client sends policy specification → Server
2. Client streams observations → Server
3. Server streams action chunks → Client
4. Client executes received actions
The test succeeds if at least one action is executed and the server records at
least one predicted timestep - demonstrating that the gRPC round-trip works
end-to-end using real (but lightweight) protocol messages.
"""
from __future__ import annotations
import threading
from concurrent import futures
import pytest
import torch
# Skip entire module if required deps are not available
pytest.importorskip("grpc")
pytest.importorskip("serial", reason="pyserial is required (install lerobot[hardware])")
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
# -----------------------------------------------------------------------------
# End-to-end test
# -----------------------------------------------------------------------------
def test_async_inference_e2e(monkeypatch):
"""Tests the full asynchronous inference pipeline."""
# Import grpc-dependent modules inside the test function
import grpc
from lerobot.async_inference.configs import PolicyServerConfig, RobotClientConfig
from lerobot.async_inference.helpers import map_robot_keys_to_lerobot_features
from lerobot.async_inference.policy_server import PolicyServer
from lerobot.async_inference.robot_client import RobotClient
from lerobot.robots.utils import make_robot_from_config
from lerobot.transport import (
services_pb2, # type: ignore
services_pb2_grpc, # type: ignore
)
from tests.mocks.mock_robot import MockRobotConfig
# Create a stub policy similar to test_policy_server.py
class MockPolicy:
"""A minimal mock for an actual policy, returning zeros."""
class _Config:
robot_type = "dummy_robot"
@property
def image_features(self):
"""Empty image features since this test doesn't use images."""
return {}
def __init__(self):
self.config = self._Config()
def to(self, *args, **kwargs):
return self
def model(self, batch):
# Return a chunk of 20 dummy actions.
batch_size = len(batch["robot_type"])
return torch.zeros(batch_size, 20, 6)
# ------------------------------------------------------------------
# 1. Create PolicyServer instance with mock policy
# ------------------------------------------------------------------
policy_server_config = PolicyServerConfig(host="localhost", port=9999)
policy_server = PolicyServer(policy_server_config)
# Replace the real policy with our fast, deterministic stub.
policy_server.policy = MockPolicy()
policy_server.actions_per_chunk = 20
policy_server.device = "cpu"
# NOTE(Steven): Smelly tests as the Server is a state machine being partially mocked. Adding these processors as a quick fix.
policy_server.preprocessor = lambda obs: obs
policy_server.postprocessor = lambda tensor: tensor
# Set up robot config and features
robot_config = MockRobotConfig()
mock_robot = make_robot_from_config(robot_config)
lerobot_features = map_robot_keys_to_lerobot_features(mock_robot)
policy_server.lerobot_features = lerobot_features
# Force server to produce deterministic action chunks in test mode
policy_server.policy_type = "act"
def _fake_get_action_chunk(_self, _obs, _type="test"):
action_dim = 6
batch_size = 1
actions_per_chunk = policy_server.actions_per_chunk
return torch.zeros(batch_size, actions_per_chunk, action_dim)
monkeypatch.setattr(PolicyServer, "_get_action_chunk", _fake_get_action_chunk, raising=True)
# Bypass potentially heavy model loading inside SendPolicyInstructions
def _fake_send_policy_instructions(self, request, context): # noqa: N802
return services_pb2.Empty()
monkeypatch.setattr(PolicyServer, "SendPolicyInstructions", _fake_send_policy_instructions, raising=True)
# Build gRPC server running a PolicyServer
server = grpc.server(futures.ThreadPoolExecutor(max_workers=1, thread_name_prefix="policy_server"))
services_pb2_grpc.add_AsyncInferenceServicer_to_server(policy_server, server)
# Use the host/port specified in the fixture's config
server_address = f"{policy_server.config.host}:{policy_server.config.port}"
server.add_insecure_port(server_address)
server.start()
# ------------------------------------------------------------------
# 2. Create a RobotClient around the MockRobot
# ------------------------------------------------------------------
client_config = RobotClientConfig(
server_address=server_address,
robot=robot_config,
chunk_size_threshold=0.0,
policy_type="test",
pretrained_name_or_path="test",
actions_per_chunk=20,
)
client = RobotClient(client_config)
assert client.start(), "Client failed initial handshake with the server"
# Track action chunks received and verify device type
action_chunks_received = {"count": 0, "actions_on_cpu": True}
original_aggregate = client._aggregate_action_queues
def counting_aggregate(*args, **kwargs):
action_chunks_received["count"] += 1
# Check that all received actions are on CPU
if args:
for timed_action in args[0]: # args[0] is the list of TimedAction
action_tensor = timed_action.get_action()
if action_tensor.device.type != "cpu":
action_chunks_received["actions_on_cpu"] = False
return original_aggregate(*args, **kwargs)
monkeypatch.setattr(client, "_aggregate_action_queues", counting_aggregate)
# Start client threads
action_thread = threading.Thread(target=client.receive_actions, daemon=True)
control_thread = threading.Thread(target=client.control_loop, args=({"task": ""}), daemon=True)
action_thread.start()
control_thread.start()
# ------------------------------------------------------------------
# 3. System exchanges a few messages
# ------------------------------------------------------------------
# Wait for 5 seconds
server.wait_for_termination(timeout=5)
assert action_chunks_received["count"] > 0, "Client did not receive any action chunks"
assert len(policy_server._predicted_timesteps) > 0, "Server did not record any predicted timesteps"
# ------------------------------------------------------------------
# 4. Stop the system
# ------------------------------------------------------------------
client.stop()
action_thread.join()
control_thread.join()
policy_server.stop()
server.stop(grace=None)
tests/async_inference/test_helpers.py
-454
View File
@@ -1,454 +0,0 @@
# Copyright 2025 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.
import math
import pickle
import time
import pytest
pytest.importorskip("grpc")
import numpy as np # noqa: E402
import torch # noqa: E402
from lerobot.async_inference.helpers import ( # noqa: E402
FPSTracker,
TimedAction,
TimedObservation,
observations_similar,
prepare_image,
prepare_raw_observation,
raw_observation_to_observation,
resize_robot_observation_image,
)
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
# ---------------------------------------------------------------------
# FPSTracker
# ---------------------------------------------------------------------
def test_fps_tracker_first_observation():
"""First observation should initialize timestamp and return 0 FPS."""
tracker = FPSTracker(target_fps=30.0)
timestamp = 1000.0
metrics = tracker.calculate_fps_metrics(timestamp)
assert tracker.first_timestamp == timestamp
assert tracker.total_obs_count == 1
assert metrics["avg_fps"] == 0.0
assert metrics["target_fps"] == 30.0
def test_fps_tracker_single_interval():
"""Two observations 1 second apart should give 1 FPS."""
tracker = FPSTracker(target_fps=30.0)
# First observation at t=0
metrics1 = tracker.calculate_fps_metrics(0.0)
assert metrics1["avg_fps"] == 0.0
# Second observation at t=1 (1 second later)
metrics2 = tracker.calculate_fps_metrics(1.0)
expected_fps = 1.0 # (2-1) observations / 1.0 seconds = 1 FPS
assert math.isclose(metrics2["avg_fps"], expected_fps, rel_tol=1e-6)
def test_fps_tracker_multiple_intervals():
"""Multiple observations should calculate correct average FPS."""
tracker = FPSTracker(target_fps=30.0)
# Simulate 5 observations over 2 seconds (should be 2 FPS average)
timestamps = [0.0, 0.5, 1.0, 1.5, 2.0]
for i, ts in enumerate(timestamps):
metrics = tracker.calculate_fps_metrics(ts)
if i == 0:
assert metrics["avg_fps"] == 0.0
elif i == len(timestamps) - 1:
# After 5 observations over 2 seconds: (5-1)/2 = 2 FPS
expected_fps = 2.0
assert math.isclose(metrics["avg_fps"], expected_fps, rel_tol=1e-6)
def test_fps_tracker_irregular_intervals():
"""FPS calculation should work with irregular time intervals."""
tracker = FPSTracker(target_fps=30.0)
# Irregular timestamps: 0, 0.1, 0.5, 2.0, 3.0 seconds
timestamps = [0.0, 0.1, 0.5, 2.0, 3.0]
for ts in timestamps:
metrics = tracker.calculate_fps_metrics(ts)
# 5 observations over 3 seconds: (5-1)/3 = 1.333... FPS
expected_fps = 4.0 / 3.0
assert math.isclose(metrics["avg_fps"], expected_fps, rel_tol=1e-6)
# ---------------------------------------------------------------------
# TimedData helpers
# ---------------------------------------------------------------------
def test_timed_action_getters():
"""TimedAction stores & returns timestamp, action tensor and timestep."""
ts = time.time()
action = torch.arange(10)
ta = TimedAction(timestamp=ts, action=action, timestep=0)
assert math.isclose(ta.get_timestamp(), ts, rel_tol=0, abs_tol=1e-6)
torch.testing.assert_close(ta.get_action(), action)
assert ta.get_timestep() == 0
def test_timed_observation_getters():
"""TimedObservation stores & returns timestamp, dict and timestep."""
ts = time.time()
obs_dict = {OBS_STATE: torch.ones(6)}
to = TimedObservation(timestamp=ts, observation=obs_dict, timestep=0)
assert math.isclose(to.get_timestamp(), ts, rel_tol=0, abs_tol=1e-6)
assert to.get_observation() is obs_dict
assert to.get_timestep() == 0
def test_timed_data_deserialization_data_getters():
"""TimedAction / TimedObservation survive a round-trip through ``pickle``.
The async-inference stack uses ``pickle.dumps`` to move these objects across
the gRPC boundary (see RobotClient.send_observation and PolicyServer.StreamActions).
This test ensures that the payload keeps its content intact after
the (de)serialization round-trip.
"""
ts = time.time()
# ------------------------------------------------------------------
# TimedAction
# ------------------------------------------------------------------
original_action = torch.randn(6)
ta_in = TimedAction(timestamp=ts, action=original_action, timestep=13)
# Serialize → bytes → deserialize
ta_bytes = pickle.dumps(ta_in) # nosec
ta_out: TimedAction = pickle.loads(ta_bytes) # nosec B301
# Identity & content checks
assert math.isclose(ta_out.get_timestamp(), ts, rel_tol=0, abs_tol=1e-6)
assert ta_out.get_timestep() == 13
torch.testing.assert_close(ta_out.get_action(), original_action)
# ------------------------------------------------------------------
# TimedObservation
# ------------------------------------------------------------------
obs_dict = {OBS_STATE: torch.arange(4).float()}
to_in = TimedObservation(timestamp=ts, observation=obs_dict, timestep=7, must_go=True)
to_bytes = pickle.dumps(to_in) # nosec
to_out: TimedObservation = pickle.loads(to_bytes) # nosec B301
assert math.isclose(to_out.get_timestamp(), ts, rel_tol=0, abs_tol=1e-6)
assert to_out.get_timestep() == 7
assert to_out.must_go is True
assert to_out.get_observation().keys() == obs_dict.keys()
torch.testing.assert_close(to_out.get_observation()[OBS_STATE], obs_dict[OBS_STATE])
# ---------------------------------------------------------------------
# observations_similar()
# ---------------------------------------------------------------------
def _make_obs(state: torch.Tensor) -> TimedObservation:
"""Create a TimedObservation with raw robot observation format."""
return TimedObservation(
timestamp=time.time(),
observation={
"shoulder": state[0].item() if len(state) > 0 else 0.0,
"elbow": state[1].item() if len(state) > 1 else 0.0,
"wrist": state[2].item() if len(state) > 2 else 0.0,
"gripper": state[3].item() if len(state) > 3 else 0.0,
},
timestep=0,
)
def test_observations_similar_true():
"""Distance below atol → observations considered similar."""
# Create mock lerobot features for the similarity check
lerobot_features = {
OBS_STATE: {
"dtype": "float32",
"shape": [4],
"names": ["shoulder", "elbow", "wrist", "gripper"],
}
}
obs1 = _make_obs(torch.zeros(4))
obs2 = _make_obs(0.5 * torch.ones(4))
assert observations_similar(obs1, obs2, lerobot_features, atol=2.0)
obs3 = _make_obs(2.0 * torch.ones(4))
assert not observations_similar(obs1, obs3, lerobot_features, atol=2.0)
# ---------------------------------------------------------------------
# raw_observation_to_observation and helpers
# ---------------------------------------------------------------------
def _create_mock_robot_observation():
"""Create a mock robot observation with motor positions and camera images."""
return {
"shoulder": 1.0,
"elbow": 2.0,
"wrist": 3.0,
"gripper": 0.5,
"laptop": np.random.randint(0, 256, size=(480, 640, 3), dtype=np.uint8),
"phone": np.random.randint(0, 256, size=(480, 640, 3), dtype=np.uint8),
}
def _create_mock_lerobot_features():
"""Create mock lerobot features mapping similar to what hw_to_dataset_features returns."""
return {
OBS_STATE: {
"dtype": "float32",
"shape": [4],
"names": ["shoulder", "elbow", "wrist", "gripper"],
},
f"{OBS_IMAGES}.laptop": {
"dtype": "image",
"shape": [480, 640, 3],
"names": ["height", "width", "channels"],
},
f"{OBS_IMAGES}.phone": {
"dtype": "image",
"shape": [480, 640, 3],
"names": ["height", "width", "channels"],
},
}
def _create_mock_policy_image_features():
"""Create mock policy image features with different resolutions."""
return {
f"{OBS_IMAGES}.laptop": PolicyFeature(
type=FeatureType.VISUAL,
shape=(3, 224, 224), # Policy expects smaller resolution
),
f"{OBS_IMAGES}.phone": PolicyFeature(
type=FeatureType.VISUAL,
shape=(3, 160, 160), # Different resolution for second camera
),
}
def test_prepare_image():
"""Test image preprocessing: int8 → float32, normalization to [0,1]."""
# Create mock int8 image data
image_int8 = torch.randint(0, 256, size=(3, 224, 224), dtype=torch.uint8)
processed = prepare_image(image_int8)
# Check dtype conversion
assert processed.dtype == torch.float32
# Check normalization range
assert processed.min() >= 0.0
assert processed.max() <= 1.0
# Check that values are scaled correctly (255 → 1.0, 0 → 0.0)
if image_int8.max() == 255:
assert torch.isclose(processed.max(), torch.tensor(1.0), atol=1e-6)
if image_int8.min() == 0:
assert torch.isclose(processed.min(), torch.tensor(0.0), atol=1e-6)
# Check memory contiguity
assert processed.is_contiguous()
def test_resize_robot_observation_image():
"""Test image resizing from robot resolution to policy resolution."""
# Create mock image: (H=480, W=640, C=3)
original_image = torch.randint(0, 256, size=(480, 640, 3), dtype=torch.uint8)
target_shape = (3, 224, 224) # (C, H, W)
resized = resize_robot_observation_image(original_image, target_shape)
# Check output shape matches target
assert resized.shape == target_shape
# Check that original image had different dimensions
assert original_image.shape != resized.shape
# Check that resizing preserves value range
assert resized.min() >= 0
assert resized.max() <= 255
def test_prepare_raw_observation():
"""Test the preparation of raw robot observation to lerobot format."""
robot_obs = _create_mock_robot_observation()
lerobot_features = _create_mock_lerobot_features()
policy_image_features = _create_mock_policy_image_features()
prepared = prepare_raw_observation(robot_obs, lerobot_features, policy_image_features)
# Check that state is properly extracted and batched
assert OBS_STATE in prepared
state = prepared[OBS_STATE]
assert isinstance(state, torch.Tensor)
assert state.shape == (1, 4) # Batched state
# Check that images are processed and resized
assert f"{OBS_IMAGES}.laptop" in prepared
assert f"{OBS_IMAGES}.phone" in prepared
laptop_img = prepared[f"{OBS_IMAGES}.laptop"]
phone_img = prepared[f"{OBS_IMAGES}.phone"]
# Check image shapes match policy requirements
assert laptop_img.shape == policy_image_features[f"{OBS_IMAGES}.laptop"].shape
assert phone_img.shape == policy_image_features[f"{OBS_IMAGES}.phone"].shape
# Check that images are tensors
assert isinstance(laptop_img, torch.Tensor)
assert isinstance(phone_img, torch.Tensor)
def test_raw_observation_to_observation_basic():
"""Test the main raw_observation_to_observation function."""
robot_obs = _create_mock_robot_observation()
lerobot_features = _create_mock_lerobot_features()
policy_image_features = _create_mock_policy_image_features()
observation = raw_observation_to_observation(robot_obs, lerobot_features, policy_image_features)
# Check that all expected keys are present
assert OBS_STATE in observation
assert f"{OBS_IMAGES}.laptop" in observation
assert f"{OBS_IMAGES}.phone" in observation
# Check state processing
state = observation[OBS_STATE]
assert isinstance(state, torch.Tensor)
assert state.shape == (1, 4) # Batched
# Check image processing
laptop_img = observation[f"{OBS_IMAGES}.laptop"]
phone_img = observation[f"{OBS_IMAGES}.phone"]
# Images should have batch dimension: (B, C, H, W)
assert laptop_img.shape == (1, 3, 224, 224)
assert phone_img.shape == (1, 3, 160, 160)
# Check image dtype and range (should be float32 in [0, 1])
assert laptop_img.dtype == torch.float32
assert phone_img.dtype == torch.float32
assert laptop_img.min() >= 0.0 and laptop_img.max() <= 1.0
assert phone_img.min() >= 0.0 and phone_img.max() <= 1.0
def test_raw_observation_to_observation_with_non_tensor_data():
"""Test that non-tensor data (like task strings) is preserved."""
robot_obs = _create_mock_robot_observation()
robot_obs["task"] = "pick up the red cube" # Add string instruction
lerobot_features = _create_mock_lerobot_features()
policy_image_features = _create_mock_policy_image_features()
observation = raw_observation_to_observation(robot_obs, lerobot_features, policy_image_features)
# Check that task string is preserved
assert "task" in observation
assert observation["task"] == "pick up the red cube"
assert isinstance(observation["task"], str)
@torch.no_grad()
def test_raw_observation_to_observation_device_handling():
"""Test that tensors are created (device placement is handled by preprocessor)."""
robot_obs = _create_mock_robot_observation()
lerobot_features = _create_mock_lerobot_features()
policy_image_features = _create_mock_policy_image_features()
observation = raw_observation_to_observation(robot_obs, lerobot_features, policy_image_features)
# Check that all expected keys produce tensors (device placement handled by preprocessor later)
for key, value in observation.items():
if isinstance(value, torch.Tensor):
assert value.device.type in ["cpu", "cuda", "mps", "xpu"], f"Tensor {key} on unexpected device"
def test_raw_observation_to_observation_deterministic():
"""Test that the function produces consistent results for the same input."""
robot_obs = _create_mock_robot_observation()
lerobot_features = _create_mock_lerobot_features()
policy_image_features = _create_mock_policy_image_features()
# Run twice with same input
obs1 = raw_observation_to_observation(robot_obs, lerobot_features, policy_image_features)
obs2 = raw_observation_to_observation(robot_obs, lerobot_features, policy_image_features)
# Results should be identical
assert set(obs1.keys()) == set(obs2.keys())
for key in obs1:
if isinstance(obs1[key], torch.Tensor):
torch.testing.assert_close(obs1[key], obs2[key])
else:
assert obs1[key] == obs2[key]
def test_image_processing_pipeline_preserves_content():
"""Test that the image processing pipeline preserves recognizable patterns."""
# Create an image with a specific pattern
original_img = np.zeros((100, 100, 3), dtype=np.uint8)
original_img[25:75, 25:75, :] = 255 # White square in center
robot_obs = {"shoulder": 1.0, "elbow": 1.0, "wrist": 1.0, "gripper": 1.0, "laptop": original_img}
lerobot_features = {
OBS_STATE: {
"dtype": "float32",
"shape": [4],
"names": ["shoulder", "elbow", "wrist", "gripper"],
},
f"{OBS_IMAGES}.laptop": {
"dtype": "image",
"shape": [100, 100, 3],
"names": ["height", "width", "channels"],
},
}
policy_image_features = {
f"{OBS_IMAGES}.laptop": PolicyFeature(
type=FeatureType.VISUAL,
shape=(3, 50, 50), # Downsamples from 100x100
)
}
observation = raw_observation_to_observation(robot_obs, lerobot_features, policy_image_features)
processed_img = observation[f"{OBS_IMAGES}.laptop"].squeeze(0) # Remove batch dim
# Check that the center region has higher values than corners
# Due to bilinear interpolation, exact values will change but pattern should remain
center_val = processed_img[:, 25, 25].mean() # Center of 50x50 image
corner_val = processed_img[:, 5, 5].mean() # Corner
assert center_val > corner_val, "Image processing should preserve recognizable patterns"
tests/async_inference/test_policy_server.py
-219
View File
@@ -1,219 +0,0 @@
# Copyright 2025 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 `PolicyServer` core logic.
Monkey-patch the `policy` attribute with a stub so that no real model inference is performed.
"""
from __future__ import annotations
import time
import pytest
import torch
from lerobot.configs.types import PolicyFeature
from lerobot.utils.constants import OBS_STATE
from tests.utils import skip_if_package_missing
# -----------------------------------------------------------------------------
# Test fixtures
# -----------------------------------------------------------------------------
class MockPolicy:
"""A minimal mock for an actual policy, returning zeros.
Refer to tests/policies for tests of the individual policies supported."""
class _Config:
robot_type = "dummy_robot"
@property
def image_features(self) -> dict[str, PolicyFeature]:
"""Empty image features since this test doesn't use images."""
return {}
def predict_action_chunk(self, observation: dict[str, torch.Tensor]) -> torch.Tensor:
"""Return a chunk of 20 dummy actions."""
batch_size = len(observation[OBS_STATE])
return torch.zeros(batch_size, 20, 6)
def __init__(self):
self.config = self._Config()
def to(self, *args, **kwargs):
# The server calls `policy.to(device)`. This stub ignores it.
return self
def model(self, batch: dict) -> torch.Tensor:
# Return a chunk of 20 dummy actions.
batch_size = len(batch["robot_type"])
return torch.zeros(batch_size, 20, 6)
@pytest.fixture
@skip_if_package_missing("grpcio", "grpc")
def policy_server():
"""Fresh `PolicyServer` instance with a stubbed-out policy model."""
# Import only when the test actually runs (after decorator check)
from lerobot.async_inference.configs import PolicyServerConfig
from lerobot.async_inference.policy_server import PolicyServer
test_config = PolicyServerConfig(host="localhost", port=9999)
server = PolicyServer(test_config)
# Replace the real policy with our fast, deterministic stub.
server.policy = MockPolicy()
server.actions_per_chunk = 20
server.device = "cpu"
# Add mock lerobot_features that the observation similarity functions need
server.lerobot_features = {
OBS_STATE: {
"dtype": "float32",
"shape": [6],
"names": ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"],
}
}
return server
# -----------------------------------------------------------------------------
# Helper utilities for tests
# -----------------------------------------------------------------------------
def _make_obs(state: torch.Tensor, timestep: int = 0, must_go: bool = False):
"""Create a TimedObservation with a given state vector."""
# Import only when needed
from lerobot.async_inference.helpers import TimedObservation
return TimedObservation(
observation={
"joint1": state[0].item() if len(state) > 0 else 0.0,
"joint2": state[1].item() if len(state) > 1 else 0.0,
"joint3": state[2].item() if len(state) > 2 else 0.0,
"joint4": state[3].item() if len(state) > 3 else 0.0,
"joint5": state[4].item() if len(state) > 4 else 0.0,
"joint6": state[5].item() if len(state) > 5 else 0.0,
},
timestamp=time.time(),
timestep=timestep,
must_go=must_go,
)
# -----------------------------------------------------------------------------
# Tests
# -----------------------------------------------------------------------------
def test_time_action_chunk(policy_server):
"""Verify that `_time_action_chunk` assigns correct timestamps and timesteps."""
start_ts = time.time()
start_t = 10
# A chunk of 3 action tensors.
action_tensors = [torch.randn(6) for _ in range(3)]
timed_actions = policy_server._time_action_chunk(start_ts, action_tensors, start_t)
assert len(timed_actions) == 3
# Check timesteps
assert [ta.get_timestep() for ta in timed_actions] == [10, 11, 12]
# Check timestamps
expected_timestamps = [
start_ts,
start_ts + policy_server.config.environment_dt,
start_ts + 2 * policy_server.config.environment_dt,
]
for ta, expected_ts in zip(timed_actions, expected_timestamps, strict=True):
assert abs(ta.get_timestamp() - expected_ts) < 1e-6
def test_maybe_enqueue_observation_must_go(policy_server):
"""An observation with `must_go=True` is always enqueued."""
obs = _make_obs(torch.zeros(6), must_go=True)
assert policy_server._enqueue_observation(obs) is True
assert policy_server.observation_queue.qsize() == 1
assert policy_server.observation_queue.get_nowait() is obs
def test_maybe_enqueue_observation_dissimilar(policy_server):
"""A dissimilar observation (not `must_go`) is enqueued."""
# Set a last predicted observation.
policy_server.last_processed_obs = _make_obs(torch.zeros(6))
# Create a new, dissimilar observation.
new_obs = _make_obs(torch.ones(6) * 5) # High norm difference
assert policy_server._enqueue_observation(new_obs) is True
assert policy_server.observation_queue.qsize() == 1
def test_maybe_enqueue_observation_is_skipped(policy_server):
"""A similar observation (not `must_go`) is skipped."""
# Set a last predicted observation.
policy_server.last_processed_obs = _make_obs(torch.zeros(6))
# Create a new, very similar observation.
new_obs = _make_obs(torch.zeros(6) + 1e-4)
assert policy_server._enqueue_observation(new_obs) is False
assert policy_server.observation_queue.empty() is True
def test_obs_sanity_checks(policy_server):
"""Unit-test the private `_obs_sanity_checks` helper."""
prev = _make_obs(torch.zeros(6), timestep=0)
# Case 1 – timestep already predicted
policy_server._predicted_timesteps.add(1)
obs_same_ts = _make_obs(torch.ones(6), timestep=1)
assert policy_server._obs_sanity_checks(obs_same_ts, prev) is False
# Case 2 – observation too similar
policy_server._predicted_timesteps.clear()
obs_similar = _make_obs(torch.zeros(6) + 1e-4, timestep=2)
assert policy_server._obs_sanity_checks(obs_similar, prev) is False
# Case 3 – genuinely new & dissimilar observation passes
obs_ok = _make_obs(torch.ones(6) * 5, timestep=3)
assert policy_server._obs_sanity_checks(obs_ok, prev) is True
def test_predict_action_chunk(monkeypatch, policy_server):
"""End-to-end test of `_predict_action_chunk` with a stubbed _get_action_chunk."""
# Import only when needed
from lerobot.async_inference.policy_server import PolicyServer
# Force server to act-style policy; patch method to return deterministic tensor
policy_server.policy_type = "act"
# NOTE(Steven): Smelly tests as the Server is a state machine being partially mocked. Adding these processors as a quick fix.
policy_server.preprocessor = lambda obs: obs
policy_server.postprocessor = lambda tensor: tensor
action_dim = 6
batch_size = 1
actions_per_chunk = policy_server.actions_per_chunk
def _fake_get_action_chunk(_self, _obs, _type="act"):
return torch.zeros(batch_size, actions_per_chunk, action_dim)
monkeypatch.setattr(PolicyServer, "_get_action_chunk", _fake_get_action_chunk, raising=True)
obs = _make_obs(torch.zeros(6), timestep=5)
timed_actions = policy_server._predict_action_chunk(obs)
assert len(timed_actions) == actions_per_chunk
assert [ta.get_timestep() for ta in timed_actions] == list(range(5, 5 + actions_per_chunk))
for i, ta in enumerate(timed_actions):
expected_ts = obs.get_timestamp() + i * policy_server.config.environment_dt
assert abs(ta.get_timestamp() - expected_ts) < 1e-6
tests/async_inference/test_robot_client.py
-271
View File
@@ -1,271 +0,0 @@
# Copyright 2025 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 `RobotClient` action-queue logic (pure Python, no gRPC).
We monkey-patch `lerobot.robots.utils.make_robot_from_config` so that
no real hardware is accessed. Only the queue-update mechanism is verified.
"""
from __future__ import annotations
import time
from queue import Queue
import pytest
import torch
# Skip entire module if required deps are not available
pytest.importorskip("grpc")
pytest.importorskip("serial", reason="pyserial is required (install lerobot[hardware])")
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
# -----------------------------------------------------------------------------
# Test fixtures
# -----------------------------------------------------------------------------
@pytest.fixture()
def robot_client():
"""Fresh `RobotClient` instance for each test case (no threads started).
Uses DummyRobot."""
# Import only when the test actually runs (after decorator check)
from lerobot.async_inference.configs import RobotClientConfig
from lerobot.async_inference.robot_client import RobotClient
from tests.mocks.mock_robot import MockRobotConfig
test_config = MockRobotConfig()
# gRPC channel is not actually used in tests, so using a dummy address
test_config = RobotClientConfig(
robot=test_config,
server_address="localhost:9999",
policy_type="test",
pretrained_name_or_path="test",
actions_per_chunk=20,
)
client = RobotClient(test_config)
# Initialize attributes that are normally set in start() method
client.chunks_received = 0
client.available_actions_size = []
yield client
if client.robot.is_connected:
client.stop()
# -----------------------------------------------------------------------------
# Helper utilities for tests
# -----------------------------------------------------------------------------
def _make_actions(start_ts: float, start_t: int, count: int):
"""Generate `count` consecutive TimedAction objects starting at timestep `start_t`."""
from lerobot.async_inference.helpers import TimedAction
fps = 30 # emulates most common frame-rate
actions = []
for i in range(count):
timestep = start_t + i
timestamp = start_ts + i * (1 / fps)
action_tensor = torch.full((6,), timestep, dtype=torch.float32)
actions.append(TimedAction(action=action_tensor, timestep=timestep, timestamp=timestamp))
return actions
# -----------------------------------------------------------------------------
# Tests
# -----------------------------------------------------------------------------
def test_update_action_queue_discards_stale(robot_client):
"""`_update_action_queue` must drop actions with `timestep` <= `latest_action`."""
# Pretend we already executed up to action #4
robot_client.latest_action = 4
# Incoming chunk contains timesteps 3..7 -> expect 5,6,7 kept.
incoming = _make_actions(start_ts=time.time(), start_t=3, count=5) # 3,4,5,6,7
robot_client._aggregate_action_queues(incoming)
# Extract timesteps from queue
resulting_timesteps = [a.get_timestep() for a in robot_client.action_queue.queue]
assert resulting_timesteps == [5, 6, 7]
@pytest.mark.parametrize(
"weight_old, weight_new",
[
(1.0, 0.0),
(0.0, 1.0),
(0.5, 0.5),
(0.2, 0.8),
(0.8, 0.2),
(0.1, 0.9),
(0.9, 0.1),
],
)
def test_aggregate_action_queues_combines_actions_in_overlap(
robot_client, weight_old: float, weight_new: float
):
"""`_aggregate_action_queues` must combine actions on overlapping timesteps according
to the provided aggregate_fn, here tested with multiple coefficients."""
from lerobot.async_inference.helpers import TimedAction
robot_client.chunks_received = 0
# Pretend we already executed up to action #4, and queue contains actions for timesteps 5..6
robot_client.latest_action = 4
current_actions = _make_actions(
start_ts=time.time(), start_t=5, count=2
) # actions are [torch.ones(6), torch.ones(6), ...]
current_actions = [
TimedAction(action=10 * a.get_action(), timestep=a.get_timestep(), timestamp=a.get_timestamp())
for a in current_actions
]
for a in current_actions:
robot_client.action_queue.put(a)
# Incoming chunk contains timesteps 3..7 -> expect 5,6,7 kept.
incoming = _make_actions(start_ts=time.time(), start_t=3, count=5) # 3,4,5,6,7
overlap_timesteps = [5, 6] # properly tested in test_aggregate_action_queues_discards_stale
nonoverlap_timesteps = [7]
robot_client._aggregate_action_queues(
incoming, aggregate_fn=lambda x1, x2: weight_old * x1 + weight_new * x2
)
queue_overlap_actions = []
queue_non_overlap_actions = []
for a in robot_client.action_queue.queue:
if a.get_timestep() in overlap_timesteps:
queue_overlap_actions.append(a)
elif a.get_timestep() in nonoverlap_timesteps:
queue_non_overlap_actions.append(a)
queue_overlap_actions = sorted(queue_overlap_actions, key=lambda x: x.get_timestep())
queue_non_overlap_actions = sorted(queue_non_overlap_actions, key=lambda x: x.get_timestep())
assert torch.allclose(
queue_overlap_actions[0].get_action(),
weight_old * current_actions[0].get_action() + weight_new * incoming[-3].get_action(),
)
assert torch.allclose(
queue_overlap_actions[1].get_action(),
weight_old * current_actions[1].get_action() + weight_new * incoming[-2].get_action(),
)
assert torch.allclose(queue_non_overlap_actions[0].get_action(), incoming[-1].get_action())
@pytest.mark.parametrize(
"chunk_size, queue_len, expected",
[
(20, 12, False), # 12 / 20 = 0.6 > g=0.5 threshold, not ready to send
(20, 8, True), # 8 / 20 = 0.4 <= g=0.5, ready to send
(10, 5, True),
(10, 6, False),
],
)
def test_ready_to_send_observation(robot_client, chunk_size: int, queue_len: int, expected: bool):
"""Validate `_ready_to_send_observation` ratio logic for various sizes."""
robot_client.action_chunk_size = chunk_size
# Clear any existing actions then fill with `queue_len` dummy entries ----
robot_client.action_queue = Queue()
dummy_actions = _make_actions(start_ts=time.time(), start_t=0, count=queue_len)
for act in dummy_actions:
robot_client.action_queue.put(act)
assert robot_client._ready_to_send_observation() is expected
@pytest.mark.parametrize(
"g_threshold, expected",
[
# The condition is `queue_size / chunk_size <= g`.
# Here, ratio = 6 / 10 = 0.6.
(0.0, False), # 0.6 <= 0.0 is False
(0.1, False),
(0.2, False),
(0.3, False),
(0.4, False),
(0.5, False),
(0.6, True), # 0.6 <= 0.6 is True
(0.7, True),
(0.8, True),
(0.9, True),
(1.0, True),
],
)
def test_ready_to_send_observation_with_varying_threshold(robot_client, g_threshold: float, expected: bool):
"""Validate `_ready_to_send_observation` with fixed sizes and varying `g`."""
# Fixed sizes for this test: ratio = 6 / 10 = 0.6
chunk_size = 10
queue_len = 6
robot_client.action_chunk_size = chunk_size
# This is the parameter we are testing
robot_client._chunk_size_threshold = g_threshold
# Fill queue with dummy actions
robot_client.action_queue = Queue()
dummy_actions = _make_actions(start_ts=time.time(), start_t=0, count=queue_len)
for act in dummy_actions:
robot_client.action_queue.put(act)
assert robot_client._ready_to_send_observation() is expected
# -----------------------------------------------------------------------------
# Regression test: robot type registry populated by robot_client imports
# -----------------------------------------------------------------------------
def test_robot_client_registers_builtin_robot_types():
"""Importing robot_client must populate RobotConfig's ChoiceRegistry.
This is a regression test for a bug introduced in #2425, where removing
robot module imports from robot_client.py caused RobotConfig's registry to
be empty, breaking CLI argument parsing with:
error: argument --robot.type: invalid choice: 'so101_follower' (choose from )
Robot types are registered via @RobotConfig.register_subclass() decorators
at import time, so all supported modules must be explicitly imported.
"""
import lerobot.async_inference.robot_client # noqa: F401
from lerobot.robots.config import RobotConfig
known_choices = RobotConfig.get_known_choices()
expected_robot_types = [
"so100_follower",
"so101_follower",
"koch_follower",
"omx_follower",
"bi_so_follower",
]
for robot_type in expected_robot_types:
assert robot_type in known_choices, (
f"Robot type '{robot_type}' is not registered in RobotConfig's ChoiceRegistry. "
f"Ensure the corresponding module is imported in robot_client.py. "
f"Known choices: {sorted(known_choices)}"
)
tests/policy_server/conftest.py
+275
View File
@@ -0,0 +1,275 @@
# 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.
"""Shared fixtures for the remote-inference test suite.
The mock policy is deterministic: chunk[t, j] = state[j] + 0.01 * t (so
tests can predict exact values), accepts the RTC kwargs, and records
every call for assertions. Pipelines mimic the
``PolicyProcessorPipeline`` surface the server uses (``__call__``,
``reset``, ``steps``); the mock postprocessor doubles actions so tests
can tell model-space from robot-space chunks.
"""
from __future__ import annotations
import socket
from dataclasses import dataclass, field
from threading import Event
import numpy as np
import pytest
import torch
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.policy_server.manifest import ModelSpec, PolicyServerManifest, ZenohSpec
from lerobot.policy_server.validation import PolicyClassification, ServingClass
ACTION_DIM = 6
CHUNK_SIZE = 20
STATE_DIM = 6
IMG_H, IMG_W = 48, 64
ACTION_NAMES = [f"joint_{i}.pos" for i in range(ACTION_DIM)]
TASK = "test task"
MODEL_ID = "mock/model"
# ---------------------------------------------------------------------------
# Mock policy & config
# ---------------------------------------------------------------------------
@dataclass
class MockPolicyConfig:
type: str = "mockchunk"
pretrained_path: str = MODEL_ID
chunk_size: int = CHUNK_SIZE
action_feature_names: list[str] = field(default_factory=lambda: list(ACTION_NAMES))
input_features: dict = field(
default_factory=lambda: {
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)),
"observation.images.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMG_H, IMG_W)),
}
)
rtc_config: object | None = None
class MockChunkPolicy:
"""Deterministic chunk policy with the RTC kwargs surface."""
name = "mockchunk"
def __init__(self, config: MockPolicyConfig | None = None):
self.config = config or MockPolicyConfig()
self.calls: list[dict] = []
self.reset_count = 0
self.rtc_initialized = False
# nn.Module surface the server touches
def to(self, *args, **kwargs):
return self
def eval(self):
return self
def reset(self):
self.reset_count += 1
def init_rtc_processor(self):
self.rtc_initialized = True
def predict_action_chunk(self, batch, inference_delay=None, prev_chunk_left_over=None):
state = batch["observation.state"]
if state.ndim == 1:
state = state.unsqueeze(0)
self.calls.append(
{
"state": state.detach().clone(),
"inference_delay": inference_delay,
"prev_chunk_left_over": None
if prev_chunk_left_over is None
else prev_chunk_left_over.detach().clone(),
"task": batch.get("task"),
}
)
steps = torch.arange(CHUNK_SIZE, dtype=torch.float32).unsqueeze(1) * 0.01
return (state[:, :ACTION_DIM].unsqueeze(1) + steps.unsqueeze(0)).clone()
# ---------------------------------------------------------------------------
# Mock processor pipelines
# ---------------------------------------------------------------------------
class MockPipeline:
"""Mimics the PolicyProcessorPipeline surface used by the server."""
def __init__(self, transform=None, steps=()):
self._transform = transform
self.steps = list(steps)
self.reset_count = 0
self.call_count = 0
def __call__(self, x):
self.call_count += 1
return self._transform(x) if self._transform is not None else x
def reset(self):
self.reset_count += 1
def make_mock_processors():
"""Identity preprocessor + doubling postprocessor (model vs robot space)."""
return MockPipeline(), MockPipeline(transform=lambda actions: actions * 2.0)
# ---------------------------------------------------------------------------
# Server fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def mock_policy():
return MockChunkPolicy()
@pytest.fixture
def shared_rtc_classification():
return PolicyClassification(
ServingClass.SHARED, supports_rtc=True, needs_queue_population=False, reason="mock"
)
def make_manifest(**overrides) -> PolicyServerManifest:
kwargs = {
"model": ModelSpec(repo_or_path=MODEL_ID, device="cpu"),
"zenoh": ZenohSpec(mode="peer"),
"default_task": TASK,
"max_sessions": 4,
"warmup_inferences": 0,
"trained_fps": 30.0,
"health_port": 0,
}
kwargs.update(overrides)
return PolicyServerManifest(**kwargs)
@pytest.fixture
def manifest():
return make_manifest()
def make_logic_server(
manifest: PolicyServerManifest | None = None,
policy: MockChunkPolicy | None = None,
classification: PolicyClassification | None = None,
processor_factory=None,
):
"""A PolicyServer with everything injected and no zenoh transport."""
from lerobot.policy_server.server import PolicyServer
policy = policy or MockChunkPolicy()
factory_calls = []
def default_factory():
pair = make_mock_processors()
factory_calls.append(pair)
return pair
server = PolicyServer(
manifest or make_manifest(),
policy=policy,
policy_cfg=policy.config,
processor_factory=processor_factory or default_factory,
classification=classification
or PolicyClassification(
ServingClass.SHARED, supports_rtc=True, needs_queue_population=False, reason="mock"
),
)
server.load_policy()
server.factory_calls = factory_calls
return server
# ---------------------------------------------------------------------------
# Client-side fixtures (hw features, observations)
# ---------------------------------------------------------------------------
@pytest.fixture
def hw_features():
return {
"observation.state": {
"dtype": "float32",
"shape": (STATE_DIM,),
"names": list(ACTION_NAMES),
},
"observation.images.front": {
"dtype": "video",
"shape": (IMG_H, IMG_W, 3),
"names": ["height", "width", "channels"],
},
}
def make_robot_obs(seed: float = 1.0) -> dict:
obs = {name: seed + 0.1 * i for i, name in enumerate(ACTION_NAMES)}
rng = np.random.default_rng(int(seed * 10))
obs["front"] = rng.integers(0, 255, size=(IMG_H, IMG_W, 3), dtype=np.uint8)
return obs
@pytest.fixture
def shutdown_event():
return Event()
# ---------------------------------------------------------------------------
# Loopback helpers
# ---------------------------------------------------------------------------
def free_tcp_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("127.0.0.1", 0))
return sock.getsockname()[1]
def make_loopback_manifest(port: int, **overrides) -> PolicyServerManifest:
return make_manifest(
zenoh=ZenohSpec(mode="peer", listen_endpoints=[f"tcp/127.0.0.1:{port}"]),
**overrides,
)
def make_remote_config(port: int, **overrides):
"""RemoteInferenceConfig dialing a loopback server (fast watchdogs)."""
from lerobot.rollout.inference.factory import RemoteInferenceConfig
kwargs = {
"connect_endpoint": f"tcp/127.0.0.1:{port}",
"zenoh_mode": "peer",
"service_model_id": MODEL_ID,
"service_task": TASK,
"jpeg_quality": 0, # raw images: byte-exact loopback
"buffer_time_s": 0.2,
"handshake_timeout_s": 2.0,
"request_timeout_s": 1.0,
"degraded_after_s": 0.3,
"reconnect_initial_backoff_s": 0.1,
"reconnect_max_backoff_s": 0.5,
"max_offline_s": 8.0,
}
kwargs.update(overrides)
return RemoteInferenceConfig(**kwargs)
tests/policy_server/test_codec.py
+412
View File
@@ -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"
tests/policy_server/test_golden_parity.py
+235
View File
@@ -0,0 +1,235 @@
# 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.
"""Golden parity contract test: remote request path == local RTC compute path.
The local side replicates exactly what ``RTCInferenceEngine._rtc_loop``
(rtc.py) does per iteration; the remote side runs the same observation
through the wire codec (encode -> decode), ``PolicyServer.run_inference_request``,
and the action-chunk codec — no network, no threads. With the same
deterministic policy and identical inputs, both ActionQueues must stay
byte-identical merge after merge.
"""
from __future__ import annotations
import numpy as np
import pytest
import torch
pytest.importorskip("msgpack")
from lerobot.policies.rtc import ActionQueue # noqa: E402
from lerobot.policies.rtc.configuration_rtc import RTCConfig # noqa: E402
from lerobot.policies.utils import prepare_observation_for_inference # noqa: E402
from lerobot.policy_server import codec # noqa: E402
from lerobot.policy_server.schema import MsgHeader, ObservationMsg # noqa: E402
from lerobot.policy_server.server import _normalize_prev_actions_length # noqa: E402
from lerobot.policy_server.session import Session # noqa: E402
from lerobot.utils.constants import OBS_STATE, OBS_STR # noqa: E402
from lerobot.utils.feature_utils import build_dataset_frame # noqa: E402
from tests.policy_server.conftest import ( # noqa: E402
ACTION_NAMES,
CHUNK_SIZE,
STATE_DIM,
TASK,
MockChunkPolicy,
make_logic_server,
make_mock_processors,
make_robot_obs,
)
# Must match make_manifest()'s default RTCConfig (enabled=True, horizon=10).
EXECUTION_HORIZON = 10
ROBOT_TYPE = "mock_robot"
# Fixed per-cycle inference-delay hints; cycle 2 exercises a non-zero delay.
DELAYS = [0, 2, 1]
# Actions consumed from both queues between cycles (makes prefixes non-trivial).
CONSUME_K = 4
STATE_ONLY_FEATURES = {
OBS_STATE: {
"dtype": "float32",
"shape": (STATE_DIM,),
"names": list(ACTION_NAMES),
},
}
def _make_queue() -> ActionQueue:
return ActionQueue(RTCConfig(enabled=True, execution_horizon=EXECUTION_HORIZON))
def _local_cycle(policy, pre, post, queue, features, obs, delay) -> None:
"""Replicates the loop body of RTCInferenceEngine._rtc_loop (rtc.py)."""
idx_before = queue.get_action_index()
prev_actions = queue.get_left_over()
obs_batch = build_dataset_frame(features, obs, prefix=OBS_STR)
obs_batch = prepare_observation_for_inference(obs_batch, torch.device("cpu"), TASK, ROBOT_TYPE)
obs_batch["task"] = [TASK]
preprocessed = pre(obs_batch)
if prev_actions is not None:
prev_actions = _normalize_prev_actions_length(prev_actions, target_steps=EXECUTION_HORIZON)
actions = policy.predict_action_chunk(
preprocessed, inference_delay=delay, prev_chunk_left_over=prev_actions
)
original = actions.squeeze(0).clone()
processed = post(actions).squeeze(0)
queue.merge(original, processed, delay, idx_before)
def _remote_cycle(server, session, queue, features, obs, delay, seq_id) -> None:
"""Replicates RemoteInferenceEngine._request_cycle with the wire codec in
the loop (encode -> decode on both legs) but no network or threads."""
idx_before = queue.get_action_index()
obs_frame = build_dataset_frame(features, obs, prefix=OBS_STR)
state = obs_frame.pop(OBS_STATE, None)
images = {k: v for k, v in obs_frame.items() if isinstance(v, np.ndarray) and v.ndim == 3}
prefix_model: np.ndarray | None = None
prefix_robot: np.ndarray | None = None
left_over = queue.get_left_over()
if left_over is not None and left_over.numel():
prefix_model = left_over[:EXECUTION_HORIZON].to(torch.float32).numpy()
processed_left_over = queue.get_processed_left_over()
if processed_left_over is not None and processed_left_over.numel():
prefix_robot = processed_left_over[:EXECUTION_HORIZON].to(torch.float32).numpy()
msg = ObservationMsg(
state=state,
images=images,
task=TASK,
inference_delay_steps=delay,
prefix_model=prefix_model,
prefix_robot=prefix_robot,
jpeg_quality=0, # raw image codec: byte-exact transport
)
decoded = codec.decode_observation(codec.encode_observation(msg))
# The float32 wire dtype must introduce zero drift: byte-exact roundtrip.
assert decoded.state.dtype == np.float32
assert decoded.state.tobytes() == np.ascontiguousarray(state).tobytes()
if prefix_model is not None:
assert decoded.prefix_model.tobytes() == np.ascontiguousarray(prefix_model).tobytes()
if prefix_robot is not None:
assert decoded.prefix_robot.tobytes() == np.ascontiguousarray(prefix_robot).tobytes()
for name, img in images.items():
assert np.array_equal(decoded.images[name], img)
reply = server.run_inference_request(session, MsgHeader(seq_id=seq_id), decoded)
chunk = codec.decode_action_chunk(codec.encode_action_chunk(reply))
# Reply leg is byte-exact too (float32 in, float32 on the wire).
assert chunk.chunk_model.tobytes() == np.ascontiguousarray(reply.chunk_model).tobytes()
assert chunk.chunk_robot.tobytes() == np.ascontiguousarray(reply.chunk_robot).tobytes()
queue.merge(
torch.from_numpy(np.ascontiguousarray(chunk.chunk_model)),
torch.from_numpy(np.ascontiguousarray(chunk.chunk_robot)),
delay,
idx_before,
)
def _drive_parity(features) -> tuple[MockChunkPolicy, MockChunkPolicy]:
"""Run DELAYS cycles through both paths, asserting queue parity after
each merge and consuming CONSUME_K actions from both queues between
cycles. Returns (local_policy, remote_policy) for call-level checks."""
policy_local = MockChunkPolicy()
pre_local, post_local = make_mock_processors()
queue_local = _make_queue()
policy_remote = MockChunkPolicy()
server = make_logic_server(policy=policy_remote)
pre_remote, post_remote = make_mock_processors()
session = Session(
session_id="parity",
client_uuid="parity-client",
task=TASK,
robot_type=ROBOT_TYPE,
rtc_enabled=True,
preprocessor=pre_remote,
postprocessor=post_remote,
)
queue_remote = _make_queue()
for cycle, delay in enumerate(DELAYS):
obs = make_robot_obs(seed=float(cycle + 1))
_local_cycle(policy_local, pre_local, post_local, queue_local, features, obs, delay)
_remote_cycle(server, session, queue_remote, features, obs, delay, seq_id=cycle + 1)
assert queue_local.queue is not None and queue_remote.queue is not None
assert torch.equal(queue_local.queue, queue_remote.queue), (
f"robot-space queues diverged (cycle {cycle})"
)
assert torch.equal(queue_local.original_queue, queue_remote.original_queue), (
f"model-space queues diverged (cycle {cycle})"
)
assert queue_local.queue.shape == (CHUNK_SIZE - min(delay, CHUNK_SIZE), len(ACTION_NAMES))
# Consume the same k actions on both sides so the next cycle's RTC
# prefixes are non-trivial (and identical).
for _ in range(CONSUME_K):
action_local = queue_local.get()
action_remote = queue_remote.get()
assert action_local is not None and action_remote is not None
assert torch.equal(action_local, action_remote)
return policy_local, policy_remote
def test_remote_path_matches_local_rtc_path_state_only():
"""3 cycles, state-only features: queues stay byte-identical."""
_drive_parity(STATE_ONLY_FEATURES)
def test_remote_path_matches_local_rtc_path_with_images(hw_features):
"""Images in the loop (raw codec) must not perturb state-driven outputs."""
_drive_parity(hw_features)
def test_policy_inputs_identical_across_paths(hw_features):
"""The strongest contract: both policies saw byte-identical inputs."""
policy_local, policy_remote = _drive_parity(hw_features)
assert len(policy_local.calls) == len(policy_remote.calls) == len(DELAYS)
for i, (local_call, remote_call) in enumerate(zip(policy_local.calls, policy_remote.calls, strict=True)):
assert torch.equal(local_call["state"], remote_call["state"]), f"state diverged (call {i})"
assert local_call["state"].dtype == remote_call["state"].dtype == torch.float32
assert local_call["inference_delay"] == remote_call["inference_delay"] == DELAYS[i]
if local_call["prev_chunk_left_over"] is None:
assert remote_call["prev_chunk_left_over"] is None
else:
assert torch.equal(local_call["prev_chunk_left_over"], remote_call["prev_chunk_left_over"])
assert local_call["prev_chunk_left_over"].shape == (EXECUTION_HORIZON, len(ACTION_NAMES))
# First cycle has no leftover; later cycles must carry a real prefix.
assert policy_local.calls[0]["prev_chunk_left_over"] is None
assert all(call["prev_chunk_left_over"] is not None for call in policy_local.calls[1:])
def test_float32_wire_dtype_is_byte_exact():
"""Round-tripping non-dyadic float32 values through the tensor codec
must reproduce the exact bytes (no dtype casts, no re-quantization)."""
rng = np.random.default_rng(7)
arr = (rng.standard_normal((CHUNK_SIZE, STATE_DIM)) * 0.1).astype(np.float32)
decoded = codec.decode_tensor(codec.encode_tensor(arr))
assert decoded.dtype == np.float32
assert decoded.shape == arr.shape
assert decoded.tobytes() == arr.tobytes()
assert torch.equal(torch.from_numpy(decoded), torch.from_numpy(arr))
tests/policy_server/test_loopback.py
+366
View File
@@ -0,0 +1,366 @@
# 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.
"""Real zenoh peer-to-peer loopback tests: PolicyServer ↔ RemoteInferenceEngine.
The server listens on a fresh loopback TCP port per test; the engine
dials it directly (``mode=peer``, no router). Mock policy values are
deterministic — chunk_robot[t, j] = 2 * (state[j] + 0.01 t) — so first
actions identify which client's observation produced them (the
per-session isolation regression). Chaos tests kill/restart the server
mid-episode and assert the engine degrades and recovers without ever
raising on the control thread.
"""
import time
from threading import Event
import pytest
import torch
pytest.importorskip("msgpack")
zenoh = pytest.importorskip("zenoh")
from lerobot.policy_server.schema import MsgHeader, obs_key # noqa: E402
from lerobot.policy_server.zenoh_utils import build_zenoh_config # noqa: E402
from lerobot.rollout.inference.factory import FallbackMode # noqa: E402
from lerobot.rollout.inference.remote import ClientState, RemoteInferenceEngine # noqa: E402
from tests.policy_server.conftest import ( # noqa: E402
ACTION_DIM,
ACTION_NAMES,
TASK,
free_tcp_port,
make_logic_server,
make_loopback_manifest,
make_remote_config,
make_robot_obs,
)
_FPS = 30.0
_TICK_S = 1.0 / _FPS
# Settle time after server.start() for zenoh declarations to propagate.
_DECLARATION_SETTLE_S = 0.5
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _wait_until(predicate, timeout_s: float, interval_s: float = 0.05) -> bool:
"""Poll ``predicate`` until true or the deadline passes (never a fixed sleep)."""
deadline = time.monotonic() + timeout_s
while time.monotonic() < deadline:
if predicate():
return True
time.sleep(interval_s)
return bool(predicate())
def _start_loopback_server(port: int, attempts: int = 3):
"""Start a fully-injected PolicyServer listening on tcp/127.0.0.1:<port>."""
last_error: Exception | None = None
for _ in range(attempts):
server = make_logic_server(make_loopback_manifest(port))
try:
server.start()
except Exception as e: # noqa: BLE001 — e.g. lingering socket on restart
last_error = e
server.stop()
time.sleep(0.5)
continue
time.sleep(_DECLARATION_SETTLE_S)
return server
raise last_error
def _make_engine(port: int, server, hw_features: dict, **config_overrides) -> RemoteInferenceEngine:
return RemoteInferenceEngine(
config=make_remote_config(port, **config_overrides),
policy_config=server._policy_cfg,
hw_features=hw_features,
ordered_action_keys=list(ACTION_NAMES),
task=TASK,
fps=_FPS,
robot_type="mock",
shutdown_event=Event(),
)
def _start_engine(engine: RemoteInferenceEngine, attempts: int = 4) -> None:
"""Engine start with handshake retries (declarations may still be settling)."""
last_error: Exception | None = None
for _ in range(attempts):
try:
engine.start()
return
except ConnectionError as e:
last_error = e
engine.stop()
time.sleep(0.3)
raise last_error
def _collect_actions(engine: RemoteInferenceEngine, n: int, timeout_s: float) -> list[torch.Tensor]:
"""Poll ``get_action`` at ~30 Hz until ``n`` actions arrive or the deadline passes."""
actions: list[torch.Tensor] = []
deadline = time.monotonic() + timeout_s
while len(actions) < n and time.monotonic() < deadline:
action = engine.get_action(None)
if action is not None:
actions.append(action)
time.sleep(_TICK_S)
return actions
# ---------------------------------------------------------------------------
# Happy path
# ---------------------------------------------------------------------------
@pytest.mark.timeout(60)
def test_end_to_end_chunks(hw_features):
port = free_tcp_port()
server = _start_loopback_server(port)
engine = _make_engine(port, server, hw_features)
try:
_start_engine(engine)
engine.notify_observation(make_robot_obs(2.0))
actions = _collect_actions(engine, n=20, timeout_s=15.0)
assert len(actions) >= 20
# chunk_robot[t, j] = 2 * (2.0 + 0.1 j + 0.01 t); the queue head is
# trimmed by the (small, loopback) delay → within 0.1 of the t=0 value.
first = actions[0]
assert first.shape == (ACTION_DIM,)
for j in range(ACTION_DIM):
expected = 2.0 * (2.0 + 0.1 * j)
assert abs(first[j].item() - expected) < 0.1, f"joint {j}: {first[j].item()} vs {expected}"
assert engine.state == ClientState.STREAMING
assert engine.ready
assert engine.failed is False
assert engine.stats["merges"] >= 1
finally:
engine.stop()
server.stop()
@pytest.mark.timeout(60)
def test_multi_client_no_cross_contamination(hw_features):
port = free_tcp_port()
server = _start_loopback_server(port)
engine_a = _make_engine(port, server, hw_features, client_uuid="client-a")
engine_b = _make_engine(port, server, hw_features, client_uuid="client-b")
try:
_start_engine(engine_a)
_start_engine(engine_b)
engine_a.notify_observation(make_robot_obs(2.0))
engine_b.notify_observation(make_robot_obs(7.0))
actions_a = _collect_actions(engine_a, n=1, timeout_s=10.0)
actions_b = _collect_actions(engine_b, n=1, timeout_s=10.0)
assert actions_a, "engine A produced no actions"
assert actions_b, "engine B produced no actions"
# Each engine's first action must reflect ITS OWN observation seed:
# 2*(2.0) = 4.0 for A, 2*(7.0) = 14.0 for B (gap 10.0 ≫ tolerance).
first_a = actions_a[0][0].item()
first_b = actions_b[0][0].item()
assert abs(first_a - 4.0) < 0.3, f"engine A got {first_a} (cross-contamination?)"
assert abs(first_b - 14.0) < 0.3, f"engine B got {first_b} (cross-contamination?)"
finally:
engine_a.stop()
engine_b.stop()
server.stop()
@pytest.mark.timeout(60)
def test_reset_roundtrip(hw_features):
port = free_tcp_port()
server = _start_loopback_server(port)
engine = _make_engine(port, server, hw_features)
try:
_start_engine(engine)
engine.notify_observation(make_robot_obs(2.0))
assert _collect_actions(engine, n=3, timeout_s=10.0), "no actions before reset"
merges_before = engine.stats["merges"]
engine.reset()
engine.notify_observation(make_robot_obs(5.0))
# New merges land after the reset (worker keeps cycling).
assert _wait_until(lambda: engine.stats["merges"] > merges_before, timeout_s=10.0)
# The queue refills with post-reset chunks.
assert _collect_actions(engine, n=1, timeout_s=5.0), "queue did not refill after reset"
# Server-side session advanced to the new episode (via the acked
# reset query, or via the episode bump in the next obs header).
def _episode_advanced() -> bool:
sessions = server.registry.snapshot()
return bool(sessions) and sessions[0].episode_id >= 1
assert _wait_until(_episode_advanced, timeout_s=8.0), "server session episode_id never advanced"
assert engine.failed is False
finally:
engine.stop()
server.stop()
# ---------------------------------------------------------------------------
# Chaos: server death / restart
# ---------------------------------------------------------------------------
@pytest.mark.timeout(60)
def test_server_death_is_safe(hw_features):
port = free_tcp_port()
server = _start_loopback_server(port)
engine = _make_engine(port, server, hw_features)
try:
_start_engine(engine)
engine.notify_observation(make_robot_obs(2.0))
assert _collect_actions(engine, n=5, timeout_s=10.0), "no actions before server death"
server.stop()
# Keep ticking at 30 Hz for ~2 s: get_action must never raise.
results = []
deadline = time.monotonic() + 2.0
while time.monotonic() < deadline:
results.append(engine.get_action(None))
time.sleep(_TICK_S)
# The local queue drains and HOLD fallback yields None.
assert all(result is None for result in results[-5:]), "queue never drained to HOLD fallback"
# max_offline_s=8 not reached → not failed, in a degraded-but-alive state.
assert engine.failed is False
assert engine.state in {ClientState.DEGRADED, ClientState.STALLED, ClientState.RECONNECTING}
finally:
engine.stop()
server.stop()
@pytest.mark.timeout(60)
def test_server_restart_recovery(hw_features):
port = free_tcp_port()
server = _start_loopback_server(port)
engine = _make_engine(port, server, hw_features, max_offline_s=45.0)
new_server = None
try:
_start_engine(engine)
engine.notify_observation(make_robot_obs(2.0))
assert _collect_actions(engine, n=3, timeout_s=10.0), "no actions before server death"
server.stop()
# Let the engine notice the death (liveliness drop / request timeout).
_wait_until(lambda: engine.state != ClientState.STREAMING, timeout_s=5.0)
new_server = _start_loopback_server(port)
# Re-handshake: bounded by the engine backoff and zenoh's TCP
# reconnect period — poll generously rather than sleeping.
reconnected = _wait_until(lambda: engine.stats["reconnects"] >= 1, timeout_s=25.0, interval_s=0.1)
assert reconnected, f"engine never re-handshook (state={engine.state})"
engine.notify_observation(make_robot_obs(2.0))
actions = _collect_actions(engine, n=3, timeout_s=10.0)
assert len(actions) >= 3, "no actions after server restart"
assert abs(actions[-1][0].item() - 4.0) < 0.3
assert engine.failed is False
finally:
engine.stop()
server.stop()
if new_server is not None:
new_server.stop()
# ---------------------------------------------------------------------------
# Robustness: unknown clients, fallback modes
# ---------------------------------------------------------------------------
@pytest.mark.timeout(60)
def test_unknown_client_dropped(hw_features):
port = free_tcp_port()
server = _start_loopback_server(port)
intruder = None
engine = None
try:
intruder = zenoh.open(build_zenoh_config(mode="peer", connect_endpoints=[f"tcp/127.0.0.1:{port}"]))
key = obs_key(server.prefix, "intruder-uuid")
header_bytes = MsgHeader().pack() # valid header, garbage body, no session
deadline = time.monotonic() + 8.0
while server.metrics["dropped_unknown_client_total"] < 1 and time.monotonic() < deadline:
intruder.put(key, b"\xde\xad\xbe\xef", attachment=header_bytes)
time.sleep(0.1)
assert server.metrics["dropped_unknown_client_total"] >= 1
assert len(server.registry) == 0
# The server stays healthy: a legitimate engine still works.
engine = _make_engine(port, server, hw_features)
_start_engine(engine)
engine.notify_observation(make_robot_obs(2.0))
actions = _collect_actions(engine, n=1, timeout_s=10.0)
assert actions, "legitimate engine got no actions after garbage traffic"
assert abs(actions[0][0].item() - 4.0) < 0.3
finally:
if engine is not None:
engine.stop()
if intruder is not None:
intruder.close()
server.stop()
@pytest.mark.timeout(60)
def test_fallback_zero(hw_features):
port = free_tcp_port()
server = _start_loopback_server(port)
engine = _make_engine(port, server, hw_features, fallback=FallbackMode.ZERO)
try:
_start_engine(engine)
engine.notify_observation(make_robot_obs(2.0))
# With ZERO fallback an empty queue already yields zeros, so wait
# for a *streamed* (nonzero ~4.0) action to prove chunks flowed.
streamed = False
deadline = time.monotonic() + 10.0
while time.monotonic() < deadline:
action = engine.get_action(None)
if action is not None and torch.count_nonzero(action) > 0:
streamed = True
break
time.sleep(_TICK_S)
assert streamed, "no streamed (nonzero) actions before server death"
server.stop()
# Drain the local queue; once dry, ZERO fallback must return an
# explicit zero command (never None) of the action dimension.
saw_zero = False
deadline = time.monotonic() + 6.0
while time.monotonic() < deadline:
action = engine.get_action(None)
assert action is not None, "FallbackMode.ZERO returned None"
if torch.count_nonzero(action) == 0:
assert action.shape == (len(ACTION_NAMES),)
saw_zero = True
break
time.sleep(_TICK_S)
assert saw_zero, "queue never drained to the zero fallback"
assert engine.failed is False
finally:
engine.stop()
server.stop()
tests/policy_server/test_manifest.py
+190
View File
@@ -0,0 +1,190 @@
# 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 policy-server manifest (defaults + __post_init__ validation)."""
from __future__ import annotations
import textwrap
import draccus
import pytest
from lerobot.policy_server.manifest import (
SERVING_MODE_AUTO,
ModelSpec,
PolicyServerManifest,
ZenohSpec,
)
def _manifest(**overrides) -> PolicyServerManifest:
kwargs: dict = {"model": ModelSpec(repo_or_path="mock/model")}
kwargs.update(overrides)
return PolicyServerManifest(**kwargs)
# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------
def test_defaults_parse():
# The default zenoh mode is "client", which requires a router address;
# peer mode is the minimal valid transport for a defaults-only manifest.
manifest = _manifest(zenoh=ZenohSpec(mode="peer"))
assert manifest.model.repo_or_path == "mock/model"
assert manifest.model.revision == "main"
assert manifest.model.device == "cuda"
assert manifest.serving_mode == SERVING_MODE_AUTO
assert manifest.max_sessions == 5
assert manifest.warmup_inferences == 2
assert manifest.trained_fps == 30.0
assert manifest.strict_fps is False
assert manifest.pin_task is False
assert manifest.session_idle_timeout_s == 300.0
assert manifest.health_port == 9100
assert manifest.debug.capture_dir is None
assert manifest.rtc.enabled is True
# Bare ZenohSpec defaults (validated only when embedded in a manifest).
assert ZenohSpec().mode == "client"
assert ZenohSpec().connect_endpoints == []
# ---------------------------------------------------------------------------
# __post_init__ validation
# ---------------------------------------------------------------------------
def test_missing_repo_or_path_raises():
with pytest.raises(ValueError, match="repo_or_path"):
PolicyServerManifest()
def test_bad_serving_mode_raises():
with pytest.raises(ValueError, match="serving_mode"):
_manifest(serving_mode="multiplexed")
@pytest.mark.parametrize("max_sessions", [0, -3])
def test_max_sessions_below_one_raises(max_sessions):
with pytest.raises(ValueError, match="max_sessions"):
_manifest(max_sessions=max_sessions)
def test_zenoh_client_mode_requires_connect_endpoints():
with pytest.raises(ValueError, match="connect_endpoints"):
_manifest(zenoh=ZenohSpec(mode="client", connect_endpoints=[]))
def test_zenoh_client_mode_with_endpoints_ok():
manifest = _manifest(zenoh=ZenohSpec(mode="client", connect_endpoints=["tcp/router:7447"]))
assert manifest.zenoh.connect_endpoints == ["tcp/router:7447"]
def test_zenoh_peer_mode_without_endpoints_ok():
manifest = _manifest(zenoh=ZenohSpec(mode="peer"))
assert manifest.zenoh.mode == "peer"
assert manifest.zenoh.connect_endpoints == []
def test_partial_tls_triple_raises():
with pytest.raises(ValueError, match="TLS"):
_manifest(zenoh=ZenohSpec(mode="peer", tls_root_ca_certificate="/certs/ca.pem"))
def test_full_tls_triple_ok():
manifest = _manifest(
zenoh=ZenohSpec(
mode="peer",
tls_root_ca_certificate="/certs/ca.pem",
tls_connect_certificate="/certs/cert.pem",
tls_connect_private_key="/certs/key.pem",
)
)
assert manifest.zenoh.tls_connect_private_key == "/certs/key.pem"
# ---------------------------------------------------------------------------
# Draccus round-trip (YAML manifest → dataclass)
# ---------------------------------------------------------------------------
def test_draccus_yaml_round_trip(tmp_path):
yaml_path = tmp_path / "server.yaml"
yaml_path.write_text(
textwrap.dedent(
"""\
model:
repo_or_path: mock/model
revision: v2.0
device: cpu
zenoh:
mode: peer
listen_endpoints:
- tcp/127.0.0.1:7447
default_task: pick the cube
pin_task: true
serving_mode: exclusive
max_sessions: 1
trained_fps: 25.0
strict_fps: true
health_port: 0
"""
)
)
manifest = draccus.parse(PolicyServerManifest, config_path=str(yaml_path), args=[])
assert isinstance(manifest, PolicyServerManifest)
assert manifest.model.repo_or_path == "mock/model"
assert manifest.model.revision == "v2.0"
assert manifest.model.device == "cpu"
assert manifest.zenoh.mode == "peer"
assert manifest.zenoh.listen_endpoints == ["tcp/127.0.0.1:7447"]
assert manifest.default_task == "pick the cube"
assert manifest.pin_task is True
assert manifest.serving_mode == "exclusive"
assert manifest.max_sessions == 1
assert manifest.trained_fps == 25.0
assert manifest.strict_fps is True
assert manifest.health_port == 0
# Untouched fields keep their defaults.
assert manifest.warmup_inferences == 2
assert manifest.session_idle_timeout_s == 300.0
def test_draccus_cli_override_on_top_of_yaml(tmp_path):
yaml_path = tmp_path / "server.yaml"
yaml_path.write_text(
textwrap.dedent(
"""\
model:
repo_or_path: mock/model
device: cpu
zenoh:
mode: peer
"""
)
)
manifest = draccus.parse(
PolicyServerManifest,
config_path=str(yaml_path),
args=["--max_sessions", "3", "--model.revision", "exp-1"],
)
assert manifest.max_sessions == 3
assert manifest.model.revision == "exp-1"
assert manifest.model.repo_or_path == "mock/model"
tests/policy_server/test_scheduler.py
+117
View File
@@ -0,0 +1,117 @@
# 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 RoundRobinScheduler fairness guarantees.
The scheduler only reads ``client_uuid`` from sessions, so minimal fakes
stand in for real Session objects.
"""
from __future__ import annotations
from collections import Counter
from lerobot.policy_server.scheduler import RoundRobinScheduler
class FakeSession:
"""The scheduler touches nothing but client_uuid."""
def __init__(self, client_uuid: str):
self.client_uuid = client_uuid
def picks(scheduler: RoundRobinScheduler, ready: list[FakeSession], n: int) -> list[str]:
served = []
for _ in range(n):
chosen = scheduler.select(list(ready))
assert len(chosen) == 1
served.append(chosen[0].client_uuid)
return served
def test_empty_ready_returns_empty():
scheduler = RoundRobinScheduler()
assert scheduler.select([]) == []
def test_empty_ready_after_serving_returns_empty():
scheduler = RoundRobinScheduler()
scheduler.select([FakeSession("a")])
assert scheduler.select([]) == []
def test_single_session_picked_repeatedly():
scheduler = RoundRobinScheduler()
only = FakeSession("solo")
for _ in range(5):
assert scheduler.select([only]) == [only]
def test_three_sessions_served_fairly_in_sorted_uuid_order():
scheduler = RoundRobinScheduler()
a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
# Pass ready in non-sorted order: the ring is sorted by uuid internally.
served = picks(scheduler, [c, a, b], 9)
assert served == ["a", "b", "c", "a", "b", "c", "a", "b", "c"]
assert Counter(served) == {"a": 3, "b": 3, "c": 3}
def test_session_leaving_between_calls_keeps_fairness():
scheduler = RoundRobinScheduler()
a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
assert scheduler.select([a, b, c]) == [a]
# 'a' leaves; remaining sessions alternate with no crash or starvation.
served = picks(scheduler, [b, c], 4)
assert served == ["b", "c", "b", "c"]
def test_departed_last_served_uuid_resumes_after_it():
scheduler = RoundRobinScheduler()
a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
picks(scheduler, [a, b, c], 2) # last served is 'b'
# 'b' leaves; the next pick is the first uuid greater than 'b'.
assert scheduler.select([a, c]) == [c]
assert scheduler.select([a, c]) == [a]
def test_wraparound_from_last_uuid_back_to_first():
scheduler = RoundRobinScheduler()
a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
assert scheduler.select([c]) == [c] # last served is the highest uuid
# Everyone is <= last served: wrap back to the first sorted uuid.
assert scheduler.select([a, b, c]) == [a]
def test_newly_ready_session_joins_ring_fairly():
scheduler = RoundRobinScheduler()
a, c = FakeSession("a"), FakeSession("c")
served = picks(scheduler, [a, c], 2)
assert served == ["a", "c"]
# 'b' becomes ready; wrap-around lands on 'a', then 'b' gets its turn.
b = FakeSession("b")
served = picks(scheduler, [a, b, c], 3)
assert served == ["a", "b", "c"]
def test_no_starvation_over_many_cycles():
scheduler = RoundRobinScheduler()
ready = [FakeSession(f"u{i:02d}") for i in range(5)]
served = picks(scheduler, ready, 50)
assert Counter(served) == {f"u{i:02d}": 10 for i in range(5)}
tests/policy_server/test_schema.py
+197
View File
@@ -0,0 +1,197 @@
# 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 wire schema: packed header and key-expression layout."""
import pytest
from lerobot.policy_server.schema import (
HEADER_SIZE,
MSG_TYPE_CHUNK,
MSG_TYPE_EVENT,
MSG_TYPE_OBS,
RESERVED_SEGMENTS,
SCHEMA_VERSION,
MsgHeader,
action_key,
client_alive_key,
client_alive_wildcard,
client_uuid_from_key,
obs_key,
obs_wildcard,
reset_key,
reset_wildcard,
sanitize_key_segment,
server_alive_key,
service_prefix,
session_key,
status_key,
)
# ---------------------------------------------------------------------------
# MsgHeader
# ---------------------------------------------------------------------------
def test_header_roundtrip_all_fields():
hdr = MsgHeader(
schema_version=3,
msg_type=MSG_TYPE_CHUNK,
seq_id=123456789,
episode_id=42,
client_mono_ns=987654321012345,
session_epoch=7,
)
out = MsgHeader.unpack(hdr.pack())
assert out == hdr
def test_header_defaults_roundtrip():
out = MsgHeader.unpack(MsgHeader().pack())
assert out.schema_version == SCHEMA_VERSION
assert out.msg_type == MSG_TYPE_OBS
assert out.seq_id == 0
assert out.episode_id == 0
assert out.client_mono_ns == 0
assert out.session_epoch == 0
def test_header_negative_client_mono_ns():
hdr = MsgHeader(msg_type=MSG_TYPE_EVENT, client_mono_ns=-123456789)
out = MsgHeader.unpack(hdr.pack())
assert out.client_mono_ns == -123456789
def test_header_max_u64_seq_id():
max_u64 = 2**64 - 1
hdr = MsgHeader(seq_id=max_u64)
out = MsgHeader.unpack(hdr.pack())
assert out.seq_id == max_u64
def test_header_size_constant_matches_pack():
assert len(MsgHeader().pack()) == HEADER_SIZE
def test_header_unpack_rejects_wrong_length():
packed = MsgHeader().pack()
with pytest.raises(ValueError, match="Bad header length"):
MsgHeader.unpack(packed[:-1])
with pytest.raises(ValueError, match="Bad header length"):
MsgHeader.unpack(packed + b"\x00")
with pytest.raises(ValueError, match="Bad header length"):
MsgHeader.unpack(b"")
# ---------------------------------------------------------------------------
# sanitize_key_segment
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("bad_char", ["/", "*", "$", "?", "#", " "])
def test_sanitize_folds_unsafe_chars_to_dash(bad_char):
assert sanitize_key_segment(f"a{bad_char}b") == "a-b"
def test_sanitize_folds_runs_to_single_dash():
assert sanitize_key_segment("a/*$?# b") == "a-b"
def test_sanitize_strips_whitespace_and_edge_dashes():
assert sanitize_key_segment(" hello ") == "hello"
assert sanitize_key_segment("/leading/trailing/") == "leading-trailing"
def test_sanitize_preserves_allowed_chars():
assert sanitize_key_segment("Az09_.-x") == "Az09_.-x"
@pytest.mark.parametrize("empty_input", ["", " ", "***", "/?#$*"])
def test_sanitize_raises_on_empty_after_sanitize(empty_input):
with pytest.raises(ValueError, match="empty"):
sanitize_key_segment(empty_input)
@pytest.mark.parametrize(
"reserved", sorted(["status", "session", "server", "obs", "action", "reset", "alive"])
)
def test_sanitize_raises_on_reserved_segments(reserved):
assert reserved in RESERVED_SEGMENTS
with pytest.raises(ValueError, match="reserved"):
sanitize_key_segment(reserved)
# ---------------------------------------------------------------------------
# service_prefix
# ---------------------------------------------------------------------------
def test_service_prefix_example():
prefix = service_prefix("lerobot/pi0_towels", "main", "fold the towel")
assert prefix == "@lerobot/lerobot-pi0_towels/main/fold-the-towel"
def test_service_prefix_defaults_for_empty_revision_and_task():
prefix = service_prefix("org/model", "", "")
assert prefix == "@lerobot/org-model/main/default"
# ---------------------------------------------------------------------------
# Key builders and wildcards
# ---------------------------------------------------------------------------
PREFIX = "@lerobot/org-model/main/default"
UUID = "client-uuid-1234"
def test_per_client_keys():
assert obs_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/obs"
assert action_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/action"
assert reset_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/reset"
assert client_alive_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/alive"
def test_service_level_keys():
assert status_key(PREFIX) == f"{PREFIX}/status"
assert session_key(PREFIX) == f"{PREFIX}/session"
assert server_alive_key(PREFIX) == f"{PREFIX}/server/alive"
def test_wildcards_are_single_depth():
assert obs_wildcard(PREFIX) == f"{PREFIX}/*/obs"
assert reset_wildcard(PREFIX) == f"{PREFIX}/*/reset"
assert client_alive_wildcard(PREFIX) == f"{PREFIX}/*/alive"
assert "**" not in obs_wildcard(PREFIX)
def test_key_builders_sanitize_client_uuid():
assert obs_key(PREFIX, "bad uuid") == f"{PREFIX}/bad-uuid/obs"
with pytest.raises(ValueError):
obs_key(PREFIX, "status")
# ---------------------------------------------------------------------------
# client_uuid_from_key
# ---------------------------------------------------------------------------
def test_client_uuid_from_obs_reset_alive_keys():
assert client_uuid_from_key(obs_key(PREFIX, UUID)) == UUID
assert client_uuid_from_key(reset_key(PREFIX, UUID)) == UUID
assert client_uuid_from_key(client_alive_key(PREFIX, UUID)) == UUID
def test_client_uuid_from_key_rejects_keys_without_client_chunk():
with pytest.raises(ValueError, match="no client chunk"):
client_uuid_from_key("obs")
tests/policy_server/test_server_logic.py
+405
View File
@@ -0,0 +1,405 @@
# 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.
"""Pure-logic PolicyServer tests (no zenoh transport).
Covers status snapshots, session open/reject/re-handshake, the
per-request inference path (determinism, RTC forwarding, echo fields,
supersession), episode boundaries in ``_serve_one``, warmup, and the
error/metrics accounting.
"""
import numpy as np
import pytest
pytest.importorskip("msgpack")
from lerobot.policy_server import codec # noqa: E402
from lerobot.policy_server.schema import MsgHeader, ObservationMsg, SessionOpenMsg # noqa: E402
from lerobot.policy_server.validation import PolicyClassification, ServingClass # noqa: E402
from tests.policy_server.conftest import ( # noqa: E402
ACTION_DIM,
ACTION_NAMES,
CHUNK_SIZE,
IMG_H,
IMG_W,
MODEL_ID,
STATE_DIM,
TASK,
MockChunkPolicy,
make_logic_server,
make_manifest,
)
CAMERA_KEY = "observation.images.front"
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def make_open_msg(client_uuid: str = "client-a", **overrides) -> SessionOpenMsg:
kwargs = {
"client_uuid": client_uuid,
"robot_type": "so101",
"policy_type": "mockchunk",
"fps": 30.0,
"action_names": list(ACTION_NAMES),
"camera_names": [CAMERA_KEY],
"state_dim": STATE_DIM,
"rtc_enabled": True,
"task": TASK,
}
kwargs.update(overrides)
return SessionOpenMsg(**kwargs)
def make_obs(**overrides) -> ObservationMsg:
kwargs = {
"state": np.arange(STATE_DIM, dtype=np.float32),
"images": {CAMERA_KEY: np.zeros((IMG_H, IMG_W, 3), dtype=np.uint8)},
"task": TASK,
"jpeg_quality": 0,
}
kwargs.update(overrides)
return ObservationMsg(**kwargs)
def open_session(server, client_uuid: str = "client-a", **overrides):
ack = server._handle_session_open(make_open_msg(client_uuid=client_uuid, **overrides))
assert ack.accepted, ack.error
return server.registry.get(client_uuid), ack
def deposit(session, obs: ObservationMsg, header: MsgHeader | None = None) -> None:
session.deposit(header or MsgHeader(episode_id=session.episode_id), codec.encode_observation(obs))
def make_exclusive_server(policy=None):
return make_logic_server(
manifest=make_manifest(serving_mode="exclusive"),
policy=policy,
classification=PolicyClassification(
ServingClass.EXCLUSIVE, supports_rtc=False, needs_queue_population=False, reason="x"
),
)
def expected_chunk(state: np.ndarray) -> np.ndarray:
steps = np.arange(CHUNK_SIZE, dtype=np.float32)[:, None] * np.float32(0.01)
return state[None, :ACTION_DIM] + steps
# ---------------------------------------------------------------------------
# status_snapshot
# ---------------------------------------------------------------------------
def test_status_snapshot_capabilities():
server = make_logic_server()
snap = server.status_snapshot()
assert snap.model_repo == MODEL_ID
assert snap.policy_type == "mockchunk"
assert snap.action_names == ACTION_NAMES
assert snap.state_dim == STATE_DIM
assert snap.chunk_size == CHUNK_SIZE
assert snap.expected_cameras == [CAMERA_KEY]
assert snap.supports_rtc is True
assert snap.warmed_up is True
assert snap.serving_mode == "shared"
assert snap.active_sessions == 0
assert snap.max_sessions == 4 # make_manifest default
# ---------------------------------------------------------------------------
# _handle_session_open
# ---------------------------------------------------------------------------
def test_session_open_happy_path():
server = make_logic_server()
session, ack = open_session(server, "client-a")
assert ack.accepted is True
assert ack.error == ""
assert ack.session_id == session.session_id
assert ack.session_id != ""
assert ack.model_repo == MODEL_ID
assert ack.policy_type == "mockchunk"
assert ack.action_names == ACTION_NAMES
assert ack.expected_cameras == [CAMERA_KEY]
assert ack.state_dim == STATE_DIM
assert ack.chunk_size == CHUNK_SIZE
assert ack.supports_rtc is True
assert ack.serving_mode == "shared"
assert ack.warmed_up is True
assert len(server.registry) == 1
assert session.rtc_enabled is True
assert session.task == TASK
assert server.metrics["sessions_opened_total"] == 1
def test_session_open_fresh_processor_pair_per_session():
server = make_logic_server()
assert len(server.factory_calls) == 0 # warmup_inferences=0: no warmup pair
session_a, _ = open_session(server, "client-a")
assert len(server.factory_calls) == 1
session_b, _ = open_session(server, "client-b")
assert len(server.factory_calls) == 2
assert session_a.preprocessor is not session_b.preprocessor
assert session_a.postprocessor is not session_b.postprocessor
def test_session_open_rejects_action_order_mismatch():
server = make_logic_server()
ack = server._handle_session_open(make_open_msg(action_names=list(reversed(ACTION_NAMES))))
assert ack.accepted is False
assert "action" in ack.error
assert len(server.registry) == 0
def test_session_open_rejects_at_capacity():
server = make_logic_server()
for i in range(4): # make_manifest max_sessions=4
open_session(server, f"client-{i}")
ack = server._handle_session_open(make_open_msg(client_uuid="client-overflow"))
assert ack.accepted is False
assert "sessions" in ack.error
assert len(server.registry) == 4
def test_rehandshake_replaces_session_without_counting_against_capacity():
server = make_logic_server()
first_ack = None
for i in range(4):
_, ack = open_session(server, f"client-{i}")
if i == 0:
first_ack = ack
# Server is full; the same client re-handshakes and must be accepted.
session, ack = open_session(server, "client-0")
assert ack.accepted is True
assert len(server.registry) == 4
assert ack.session_id != first_ack.session_id
assert server.registry.get("client-0").session_id == session.session_id
def test_session_open_rtc_downgrade():
server = make_logic_server(
classification=PolicyClassification(
ServingClass.SHARED, supports_rtc=False, needs_queue_population=False, reason="x"
)
)
session, ack = open_session(server, "client-a", rtc_enabled=True)
assert ack.accepted is True
assert ack.supports_rtc is False
assert session.rtc_enabled is False
assert any("RTC" in w for w in ack.warnings)
# ---------------------------------------------------------------------------
# run_inference_request
# ---------------------------------------------------------------------------
def test_inference_deterministic_chunks():
server = make_logic_server()
session, _ = open_session(server)
state = np.arange(STATE_DIM, dtype=np.float32)
reply = server.run_inference_request(session, MsgHeader(), make_obs(state=state))
assert reply.chunk_model.shape == (CHUNK_SIZE, ACTION_DIM)
assert reply.chunk_robot.shape == (CHUNK_SIZE, ACTION_DIM)
np.testing.assert_allclose(reply.chunk_model[0], state, rtol=0, atol=0)
np.testing.assert_allclose(reply.chunk_model, expected_chunk(state), rtol=1e-6)
np.testing.assert_allclose(reply.chunk_robot, 2.0 * reply.chunk_model, rtol=0, atol=0)
def test_inference_delay_forwarded_to_policy():
server = make_logic_server()
session, _ = open_session(server)
policy = server._policy
server.run_inference_request(session, MsgHeader(), make_obs(inference_delay_steps=3))
assert policy.calls[-1]["inference_delay"] == 3
assert policy.calls[-1]["prev_chunk_left_over"] is None
def test_prefix_model_forwarded_padded_to_execution_horizon():
server = make_logic_server()
session, _ = open_session(server)
policy = server._policy
prefix = (np.arange(3 * ACTION_DIM, dtype=np.float32).reshape(3, ACTION_DIM)) + 1.0
server.run_inference_request(session, MsgHeader(), make_obs(prefix_model=prefix))
received = policy.calls[-1]["prev_chunk_left_over"]
assert received is not None
horizon = server._manifest.rtc.execution_horizon # 10 by default
assert received.shape == (horizon, ACTION_DIM)
np.testing.assert_allclose(received[:3].numpy(), prefix, rtol=0, atol=0)
np.testing.assert_allclose(received[3:].numpy(), np.zeros((horizon - 3, ACTION_DIM)), rtol=0, atol=0)
def test_prefix_model_truncated_to_execution_horizon():
server = make_logic_server()
session, _ = open_session(server)
policy = server._policy
horizon = server._manifest.rtc.execution_horizon
prefix = np.ones((horizon + 5, ACTION_DIM), dtype=np.float32)
server.run_inference_request(session, MsgHeader(), make_obs(prefix_model=prefix))
assert policy.calls[-1]["prev_chunk_left_over"].shape == (horizon, ACTION_DIM)
def test_reply_echo_fields():
server = make_logic_server()
session, _ = open_session(server)
header = MsgHeader(seq_id=7, episode_id=2, client_mono_ns=123_456_789)
reply = server.run_inference_request(session, header, make_obs())
assert reply.seq_id_echo == 7
assert reply.episode_id_echo == 2
assert reply.client_mono_ns_echo == 123_456_789
def test_superseded_seqs_reported_then_reset():
server = make_logic_server()
session, _ = open_session(server)
deposit(session, make_obs(), MsgHeader(seq_id=1))
deposit(session, make_obs(), MsgHeader(seq_id=2)) # supersedes seq 1
item = session.take()
assert item.header.seq_id == 2 # latest-only mailbox
reply = server.run_inference_request(session, item.header, codec.decode_observation(item.payload))
assert reply.superseded_seqs == 1
deposit(session, make_obs(), MsgHeader(seq_id=3))
item = session.take()
reply = server.run_inference_request(session, item.header, codec.decode_observation(item.payload))
assert reply.superseded_seqs == 0
# ---------------------------------------------------------------------------
# _serve_one
# ---------------------------------------------------------------------------
def test_serve_one_episode_boundary_resets_session_pipelines():
server = make_logic_server()
session, _ = open_session(server)
# Fresh sessions start at the -1 sentinel so their first request
# always lands on the episode-boundary branch (mid-episode reconnects
# can never inherit stale state).
assert session.episode_id == -1
deposit(session, make_obs(episode_start=True), MsgHeader(episode_id=1))
server._serve_one(session)
assert session.preprocessor.reset_count == 1
assert session.postprocessor.reset_count == 1
assert session.episode_id == 1
# Shared mode never resets the policy itself.
assert server._policy.reset_count == 0
def test_serve_one_no_boundary_no_reset():
server = make_logic_server()
session, _ = open_session(server)
# First request always resets (the -1 sentinel) and syncs the episode.
deposit(session, make_obs(), MsgHeader(episode_id=0))
server._serve_one(session)
assert session.preprocessor.reset_count == 1
assert session.episode_id == 0
# Same-episode follow-up: no further reset.
deposit(session, make_obs(), MsgHeader(episode_id=0, seq_id=2))
server._serve_one(session)
assert session.preprocessor.reset_count == 1
assert session.postprocessor.reset_count == 1
def test_serve_one_exclusive_mode_resets_policy_on_boundary():
policy = MockChunkPolicy()
server = make_exclusive_server(policy=policy)
assert server.status_snapshot().serving_mode == "exclusive"
assert server.status_snapshot().max_sessions == 1 # exclusive forces 1
session, _ = open_session(server, rtc_enabled=False)
# Exclusive session open already resets the policy to fresh state.
base_resets = policy.reset_count
assert base_resets >= 1
deposit(session, make_obs(episode_start=True), MsgHeader(episode_id=1))
server._serve_one(session)
assert policy.reset_count == base_resets + 1
assert session.episode_id == 1
def test_serve_one_inference_error_counted_not_propagated(monkeypatch):
server = make_logic_server()
session, _ = open_session(server)
def boom(*args, **kwargs):
raise RuntimeError("boom")
monkeypatch.setattr(server._policy, "predict_action_chunk", boom)
deposit(session, make_obs())
server._serve_one(session) # must not raise
assert server.metrics["errors_total"] == 1
assert server.metrics["requests_total"] == 0
assert session.stats.errors == 1
def test_serve_one_increments_requests_total():
server = make_logic_server()
session, _ = open_session(server)
for seq in (1, 2):
deposit(session, make_obs(), MsgHeader(seq_id=seq))
server._serve_one(session)
assert server.metrics["requests_total"] == 2
assert server.metrics["errors_total"] == 0
assert session.stats.requests == 2
# ---------------------------------------------------------------------------
# Episode reset semantics (session-level, as used by _on_reset_query)
# ---------------------------------------------------------------------------
def test_session_reset_episode_clears_state():
server = make_logic_server()
session, _ = open_session(server)
deposit(session, make_obs())
assert session.has_pending()
session.reset_episode(5)
assert not session.has_pending() # mailbox cleared
assert session.episode_id == 5
assert session.preprocessor.reset_count == 1
assert session.postprocessor.reset_count == 1
session.reset_episode() # no explicit id: increments
assert session.episode_id == 6
# ---------------------------------------------------------------------------
# Warmup
# ---------------------------------------------------------------------------
def test_warmup_runs_before_any_session():
policy = MockChunkPolicy()
server = make_logic_server(manifest=make_manifest(warmup_inferences=2), policy=policy)
assert len(policy.calls) == 2
assert len(server.registry) == 0 # warmup session is not registered
for call in policy.calls:
assert tuple(call["state"].shape) == (1, STATE_DIM)
assert float(call["state"].abs().sum()) == 0.0 # synthetic zeros
assert server.status_snapshot().warmed_up is True
def test_synthetic_observation_matches_input_features():
server = make_logic_server()
obs = server._synthetic_observation()
assert obs.state.shape == (STATE_DIM,)
assert obs.state.dtype == np.float32
assert set(obs.images) == {CAMERA_KEY}
assert obs.images[CAMERA_KEY].shape == (IMG_H, IMG_W, 3)
assert obs.images[CAMERA_KEY].dtype == np.uint8
assert obs.jpeg_quality == 0
tests/policy_server/test_session.py
+298
View File
@@ -0,0 +1,298 @@
# 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 Session (latest-only mailbox, episode reset, close,
processor-step introspection) and SessionRegistry (thread-safe map)."""
from __future__ import annotations
import threading
from lerobot.policy_server.schema import MsgHeader
from lerobot.policy_server.session import Session, SessionRegistry
from lerobot.processor import NormalizerProcessorStep, RelativeActionsProcessorStep
from tests.policy_server.conftest import TASK, MockPipeline, make_mock_processors
def make_session(
client_uuid: str = "client-a",
preprocessor: MockPipeline | None = None,
postprocessor: MockPipeline | None = None,
publisher=None,
) -> Session:
default_pre, default_post = make_mock_processors()
return Session(
session_id=f"sess-{client_uuid}",
client_uuid=client_uuid,
task=TASK,
robot_type="mock_robot",
rtc_enabled=False,
preprocessor=preprocessor if preprocessor is not None else default_pre,
postprocessor=postprocessor if postprocessor is not None else default_post,
action_publisher=publisher,
)
# ---------------------------------------------------------------------------
# Mailbox: latest-only deposit / take
# ---------------------------------------------------------------------------
def test_deposit_then_take_returns_item():
session = make_session()
header = MsgHeader(seq_id=7)
session.deposit(header, b"payload-7")
item = session.take()
assert item is not None
assert item.header.seq_id == 7
assert item.payload == b"payload-7"
assert item.recv_mono > 0
def test_second_deposit_supersedes_and_take_returns_newer():
session = make_session()
session.deposit(MsgHeader(seq_id=1), b"old")
session.deposit(MsgHeader(seq_id=2), b"new")
assert session.stats.superseded == 1
assert session.stats.superseded_since_reply == 1
item = session.take()
assert item is not None
assert item.header.seq_id == 2
assert item.payload == b"new"
def test_deposit_after_take_is_not_superseded():
session = make_session()
session.deposit(MsgHeader(seq_id=1), b"one")
session.take()
session.deposit(MsgHeader(seq_id=2), b"two")
assert session.stats.superseded == 0
assert session.stats.superseded_since_reply == 0
def test_take_clears_mailbox_second_take_is_none():
session = make_session()
session.deposit(MsgHeader(seq_id=1), b"one")
assert session.take() is not None
assert session.take() is None
def test_has_pending_transitions():
session = make_session()
assert not session.has_pending()
session.deposit(MsgHeader(seq_id=1), b"one")
assert session.has_pending()
session.take()
assert not session.has_pending()
def test_deposit_marks_alive_and_clears_token_drop():
session = make_session()
session.alive = False
session.token_dropped_mono = 123.4
before = session.last_seen_mono
session.deposit(MsgHeader(seq_id=1), b"one")
assert session.alive is True
assert session.token_dropped_mono is None
assert session.last_seen_mono >= before
# ---------------------------------------------------------------------------
# Episode boundary
# ---------------------------------------------------------------------------
def test_reset_episode_resets_pipelines_clears_mailbox_and_increments():
pre, post = make_mock_processors()
session = make_session(preprocessor=pre, postprocessor=post)
session.deposit(MsgHeader(seq_id=1), b"stale")
assert session.episode_id == 0
session.reset_episode()
assert not session.has_pending()
assert pre.reset_count == 1
assert post.reset_count == 1
assert session.episode_id == 1
session.reset_episode()
assert session.episode_id == 2
assert pre.reset_count == 2
assert post.reset_count == 2
def test_reset_episode_with_explicit_id():
session = make_session()
session.reset_episode(episode_id=7)
assert session.episode_id == 7
# ---------------------------------------------------------------------------
# close()
# ---------------------------------------------------------------------------
class FakePublisher:
def __init__(self, raise_on_undeclare: bool = False):
self.undeclare_calls = 0
self.raise_on_undeclare = raise_on_undeclare
def undeclare(self):
self.undeclare_calls += 1
if self.raise_on_undeclare:
raise RuntimeError("transport already closed")
def test_close_clears_mailbox_and_undeclares_publisher_exactly_once():
publisher = FakePublisher()
session = make_session(publisher=publisher)
session.deposit(MsgHeader(seq_id=1), b"stale")
session.close()
assert not session.has_pending()
assert publisher.undeclare_calls == 1
assert session.action_publisher is None
# Idempotent: a second close must not undeclare again.
session.close()
assert publisher.undeclare_calls == 1
def test_close_tolerates_undeclare_raising():
publisher = FakePublisher(raise_on_undeclare=True)
session = make_session(publisher=publisher)
session.deposit(MsgHeader(seq_id=1), b"stale")
session.close() # must not raise
assert publisher.undeclare_calls == 1
assert not session.has_pending()
assert session.action_publisher is None
def test_close_without_publisher_is_noop():
session = make_session(publisher=None)
session.close() # must not raise
assert session.action_publisher is None
# ---------------------------------------------------------------------------
# Processor-step introspection
# ---------------------------------------------------------------------------
def test_relative_and_normalizer_steps_detected():
relative = RelativeActionsProcessorStep(enabled=True)
normalizer = NormalizerProcessorStep(features={}, norm_map={})
pre = MockPipeline(steps=[relative, normalizer])
session = make_session(preprocessor=pre)
assert session.relative_step is relative
assert session.normalizer_step is normalizer
def test_disabled_relative_step_is_not_detected():
relative = RelativeActionsProcessorStep(enabled=False)
pre = MockPipeline(steps=[relative])
session = make_session(preprocessor=pre)
assert session.relative_step is None
assert session.normalizer_step is None
def test_empty_pipeline_yields_no_introspected_steps():
session = make_session()
assert session.relative_step is None
assert session.normalizer_step is None
# ---------------------------------------------------------------------------
# SessionRegistry
# ---------------------------------------------------------------------------
def test_registry_add_get_remove_len_snapshot():
registry = SessionRegistry()
assert len(registry) == 0
assert registry.get("missing") is None
assert registry.snapshot() == []
session_a = make_session("uuid-a")
session_b = make_session("uuid-b")
assert registry.add(session_a) is None
assert registry.add(session_b) is None
assert len(registry) == 2
assert registry.get("uuid-a") is session_a
assert registry.get("uuid-b") is session_b
assert set(registry.snapshot()) == {session_a, session_b}
removed = registry.remove("uuid-a")
assert removed is session_a
assert len(registry) == 1
assert registry.get("uuid-a") is None
def test_registry_remove_missing_returns_none():
registry = SessionRegistry()
assert registry.remove("never-added") is None
def test_registry_add_returns_displaced_same_uuid_session():
registry = SessionRegistry()
first = make_session("uuid-x")
second = make_session("uuid-x")
assert registry.add(first) is None
displaced = registry.add(second)
assert displaced is first
assert registry.get("uuid-x") is second
assert len(registry) == 1
def test_registry_thread_safety_smoke():
registry = SessionRegistry()
errors: list[Exception] = []
def worker(prefix: str) -> None:
try:
for i in range(200):
session = make_session(f"{prefix}-{i}")
registry.add(session)
assert registry.get(session.client_uuid) is session
assert registry.remove(session.client_uuid) is session
except Exception as exc: # noqa: BLE001 — surfaced to the main thread
errors.append(exc)
threads = [threading.Thread(target=worker, args=(p,)) for p in ("alpha", "beta")]
for t in threads:
t.start()
for t in threads:
t.join(timeout=10)
assert not t.is_alive()
assert errors == []
assert len(registry) == 0
assert registry.snapshot() == []
tests/policy_server/test_validation.py
+323
View File
@@ -0,0 +1,323 @@
# 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 serving-mode classification and session-open validation.
Uses tiny fake policy classes (deliberately NOT subclassing
``PreTrainedPolicy``): classification keys off the ``name`` attribute and
the presence of a ``predict_action_chunk`` override, never off the class
hierarchy.
"""
from __future__ import annotations
from types import SimpleNamespace
import pytest
from lerobot.policy_server.schema import (
MIN_SUPPORTED_SCHEMA_VERSION,
SCHEMA_VERSION,
SessionOpenMsg,
StatusMsg,
)
from lerobot.policy_server.validation import (
PolicyClassification,
ServingClass,
classify_policy,
resolve_serving_mode,
validate_session_open,
)
from tests.policy_server.conftest import ACTION_NAMES, STATE_DIM, TASK, make_manifest
# ---------------------------------------------------------------------------
# Fake policy classes (classification only needs `name`, an optional
# `predict_action_chunk` method, and `.config` for smolvla)
# ---------------------------------------------------------------------------
def _fake_policy(name: str, *, chunk_api: bool = True, n_obs_steps: int | None = None):
"""Build a minimal fake policy instance with a class-level chunk method."""
namespace = {"name": name}
if chunk_api:
namespace["predict_action_chunk"] = lambda self, batch, **kwargs: None
cls = type(f"Fake_{name}_Policy", (), namespace)
policy = cls()
if n_obs_steps is not None:
policy.config = SimpleNamespace(n_obs_steps=n_obs_steps)
return policy
# ---------------------------------------------------------------------------
# classify_policy
# ---------------------------------------------------------------------------
def test_classify_act_is_shared_without_rtc():
classification = classify_policy(_fake_policy("act"))
assert classification.serving_class is ServingClass.SHARED
assert classification.supports_rtc is False
assert classification.needs_queue_population is False
@pytest.mark.parametrize("name", ["pi0", "pi05"])
def test_classify_pi_families_are_shared_with_rtc(name):
classification = classify_policy(_fake_policy(name))
assert classification.serving_class is ServingClass.SHARED
assert classification.supports_rtc is True
assert classification.needs_queue_population is False
def test_classify_smolvla_single_obs_step_is_shared():
classification = classify_policy(_fake_policy("smolvla", n_obs_steps=1))
assert classification.serving_class is ServingClass.SHARED
assert classification.supports_rtc is True
def test_classify_smolvla_with_history_is_exclusive():
classification = classify_policy(_fake_policy("smolvla", n_obs_steps=2))
assert classification.serving_class is ServingClass.EXCLUSIVE
assert classification.supports_rtc is True
assert classification.needs_queue_population is False
def test_classify_diffusion_is_exclusive_with_queue_population():
classification = classify_policy(_fake_policy("diffusion"))
assert classification.serving_class is ServingClass.EXCLUSIVE
assert classification.supports_rtc is False
assert classification.needs_queue_population is True
def test_classify_without_chunk_api_is_refused():
classification = classify_policy(_fake_policy("act", chunk_api=False))
assert classification.serving_class is ServingClass.REFUSED
assert classification.supports_rtc is False
assert "predict_action_chunk" in classification.reason
def test_classify_unknown_name_with_chunk_api_is_exclusive():
classification = classify_policy(_fake_policy("totally_new_policy"))
assert classification.serving_class is ServingClass.EXCLUSIVE
assert classification.supports_rtc is False
assert classification.needs_queue_population is False
assert "verified" in classification.reason
# ---------------------------------------------------------------------------
# resolve_serving_mode
# ---------------------------------------------------------------------------
def _classification(serving_class: ServingClass, reason: str = "test") -> PolicyClassification:
return PolicyClassification(
serving_class, supports_rtc=False, needs_queue_population=False, reason=reason
)
def test_resolve_auto_maps_shared_to_shared():
mode, max_sessions = resolve_serving_mode(
_classification(ServingClass.SHARED), make_manifest(serving_mode="auto", max_sessions=4)
)
assert mode == "shared"
assert max_sessions == 4
def test_resolve_auto_maps_exclusive_to_exclusive():
mode, max_sessions = resolve_serving_mode(
_classification(ServingClass.EXCLUSIVE), make_manifest(serving_mode="auto", max_sessions=4)
)
assert mode == "exclusive"
assert max_sessions == 1
def test_resolve_forced_shared_rejected_for_non_verified_policy():
with pytest.raises(ValueError, match="unsafe"):
resolve_serving_mode(_classification(ServingClass.EXCLUSIVE), make_manifest(serving_mode="shared"))
def test_resolve_forced_shared_allowed_for_verified_policy():
mode, max_sessions = resolve_serving_mode(
_classification(ServingClass.SHARED), make_manifest(serving_mode="shared", max_sessions=4)
)
assert mode == "shared"
assert max_sessions == 4
def test_resolve_forced_exclusive_allowed_for_shared_policy():
mode, _ = resolve_serving_mode(
_classification(ServingClass.SHARED), make_manifest(serving_mode="exclusive")
)
assert mode == "exclusive"
def test_resolve_exclusive_forces_single_session():
mode, max_sessions = resolve_serving_mode(
_classification(ServingClass.SHARED), make_manifest(serving_mode="exclusive", max_sessions=4)
)
assert mode == "exclusive"
assert max_sessions == 1
def test_resolve_refused_raises_with_reason():
with pytest.raises(ValueError, match="no chunk API here"):
resolve_serving_mode(
_classification(ServingClass.REFUSED, reason="no chunk API here"), make_manifest()
)
# ---------------------------------------------------------------------------
# validate_session_open
# ---------------------------------------------------------------------------
EXPECTED_CAMERAS = ["observation.images.front"]
@pytest.fixture
def capabilities() -> StatusMsg:
return StatusMsg(
model_repo="mock/model",
policy_type="mockchunk",
action_names=list(ACTION_NAMES),
expected_cameras=list(EXPECTED_CAMERAS),
state_dim=STATE_DIM,
chunk_size=20,
trained_fps=30.0,
supports_rtc=True,
min_schema_version=MIN_SUPPORTED_SCHEMA_VERSION,
max_schema_version=SCHEMA_VERSION,
max_sessions=4,
)
def _open_msg(**overrides) -> SessionOpenMsg:
kwargs: dict = {
"client_uuid": "client-1",
"policy_type": "mockchunk",
"fps": 30.0,
"action_names": list(ACTION_NAMES),
"camera_names": list(EXPECTED_CAMERAS),
"state_dim": STATE_DIM,
"schema_version": SCHEMA_VERSION,
"task": TASK,
}
kwargs.update(overrides)
return SessionOpenMsg(**kwargs)
def test_validate_happy_path(capabilities):
result = validate_session_open(_open_msg(), capabilities, make_manifest(), active_sessions=0)
assert result.ok
assert result.error == ""
assert result.warnings == []
assert result.rtc_downgraded is False
def test_validate_action_name_order_is_a_hard_reject(capabilities):
# Same set of names, different order: chunk columns would map to the
# wrong motors, so this must be a hard reject.
result = validate_session_open(
_open_msg(action_names=list(reversed(ACTION_NAMES))),
capabilities,
make_manifest(),
active_sessions=0,
)
assert not result.ok
assert "action" in result.error
assert "mismatch" in result.error
def test_validate_missing_camera_rejected(capabilities):
result = validate_session_open(
_open_msg(camera_names=[]), capabilities, make_manifest(), active_sessions=0
)
assert not result.ok
assert "observation.images.front" in result.error
def test_validate_wrong_state_dim_rejected(capabilities):
result = validate_session_open(
_open_msg(state_dim=STATE_DIM + 1), capabilities, make_manifest(), active_sessions=0
)
assert not result.ok
assert "state dim" in result.error
def test_validate_schema_version_out_of_range_rejected(capabilities):
result = validate_session_open(
_open_msg(schema_version=SCHEMA_VERSION + 99),
capabilities,
make_manifest(),
active_sessions=0,
)
assert not result.ok
assert "schema_version" in result.error
def test_validate_at_capacity_rejected_with_load(capabilities):
result = validate_session_open(
_open_msg(), capabilities, make_manifest(), active_sessions=capabilities.max_sessions
)
assert not result.ok
assert "full" in result.error
assert f"{capabilities.max_sessions}/{capabilities.max_sessions}" in result.error
def test_validate_pinned_task_rejects_other_task(capabilities):
result = validate_session_open(
_open_msg(task="another task"),
capabilities,
make_manifest(pin_task=True),
active_sessions=0,
)
assert not result.ok
assert "pinned" in result.error
def test_validate_fps_mismatch_strict_rejects(capabilities):
result = validate_session_open(
_open_msg(fps=15.0), capabilities, make_manifest(strict_fps=True), active_sessions=0
)
assert not result.ok
assert "fps" in result.error
def test_validate_fps_mismatch_lenient_warns_only(capabilities):
result = validate_session_open(
_open_msg(fps=15.0), capabilities, make_manifest(strict_fps=False), active_sessions=0
)
assert result.ok
assert len(result.warnings) == 1
assert "fps" in result.warnings[0]
def test_validate_rtc_downgraded_when_unsupported(capabilities):
capabilities.supports_rtc = False
result = validate_session_open(
_open_msg(rtc_enabled=True), capabilities, make_manifest(), active_sessions=0
)
assert result.ok
assert result.rtc_downgraded is True
assert any("RTC" in warning for warning in result.warnings)
def test_validate_empty_capability_action_names_skips_action_check(capabilities):
capabilities.action_names = []
result = validate_session_open(
_open_msg(action_names=["whatever.pos"]),
capabilities,
make_manifest(),
active_sessions=0,
)
assert result.ok
tests/test_rollout.py
+67
View File
@@ -348,3 +348,70 @@ def test_rollout_context_fields():
field_names = {f.name for f in dataclasses.fields(RolloutContext)}
assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
# ---------------------------------------------------------------------------
# Remote inference config & factory dispatch
# ---------------------------------------------------------------------------
def test_create_inference_engine_remote_requires_policy_config():
from lerobot.rollout.inference.factory import RemoteInferenceConfig, create_inference_engine
with pytest.raises(ValueError, match="policy_config"):
create_inference_engine(
RemoteInferenceConfig(),
policy=None,
preprocessor=None,
postprocessor=None,
robot_wrapper=MagicMock(robot_type="mock"),
hw_features={},
dataset_features={},
ordered_action_keys=["k"],
task="t",
fps=30.0,
device=None,
policy_config=None,
)
def test_remote_config_draccus_registration():
from lerobot.rollout.inference.factory import InferenceEngineConfig, RemoteInferenceConfig
assert RemoteInferenceConfig().type == "remote"
assert InferenceEngineConfig.get_choice_class("remote") is RemoteInferenceConfig
assert "remote" in dict(InferenceEngineConfig.get_known_choices())
def test_fallback_mode_values():
from lerobot.rollout.inference.factory import FallbackMode
assert FallbackMode.HOLD.value == "hold"
assert FallbackMode.REPEAT_LAST.value == "repeat_last"
assert FallbackMode.ZERO.value == "zero"
assert {mode.value for mode in FallbackMode} == {"hold", "repeat_last", "zero"}
def test_local_backends_require_loaded_policy():
from lerobot.rollout.inference.factory import (
RTCInferenceConfig,
SyncInferenceConfig,
create_inference_engine,
)
common = {
"policy": None,
"preprocessor": None,
"postprocessor": None,
"robot_wrapper": MagicMock(robot_type="mock"),
"hw_features": {},
"dataset_features": {},
"ordered_action_keys": ["k"],
"task": "t",
"fps": 30.0,
"device": "cpu",
}
with pytest.raises(ValueError, match="requires a loaded policy"):
create_inference_engine(SyncInferenceConfig(), **common)
with pytest.raises(ValueError, match="requires a loaded policy"):
create_inference_engine(RTCInferenceConfig(), **common)