feat(eval): implement docker runtime with HTTP policy inference server

Add docker_runtime.py (host-side) and lerobot_eval_worker.py (container-side)
for --eval.runtime=docker. Policy loads once on the host GPU; Docker containers
run env-only workers that call back via HTTP for action chunks, maximising GPU
utilisation across parallel benchmark tasks.

- _InferenceServer: HTTP server wrapping predict_action_chunk with a single lock
- run_eval_in_docker: spawns instance_count containers, collects + merges per-task
  JSON, writes eval_info.json compatible with _aggregate_eval_from_per_task
- lerobot-eval-worker CLI: make_env → shard tasks → run episodes → write JSON
- EvalDockerConfig: add port field (default 50051)
- pyproject.toml: add lerobot-eval-worker entry point

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

pyproject.toml

@@ -247,6 +247,7 @@ lerobot-replay="lerobot.scripts.lerobot_replay:main"
 lerobot-setup-motors="lerobot.scripts.lerobot_setup_motors:main"
 lerobot-teleoperate="lerobot.scripts.lerobot_teleoperate:main"
 lerobot-eval="lerobot.scripts.lerobot_eval:main"
+lerobot-eval-worker="lerobot.scripts.lerobot_eval_worker:main"
 lerobot-train="lerobot.scripts.lerobot_train:main"
 lerobot-train-tokenizer="lerobot.scripts.lerobot_train_tokenizer:main"
 lerobot-dataset-viz="lerobot.scripts.lerobot_dataset_viz:main"

src/lerobot/configs/default.py

@@ -64,6 +64,8 @@ class EvalDockerConfig:
     gpus: str | None = "all"
     # Docker --shm-size value (increase when using larger eval.batch_size values).
     shm_size: str = "8g"
+    # Port on which the host HTTP policy inference server listens.
+    port: int = 50051
 
 
 @dataclass

src/lerobot/envs/docker_runtime.py

@@ -0,0 +1,312 @@
+#!/usr/bin/env python
+
+# 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.
+
+"""Docker runtime for lerobot-eval.
+
+The policy stays on the host GPU; gym environments run inside Docker containers.
+Each container runs `lerobot-eval-worker`, which calls back to a host HTTP inference
+server for action chunks.
+
+Architecture:
+    host (GPU):
+        1. Load policy + preprocessors from EvalPipelineConfig.
+        2. Start HTTP policy-inference server (one lock — serialises GPU calls).
+        3. Spawn ``instance_count`` Docker containers (one per shard).
+        4. Wait; collect per-task JSON written to the mounted output volume.
+        5. Merge shards → aggregate → write eval_info.json.
+
+    container (CPU only):
+        1. make_env(cfg.env) → shard tasks by (instance_id, instance_count).
+        2. For each task: run n_episodes, POST obs to /predict_chunk, step env.
+        3. Write per-task JSON to /results/worker_{instance_id}.json.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import pickle  # nosec B403 — internal serialisation only
+import platform
+import subprocess  # nosec B404
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+
+from lerobot.envs.factory import make_env_pre_post_processors
+from lerobot.policies.factory import make_policy, make_pre_post_processors
+from lerobot.utils.utils import get_safe_torch_device
+
+if TYPE_CHECKING:
+    from lerobot.configs.eval import EvalPipelineConfig
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# HTTP inference server (host side)
+# ---------------------------------------------------------------------------
+
+
+class _PolicyInferenceHandler(BaseHTTPRequestHandler):
+    """POST /predict_chunk → pickled numpy action chunk."""
+
+    server: _InferenceServer
+
+    def do_POST(self) -> None:
+        if self.path != "/predict_chunk":
+            self.send_error(404)
+            return
+        length = int(self.headers["Content-Length"])
+        body = self.rfile.read(length)
+        payload: dict = pickle.loads(body)  # nosec B301
+        obs_t: dict = payload["obs_t"]
+
+        with self.server._lock:
+            chunk_np = self.server._predict(obs_t)
+
+        resp = pickle.dumps(chunk_np)  # nosec B301
+        self.send_response(200)
+        self.send_header("Content-Type", "application/octet-stream")
+        self.send_header("Content-Length", str(len(resp)))
+        self.end_headers()
+        self.wfile.write(resp)
+
+    def log_message(self, fmt: str, *args: Any) -> None:  # noqa: ANN401
+        pass  # suppress per-request logs
+
+
+class _InferenceServer(HTTPServer):
+    """Wraps the loaded policy behind a trivial HTTP interface."""
+
+    def __init__(
+        self,
+        addr: tuple[str, int],
+        policy: Any,
+        env_preprocessor: Any,
+        preprocessor: Any,
+        postprocessor: Any,
+    ) -> None:
+        super().__init__(addr, _PolicyInferenceHandler)
+        self._policy = policy
+        self._env_preprocessor = env_preprocessor
+        self._preprocessor = preprocessor
+        self._postprocessor = postprocessor
+        self._lock = threading.Lock()
+        self._device = torch.device(str(policy.config.device))
+
+    def _predict(self, obs_t: dict) -> np.ndarray:
+        """Apply full preprocessing pipeline and return (T, A) numpy chunk."""
+        obs = self._env_preprocessor(obs_t)
+        obs = self._preprocessor(obs)
+        obs_gpu: dict = {k: v.to(self._device) if isinstance(v, torch.Tensor) else v for k, v in obs.items()}
+        with torch.no_grad():
+            chunk: torch.Tensor = self._policy.predict_action_chunk(obs_gpu)  # (B, T, A)
+
+        # Postprocessor expects (B, A); apply it treating each timestep as a batch element.
+        # For linear transforms (unnormalize) this is identical to applying it to (B, T, A).
+        batch, n_steps, action_dim = chunk.shape
+        chunk_2d = chunk.reshape(batch * n_steps, action_dim)  # (B*T, A)
+        chunk_2d = self._postprocessor(chunk_2d)  # (B*T, A)
+        # Return only the first env's chunk — batch_size=1 per container.
+        return chunk_2d[:n_steps].cpu().numpy()  # (T, A)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _get_host_ip() -> str:
+    """Return the IP that containers can use to reach the host."""
+    if platform.system() in ("Darwin", "Windows"):
+        return "host.docker.internal"
+    return "172.17.0.1"  # Linux Docker bridge default gateway
+
+
+def _resolve_image(cfg: EvalPipelineConfig) -> str:
+    """Return the Docker image name to use for the env containers."""
+    if cfg.eval.docker.image:
+        return cfg.eval.docker.image
+    return f"lerobot-benchmark-{cfg.env.type}"
+
+
+def _env_argv() -> list[str]:
+    """Extract --env.* args from sys.argv to forward verbatim to the worker."""
+    return [arg for arg in sys.argv[1:] if arg.startswith("--env.")]
+
+
+def _spawn_container(
+    *,
+    image: str,
+    instance_id: int,
+    instance_count: int,
+    server_address: str,
+    n_episodes: int,
+    seed: int,
+    output_dir: Path,
+    docker_cfg: Any,
+    env_argv: list[str],
+) -> subprocess.Popen:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    container_results = "/results"
+
+    cmd: list[str] = ["docker", "run", "--rm"]
+    if docker_cfg.gpus:
+        cmd += [f"--gpus={docker_cfg.gpus}"]
+    cmd += [f"--shm-size={docker_cfg.shm_size}"]
+    cmd += ["-v", f"{output_dir.resolve()}:{container_results}"]
+    # Allow containers on Linux to resolve host.docker.internal.
+    cmd += ["--add-host=host.docker.internal:host-gateway"]
+    cmd.append(image)
+
+    cmd += [
+        "lerobot-eval-worker",
+        *env_argv,
+        f"--server_address={server_address}",
+        f"--n_episodes={n_episodes}",
+        f"--seed={seed}",
+        f"--instance_id={instance_id}",
+        f"--instance_count={instance_count}",
+        f"--output_path={container_results}/worker_{instance_id}.json",
+    ]
+
+    logger.info(
+        "Spawning container %d/%d: %s",
+        instance_id + 1,
+        instance_count,
+        " ".join(cmd),
+    )
+    return subprocess.Popen(cmd)  # nosec B603 B607
+
+
+# ---------------------------------------------------------------------------
+# Public entry point
+# ---------------------------------------------------------------------------
+
+
+def run_eval_in_docker(cfg: EvalPipelineConfig) -> None:
+    """Run eval with env in Docker containers and policy on the host GPU.
+
+    Writes ``eval_info.json`` to ``cfg.output_dir``.  Called by
+    ``lerobot_eval._run_eval_worker`` when ``eval.runtime == "docker"``.
+    """
+    # Import here to avoid circular import at module level.
+    from lerobot.scripts.lerobot_eval import _aggregate_eval_from_per_task
+
+    start_t = time.time()
+    output_dir = Path(cfg.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    docker_cfg = cfg.eval.docker
+
+    # Optionally pull the image before starting.
+    image = _resolve_image(cfg)
+    if docker_cfg.pull:
+        logger.info("Pulling Docker image: %s", image)
+        subprocess.run(["docker", "pull", image], check=True)  # nosec B603 B607
+
+    # ── Load policy + all preprocessors on the host GPU ──────────────────
+    device = get_safe_torch_device(cfg.policy.device, log=True)
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    policy = make_policy(cfg=cfg.policy, env_cfg=cfg.env, rename_map=cfg.rename_map)
+    policy.eval()
+
+    preprocessor_overrides: dict = {
+        "device_processor": {"device": str(device)},
+        "rename_observations_processor": {"rename_map": cfg.rename_map},
+    }
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=cfg.policy,
+        pretrained_path=cfg.policy.pretrained_path,
+        preprocessor_overrides=preprocessor_overrides,
+    )
+    env_preprocessor, _env_postprocessor = make_env_pre_post_processors(
+        env_cfg=cfg.env,
+        policy_cfg=cfg.policy,
+    )
+
+    # ── Start HTTP inference server ───────────────────────────────────────
+    port = docker_cfg.port
+    server = _InferenceServer(
+        ("0.0.0.0", port),  # nosec B104 — only alive for the duration of eval
+        policy=policy,
+        env_preprocessor=env_preprocessor,
+        preprocessor=preprocessor,
+        postprocessor=postprocessor,
+    )
+    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+    server_thread.start()
+    logger.info("Policy inference server running on port %d", port)
+
+    host_ip = _get_host_ip()
+    server_address = f"{host_ip}:{port}"
+    instance_count = cfg.eval.instance_count
+    env_argv = _env_argv()
+
+    # ── Spawn containers ──────────────────────────────────────────────────
+    container_dirs: list[Path] = []
+    procs: list[subprocess.Popen] = []
+    try:
+        for i in range(instance_count):
+            shard_dir = output_dir / "shards" / str(i)
+            container_dirs.append(shard_dir)
+            proc = _spawn_container(
+                image=image,
+                instance_id=i,
+                instance_count=instance_count,
+                server_address=server_address,
+                n_episodes=cfg.eval.n_episodes,
+                seed=cfg.seed,
+                output_dir=shard_dir,
+                docker_cfg=docker_cfg,
+                env_argv=env_argv,
+            )
+            procs.append(proc)
+
+        failed: list[tuple[int, int]] = []
+        for i, proc in enumerate(procs):
+            rc = proc.wait()
+            if rc != 0:
+                failed.append((i, rc))
+                logger.error("Container %d/%d exited with code %d", i + 1, instance_count, rc)
+        if failed:
+            raise RuntimeError(f"Docker eval containers failed (instance_id, exit_code): {failed}")
+
+    finally:
+        server.shutdown()
+
+    # ── Collect and merge per-task results ───────────────────────────────
+    per_task: list[dict] = []
+    for i, shard_dir in enumerate(container_dirs):
+        result_file = shard_dir / f"worker_{i}.json"
+        with open(result_file) as f:
+            shard_data: dict = json.load(f)
+        per_task.extend(shard_data.get("per_task", []))
+
+    per_task.sort(key=lambda x: (x["task_group"], x["task_id"]))
+
+    info = _aggregate_eval_from_per_task(per_task, total_eval_s=time.time() - start_t)
+    with open(output_dir / "eval_info.json", "w") as f:
+        json.dump(info, f, indent=2)
+
+    logger.info("Docker eval complete. Results: %s/eval_info.json", output_dir)

src/lerobot/scripts/lerobot_eval_worker.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python
+
+# 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.
+
+"""Docker eval worker — runs inside a benchmark container.
+
+Runs gym episodes for a sharded subset of the configured env's tasks, calling
+a remote HTTP policy inference server (running on the host GPU) for action chunks.
+
+Usage (normally invoked by docker_runtime.run_eval_in_docker, not directly):
+    lerobot-eval-worker \\
+        --env.type=libero_plus \\
+        --server_address=host.docker.internal:50051 \\
+        --n_episodes=5 \\
+        --seed=1000 \\
+        --instance_id=0 \\
+        --instance_count=2 \\
+        --output_path=/results/worker_0.json
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import pickle  # nosec B403 — internal serialisation only
+import urllib.request
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+from lerobot import envs  # noqa: F401 — registers all env subclasses
+from lerobot.configs import parser
+from lerobot.envs.configs import EnvConfig
+from lerobot.envs.factory import make_env
+from lerobot.envs.utils import add_envs_task, preprocess_observation
+from lerobot.utils.utils import init_logging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EvalWorkerConfig:
+    env: EnvConfig
+    # Address of the policy inference HTTP server on the host.
+    server_address: str = "host.docker.internal:50051"
+    # Number of episodes to run per task.
+    n_episodes: int = 1
+    # Starting random seed; episode i of a task uses seed + i.
+    seed: int = 0
+    # 0-indexed shard id for this worker.
+    instance_id: int = 0
+    # Total number of shards (workers).
+    instance_count: int = 1
+    # Path (inside the container) to write the JSON per-task results.
+    output_path: Path = field(default_factory=lambda: Path("/results/worker.json"))
+    # Timeout in seconds for each HTTP request to the policy server.
+    server_timeout: float = 120.0
+
+
+def _call_server(server_address: str, obs_t: dict, timeout: float) -> np.ndarray:
+    """POST pickled obs to /predict_chunk, return numpy chunk (T, action_dim)."""
+    body = pickle.dumps({"obs_t": obs_t})  # nosec B301
+    req = urllib.request.Request(
+        f"http://{server_address}/predict_chunk",
+        data=body,
+        method="POST",
+        headers={"Content-Type": "application/octet-stream"},
+    )
+    with urllib.request.urlopen(req, timeout=timeout) as resp:  # nosec B310
+        return pickle.loads(resp.read())  # nosec B301
+
+
+def run_worker(cfg: EvalWorkerConfig) -> dict:
+    """Run cfg.n_episodes episodes per assigned task. Returns per-task results dict."""
+    # Build envs: {task_group: {task_id: vec_env}}
+    envs_dict = make_env(cfg.env, n_envs=1)
+
+    # Flatten to list of (task_group, task_id, env)
+    tasks = [
+        (task_group, task_id, vec)
+        for task_group, group in envs_dict.items()
+        for task_id, vec in group.items()
+    ]
+
+    # Shard: this worker handles tasks where index % instance_count == instance_id
+    if cfg.instance_count > 1:
+        total = len(tasks)
+        tasks = [t for idx, t in enumerate(tasks) if idx % cfg.instance_count == cfg.instance_id]
+        logger.info(
+            "Shard %d/%d: %d/%d tasks assigned.",
+            cfg.instance_id + 1,
+            cfg.instance_count,
+            len(tasks),
+            total,
+        )
+
+    per_task: list[dict] = []
+
+    for task_group, task_id, env in tasks:
+        sum_rewards: list[float] = []
+        max_rewards: list[float] = []
+        successes: list[bool] = []
+
+        for ep_idx in range(cfg.n_episodes):
+            obs, _info = env.reset(seed=[cfg.seed + ep_idx])
+            obs_t = preprocess_observation(obs)
+            obs_t = add_envs_task(env, obs_t)
+
+            action_buffer: list[np.ndarray] = []  # each element: (1, action_dim)
+            ep_rewards: list[float] = []
+            ep_success = False
+            done = np.zeros(1, dtype=bool)
+
+            while not np.all(done):
+                if not action_buffer:
+                    chunk_np = _call_server(cfg.server_address, obs_t, cfg.server_timeout)
+                    # chunk_np: (T, action_dim) — split into per-step slices of shape (1, action_dim)
+                    action_buffer = [chunk_np[i : i + 1] for i in range(chunk_np.shape[0])]
+
+                action_np = action_buffer.pop(0)  # (1, action_dim)
+                obs, reward, terminated, truncated, info = env.step(action_np)
+
+                done = terminated | truncated | done
+                ep_rewards.append(float(np.mean(reward)))
+
+                if "final_info" in info:
+                    final_info = info["final_info"]
+                    if isinstance(final_info, dict) and "is_success" in final_info:
+                        ep_success = bool(final_info["is_success"][0])
+
+                if not np.all(done):
+                    obs_t = preprocess_observation(obs)
+                    obs_t = add_envs_task(env, obs_t)
+
+            sum_rewards.append(float(np.sum(ep_rewards)))
+            max_rewards.append(float(np.max(ep_rewards)) if ep_rewards else 0.0)
+            successes.append(ep_success)
+            logger.info(
+                "Task %s[%d] ep %d/%d — success=%s",
+                task_group,
+                task_id,
+                ep_idx + 1,
+                cfg.n_episodes,
+                ep_success,
+            )
+
+        per_task.append(
+            {
+                "task_group": task_group,
+                "task_id": task_id,
+                "metrics": {
+                    "sum_rewards": sum_rewards,
+                    "max_rewards": max_rewards,
+                    "successes": successes,
+                    "video_paths": [],
+                },
+            }
+        )
+        env.close()
+
+    return {"per_task": per_task}
+
+
+def worker_main(cfg: EvalWorkerConfig) -> None:
+    results = run_worker(cfg)
+    output = Path(cfg.output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(results, indent=2))
+    logger.info("Worker %d wrote results to %s", cfg.instance_id, output)
+
+
+def main() -> None:
+    init_logging()
+    cfg = parser.parse(EvalWorkerConfig)
+    worker_main(cfg)
+
+
+if __name__ == "__main__":
+    main()