diff --git a/configs/README.md b/configs/README.md index 859a8528..d7b5a8ee 100644 --- a/configs/README.md +++ b/configs/README.md @@ -31,7 +31,7 @@ Benchmark names link to their config directory with available YAML files and usa | [Kinetix](benchmarks/kinetix/) | ![◇](https://img.shields.io/badge/◇-blue) | [2410.23208](https://arxiv.org/abs/2410.23208) | [`kinetix`](https://ghcr.io/allenai/vla-evaluation-harness/kinetix) | 3.11 | Physics-based 2D manipulation (JAX) | | [MolmoSpaces-Bench](benchmarks/molmospaces/) | ![✓](https://img.shields.io/badge/✓-teal) | [2603.16861](https://arxiv.org/abs/2603.16861) | [`molmospaces`](https://ghcr.io/allenai/vla-evaluation-harness/molmospaces) | 3.11 | Spatial reasoning (AI2-THOR) | | [RLBench](benchmarks/rlbench/) | ![◇](https://img.shields.io/badge/◇-blue) | [1909.12271](https://arxiv.org/abs/1909.12271) | 🔒 `rlbench` | 3.8 | Vision-guided manipulation (CoppeliaSim) | -| [RoboCasa](benchmarks/robocasa/) | ![◇](https://img.shields.io/badge/◇-blue) | [2406.02523](https://arxiv.org/abs/2406.02523) | [`robocasa`](https://ghcr.io/allenai/vla-evaluation-harness/robocasa) | 3.11 | Kitchen manipulation (MuJoCo) | +| [RoboCasa365](benchmarks/robocasa/) | ![◇](https://img.shields.io/badge/◇-blue) | [2603.04356](https://arxiv.org/abs/2603.04356) | [`robocasa`](https://ghcr.io/allenai/vla-evaluation-harness/robocasa) | 3.11 | 50-task kitchen manipulation (MuJoCo) | | [VLABench](benchmarks/vlabench/) | ![◇](https://img.shields.io/badge/◇-blue) | [2502.09858](https://arxiv.org/abs/2502.09858) | [`vlabench`](https://ghcr.io/allenai/vla-evaluation-harness/vlabench) | 3.10 | Language-conditioned manipulation (SAPIEN) | | [MIKASA-Robo](benchmarks/mikasa/) | ![◇](https://img.shields.io/badge/◇-blue) | [2502.07007](https://arxiv.org/abs/2502.07007) | [`mikasa-robo`](https://ghcr.io/allenai/vla-evaluation-harness/mikasa-robo) | 3.10 | Robot manipulation (MuJoCo) | | [RoboCerebra](benchmarks/robocerebra/) | ![◇](https://img.shields.io/badge/◇-blue) | [2502.02853](https://arxiv.org/abs/2502.02853) | [`robocerebra`](https://ghcr.io/allenai/vla-evaluation-harness/robocerebra) | 3.8 | Cognitive manipulation (MuJoCo) | diff --git a/configs/benchmarks/robocasa/README.md b/configs/benchmarks/robocasa/README.md index e6d1a722..24dd1701 100644 --- a/configs/benchmarks/robocasa/README.md +++ b/configs/benchmarks/robocasa/README.md @@ -1,16 +1,136 @@ --- -smoke_config: eval.yaml +smoke_config: smoke.yaml --- -# RoboCasa +# RoboCasa365 -Kitchen manipulation benchmark (MuJoCo/robosuite). -[Paper](https://arxiv.org/abs/2406.02523) | [GitHub](https://github.com/robocasa/robocasa) +Official 50-task multi-task benchmark using RoboCasa's Panda-Omron Gymnasium wrapper. +[Paper](https://robocasa.ai/assets/robocasa365_iclr26.pdf) · +[Code](https://github.com/robocasa/robocasa) · +[Protocol](https://robocasa.ai/docs/build/html/benchmarking/multitask_learning.html) · +[Leaderboard](https://robocasa.ai/leaderboard.html) **Docker image:** `ghcr.io/allenai/vla-evaluation-harness/robocasa:latest` ## Configs | File | Description | Tasks | Episodes/task | -|------|-------------|:-----:|:-------------:| -| `eval.yaml` | Full RoboCasa evaluation | 24 | 50 | +| --- | --- | ---: | ---: | +| `rc365.yaml` | Official multi-task protocol | 50 | 50 | +| `rc365_s2b_qualification.yaml` | Base config for one S2B qualification shard | 1 | 1 | +| `smoke.yaml` | Contract smoke test | 2 | 1 | + +The adapter reads task membership and task-specific horizons from RoboCasa's registry. +It evaluates the target50 tasks in pretraining kitchens (`split: pretrain`), matching the official multi-task leaderboard protocol. +Success is sampled and accumulated at the end of each 16-action chunk, while every episode continues to its task-specific horizon. +The wire contract preserves all 12 Panda-Omron dimensions in the official dataset order instead of padding a 7-D arm action. + +## Qualification-slice parity + +The parity script compares the isolated RC365 adapter with the read-only +`rc365_s2b.environments.OfficialRoboCasaEnv` for the six System 1 +qualification tasks at seed 0. It checks repeated-reset determinism, raw +observation keys, shapes, dtypes, values, and simulator state across the same +recorded 32-step action sequence. + +```bash +SIM_PYTHON=/path/to/system1/.venv/bin/python +REFERENCE_SRC=/path/to/projects/rc365-s2b/src +MUJOCO_GL=egl NUMBA_DISABLE_JIT=1 "$SIM_PYTHON" \ + scripts/check_rc365_adapter_parity.py \ + --reference-src "$REFERENCE_SRC" +``` + +The output is `results/rc365_adapter_qualification_parity.json`. The script +redirects RoboCasa's generated temporary XML files to `/tmp`; it does not +modify the simulator environment or the reference source tree. Full GR00T or +MLLM inference is outside adapter parity and remains GPU or Slurm-deferred. + +## Implementation boundary + +- Reused upstream: the registered RoboCasa Gym environment, task registry, pretrain split, per-task horizon, observation schema, and strict success predicate. +- Implemented here: canonical observation mapping and lossless 12-D named-action decoding. +- Dependency boundary: Docker verifies RoboCasa `1.0.1` and robosuite `1.5.2` at tested patch revisions because those upstream fixes have no new semantic tags. + +## CPU rendering + +Set `VLA_EVAL_RENDER=cpu` to select Mesa llvmpipe through OSMesa. Set +`docker.gpus: none`, or pass `--gpus none`, so the benchmark container receives +no GPU device. The model server remains on the host GPU. GPU rendering through +the EGL device platform remains the default. + +```bash +VLA_EVAL_RENDER=cpu vla-eval run \ + --config configs/benchmarks/robocasa/smoke.yaml \ + --gpus none +``` + +CPU mode sets `MUJOCO_GL=osmesa` and `PYOPENGL_PLATFORM=osmesa`, forces +software GL, and clears `EGL_PLATFORM` and `MUJOCO_EGL_DEVICE_ID` before +MuJoCo is imported. + +Container verification requires a Docker host: + +```bash +docker build -f docker/Dockerfile.base -t vla-eval-base:rc365-cpu . +docker build \ + --build-arg BASE_IMAGE=vla-eval-base:rc365-cpu \ + -f docker/Dockerfile.robocasa \ + -t vla-eval-robocasa:rc365-cpu . +docker run --rm \ + -e VLA_EVAL_RENDER=cpu \ + -e NVIDIA_VISIBLE_DEVICES=void \ + -e CUDA_VISIBLE_DEVICES=-1 \ + --entrypoint /opt/conda/bin/conda \ + vla-eval-robocasa:rc365-cpu \ + run --no-capture-output -n robocasa \ + python /workspace/scripts/check_robocasa_cpu_render.py +``` + +## RC365 S2B qualification + +The qualification runner expands the reference seed manifest in stable task +and seed order, then interleaves `gold-s2`, `global-s1`, and `random-valid`. +Each invocation owns exactly one task, seed, and condition. Qualification-only +success checks and 16-step chunk telemetry live in +`RoboCasaS2BQualificationBenchmark`, leaving the upstream adapters unchanged. + +```bash +export RC365_S2B_ROOT=/path/to/projects/rc365-s2b +export RC365_PHASE_MANIFEST=/path/to/phase_manifest.jsonl +export ROBOCASA_GR00T_N15_CKPT=/path/to/checkpoint-120000 +export ROBOCASA_MODALITY_JSON=/path/to/modality.json +export VLA_EVAL_ROBOCASA_IMAGE=vla-eval-robocasa:rc365-cpu + +uv run python -m vla_eval.rc365_s2b_qualification resolve \ + --rung dev --array-index 0 \ + --seed-manifest "$RC365_S2B_ROOT/config/qualification_seeds.json" + +uv run python -m vla_eval.rc365_s2b_qualification run \ + --rung dev --array-index 0 \ + --seed-manifest "$RC365_S2B_ROOT/config/qualification_seeds.json" \ + --phase-manifest "$RC365_PHASE_MANIFEST" \ + --render cpu --dev +``` + +Outputs are one-record JSONL files under +`results/rc365_s2b_qualification///`. Records use schema +`rc365-s2b-exec-episode-v1` and include strict success, termination, steps, +chunks, System 2 calls, run configuration, and provenance. + +```bash +rc365-s2b score-qualification \ + results/rc365_s2b_qualification/dev/gold-s2 \ + results/rc365_s2b_qualification/dev/global-s1 \ + results/rc365_s2b_qualification/dev/random-valid +``` + +Submit the Slurm template with the rung's array bound. CPU rendering requests +one GPU for the policy. GPU rendering needs two GPUs. + +```bash +# dev: 0-17, training-dev: 0-89, frozen: 0-449 +sbatch --array=0-17 \ + --export=ALL,QUALIFICATION_RUNG=dev,VLA_EVAL_RENDER=cpu,VLA_EVAL_ROBOCASA_IMAGE=vla-eval-robocasa:rc365-cpu \ + scripts/run_rc365_s2b_qualification.sbatch +``` diff --git a/configs/benchmarks/robocasa/eval.yaml b/configs/benchmarks/robocasa/eval.yaml index e530fbf2..66b0c83b 100644 --- a/configs/benchmarks/robocasa/eval.yaml +++ b/configs/benchmarks/robocasa/eval.yaml @@ -1,14 +1,11 @@ -# RoboCasa kitchen manipulation evaluation (6 tasks × 5 episodes) -# -# RoboCasa provides 365 kitchen tasks on 2500+ procedurally-generated scenes. -# Default config uses a small subset of atomic tasks for quick evaluation. -# -# Requires ~10GB of kitchen assets from HuggingFace. +# RoboCasa kitchen manipulation, default atomic subset (6 tasks x 5 episodes). +# Needs ~10GB of kitchen assets from HuggingFace. server: url: "ws://localhost:8000" docker: image: ghcr.io/allenai/vla-evaluation-harness/robocasa:latest + env: [VLA_EVAL_RENDER] output_dir: "./results" @@ -21,4 +18,3 @@ benchmarks: max_steps: 500 split: pretrain seed: null - diff --git a/configs/benchmarks/robocasa/rc365.yaml b/configs/benchmarks/robocasa/rc365.yaml new file mode 100644 index 00000000..337b40c8 --- /dev/null +++ b/configs/benchmarks/robocasa/rc365.yaml @@ -0,0 +1,19 @@ +# RoboCasa365 official multi-task evaluation: 50 tasks x 50 seeded episodes. +server: + url: "ws://localhost:8000" + +docker: + image: ghcr.io/allenai/vla-evaluation-harness/robocasa:latest + env: [VLA_EVAL_RENDER] + +output_dir: "./results" + +benchmarks: + - benchmark: "vla_eval.benchmarks.robocasa.rc365:RoboCasa365Benchmark" + episodes_per_task: 50 + params: + split: pretrain + seed: 0 + camera_size: 256 + enable_render: true + success_check_interval: 16 diff --git a/configs/benchmarks/robocasa/smoke.yaml b/configs/benchmarks/robocasa/smoke.yaml new file mode 100644 index 00000000..741dbbec --- /dev/null +++ b/configs/benchmarks/robocasa/smoke.yaml @@ -0,0 +1,21 @@ +# Short contract smoke test; this is not a leaderboard protocol. +server: + url: "ws://localhost:8000" + +docker: + image: ghcr.io/allenai/vla-evaluation-harness/robocasa:latest + env: [VLA_EVAL_RENDER] + +output_dir: "./results" + +benchmarks: + - benchmark: "vla_eval.benchmarks.robocasa.rc365:RoboCasa365Benchmark" + episodes_per_task: 1 + max_steps: 5 + params: + tasks: [KettleBoiling, ArrangeBreadBasket] + split: pretrain + seed: 0 + camera_size: 256 + enable_render: true + success_check_interval: 16 diff --git a/configs/model_servers/groot/README.md b/configs/model_servers/groot/README.md index 19c8f1ac..29796168 100644 --- a/configs/model_servers/groot/README.md +++ b/configs/model_servers/groot/README.md @@ -2,7 +2,7 @@ smoke_config: groot.yaml --- -# GR00T N1.6 +# GR00T 3B dual-system VLA from NVIDIA. [Paper](https://arxiv.org/abs/2503.14734) | [GitHub](https://github.com/NVIDIA/Isaac-GR00T) @@ -14,3 +14,8 @@ smoke_config: groot.yaml | `libero.yaml` | LIBERO | `nvidia/GR00T-N1.6-3B` | | `simpler_widowx.yaml` | SimplerEnv WidowX | `nvidia/GR00T-N1.6-bridge` | | `simpler_google_robot.yaml` | SimplerEnv GR | `nvidia/GR00T-N1.6-fractal` | +| `robocasa_n15.yaml` | RoboCasa365 | `robocasa/robocasa365_checkpoints` | + +For RoboCasa365, set `ROBOCASA_GR00T_N15_CKPT` to `gr00t_n1-5/multitask_learning/checkpoint-120000` from `robocasa/robocasa365_checkpoints`. +The pinned RoboCasa GR00T fork directly imports FlashAttention, so its script pins the exercised FlashAttention build instead of using the Transformers `kernels` extra. +The configured seed fixes the diffusion-noise sequence. diff --git a/configs/model_servers/groot/robocasa_n15.yaml b/configs/model_servers/groot/robocasa_n15.yaml new file mode 100644 index 00000000..9c6afd92 --- /dev/null +++ b/configs/model_servers/groot/robocasa_n15.yaml @@ -0,0 +1,10 @@ +# RoboCasa365 GR00T N1.5 multitask checkpoint. +# Download robocasa/robocasa365_checkpoints and point the environment variable +# at gr00t_n1-5/multitask_learning/checkpoint-120000. +script: "src/vla_eval/model_servers/robocasa_groot.py" +args: + model_path: "${oc.env:ROBOCASA_GR00T_N15_CKPT,${oc.env:VLA_EVAL_ASSETS_CACHE,${oc.env:XDG_CACHE_HOME,${oc.env:HOME}/.cache}/vla-eval/assets}/robocasa/gr00t_n1-5/multitask_learning/checkpoint-120000}" + chunk_size: 16 + denoising_steps: 4 + seed: 0 + port: 8000 diff --git a/configs/model_servers/rc365_s2b/README.md b/configs/model_servers/rc365_s2b/README.md new file mode 100644 index 00000000..cdc9d872 --- /dev/null +++ b/configs/model_servers/rc365_s2b/README.md @@ -0,0 +1,12 @@ +# RC365 S2B hierarchical policy + +One stateful policy wrapping the S2B System 2 planner and the GR00T +System 1 executor. System 2 is re-queried after every 16-step action +chunk: an unchanged call continues the current instruction, a new call +switches it, and finish_task terminates the episode (surfaced to the +runner as `terminate_episode`). + +`hierarchical.yaml` selects the System 2 implementation (`gold`, +`global_only`, `random_valid`, or `mllm`) and the checkpoint paths. +Set `VLA_EVAL_RENDER=cpu` to render through OSMesa and keep the GPU +exclusively for the policy. diff --git a/configs/model_servers/rc365_s2b/hierarchical.yaml b/configs/model_servers/rc365_s2b/hierarchical.yaml new file mode 100644 index 00000000..7a56b1a9 --- /dev/null +++ b/configs/model_servers/rc365_s2b/hierarchical.yaml @@ -0,0 +1,11 @@ +# RC365 S2B hierarchical policy. The reference package must be on PYTHONPATH. +script: "src/vla_eval/model_servers/rc365_s2b.py" +args: + checkpoint: "${oc.env:ROBOCASA_GR00T_N15_CKPT,${oc.env:VLA_EVAL_ASSETS_CACHE,${oc.env:XDG_CACHE_HOME,${oc.env:HOME}/.cache}/vla-eval/assets}/robocasa/gr00t_n1-5/multitask_learning/checkpoint-120000}" + modality_path: "${oc.env:ROBOCASA_MODALITY_JSON,/path/to/modality.json}" + registry_path: "${oc.env:RC365_S2B_ROOT,/path/to/rc365-s2b}/config/registry_composite.json" + system2: global-only + device: cuda + denoising_steps: 4 + seed: 0 + port: 8000 diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index 8beba8ef..bd5d38c2 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -1,6 +1,6 @@ # Common base image for all VLA evaluation harness benchmarks # -# Provides: CUDA 12.1 runtime, EGL/Vulkan rendering, Miniforge, uv +# Provides: CUDA 12.1 runtime, NVIDIA or Mesa rendering, Miniforge, uv # FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04 @@ -18,8 +18,9 @@ ENV DEBIAN_FRONTEND=noninteractive \ RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential cmake git wget curl ca-certificates unzip \ # EGL / OpenGL - libegl1 libegl1-mesa-dev libgl1 libgl1-mesa-glx libgles2 \ + libegl1 libegl-mesa0 libegl1-mesa-dev libgl1 libgl1-mesa-glx libglx-mesa0 libgles2 \ libglfw3 libosmesa6 libosmesa6-dev libglew-dev libglvnd-dev \ + mesa-utils \ # Vulkan libvulkan1 mesa-vulkan-drivers \ # X11 / misc diff --git a/docker/Dockerfile.robocasa b/docker/Dockerfile.robocasa index 93296950..f2e623a1 100644 --- a/docker/Dockerfile.robocasa +++ b/docker/Dockerfile.robocasa @@ -7,7 +7,8 @@ RUN conda create -n robocasa python=3.11 pip -y && conda clean -afy SHELL ["conda", "run", "-n", "robocasa", "/bin/bash", "-c"] -ARG ROBOCASA_ROBOSUITE_REF=85abee228d1c43ab1939bce33028099945d453b4 +ARG ROBOCASA_ROBOSUITE_VERSION=1.5.2 +ARG ROBOCASA_ROBOSUITE_REF=5ce6643f3092639d08f7b0f90ed1c6a84f50552c RUN mkdir -p /app/robosuite && \ cd /app/robosuite && \ git init -q && \ @@ -15,9 +16,11 @@ RUN mkdir -p /app/robosuite && \ git fetch -q --depth 1 origin "${ROBOCASA_ROBOSUITE_REF}" && \ git checkout -q FETCH_HEAD && \ uv pip install --no-cache-dir -e . && \ + test "$(python -c 'from importlib.metadata import version; print(version("robosuite"))')" = "${ROBOCASA_ROBOSUITE_VERSION}" && \ rm -rf /app/robosuite/.git -ARG ROBOCASA_REF=v1.0 +ARG ROBOCASA_VERSION=1.0.1 +ARG ROBOCASA_REF=b4684e6ee37d377cc392e98302a6b916d588b415 RUN mkdir -p /app/robocasa && \ cd /app/robocasa && \ git init -q && \ @@ -25,6 +28,7 @@ RUN mkdir -p /app/robocasa && \ git fetch -q --depth 1 origin "${ROBOCASA_REF}" && \ git checkout -q FETCH_HEAD && \ uv pip install --no-cache-dir -e . && \ + test "$(python -c 'from importlib.metadata import version; print(version("robocasa"))')" = "${ROBOCASA_VERSION}" && \ rm -rf /app/robocasa/.git # Setup macros (non-interactive) and download kitchen assets (~10 GB) @@ -39,6 +43,7 @@ ARG HARNESS_VERSION=0.0.0 ENV SETUPTOOLS_SCM_PRETEND_VERSION=${HARNESS_VERSION} RUN uv pip install --no-cache-dir -e . COPY configs/ configs/ +COPY scripts/check_robocasa_cpu_render.py scripts/ ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "robocasa", "vla-eval"] -CMD ["run", "--config", "/workspace/configs/robocasa_eval.yaml"] +CMD ["run", "--config", "/workspace/configs/benchmarks/robocasa/rc365.yaml"] diff --git a/scripts/check_robocasa_cpu_render.py b/scripts/check_robocasa_cpu_render.py new file mode 100644 index 00000000..c2316fb4 --- /dev/null +++ b/scripts/check_robocasa_cpu_render.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Create, step, and render the three RC365 cameras with Mesa llvmpipe.""" + +from __future__ import annotations + +import json +import os + +os.environ["VLA_EVAL_RENDER"] = "cpu" +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1") + +import numpy as np + +from vla_eval.benchmarks.robocasa.rc365 import ACTION_DIM, VIDEO_KEYS, RoboCasa365Benchmark +from vla_eval.recording import NullEpisodeRecorder + + +def main() -> None: + benchmark = RoboCasa365Benchmark( + tasks=["RinseSinkBasin"], + split="pretrain", + seed=0, + max_steps=16, + ) + benchmark._recorder = NullEpisodeRecorder() + try: + task = {"name": "RinseSinkBasin", "episode_idx": 0} + raw = benchmark.reset(task) + first = benchmark.make_obs(raw, task) + stepped = benchmark.step({"actions": np.zeros(ACTION_DIM, dtype=np.float32)}) + second = benchmark.make_obs(stepped.obs, task) + for observation in (first, second): + images = observation["images"] + missing = [key for key in VIDEO_KEYS if key not in images] + if missing: + raise RuntimeError(f"missing rendered cameras: {missing}") + for key in VIDEO_KEYS: + image = np.asarray(images[key]) + if image.shape != (256, 256, 3) or image.dtype != np.uint8: + raise RuntimeError(f"unexpected {key} image: shape={image.shape}, dtype={image.dtype}") + print( + json.dumps( + { + "camera_shapes": {key: list(np.asarray(second["images"][key]).shape) for key in VIDEO_KEYS}, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "egl_platform": os.environ.get("EGL_PLATFORM"), + "mujoco_gl": os.environ.get("MUJOCO_GL"), + "pyopengl_platform": os.environ.get("PYOPENGL_PLATFORM"), + "pass": True, + }, + sort_keys=True, + ) + ) + finally: + benchmark.cleanup() + + +if __name__ == "__main__": + main() diff --git a/src/vla_eval/benchmarks/robocasa/__init__.py b/src/vla_eval/benchmarks/robocasa/__init__.py index 29b0fe1c..8564f3b3 100644 --- a/src/vla_eval/benchmarks/robocasa/__init__.py +++ b/src/vla_eval/benchmarks/robocasa/__init__.py @@ -1,5 +1,6 @@ -"""RoboCasa kitchen manipulation benchmark.""" +"""RoboCasa and RoboCasa365 kitchen manipulation benchmarks.""" from vla_eval.benchmarks.robocasa.benchmark import RoboCasaBenchmark +from vla_eval.benchmarks.robocasa.rc365 import RoboCasa365Benchmark -__all__ = ["RoboCasaBenchmark"] +__all__ = ["RoboCasa365Benchmark", "RoboCasaBenchmark"] diff --git a/src/vla_eval/benchmarks/robocasa/benchmark.py b/src/vla_eval/benchmarks/robocasa/benchmark.py index a04d9fcb..6bd8320c 100644 --- a/src/vla_eval/benchmarks/robocasa/benchmark.py +++ b/src/vla_eval/benchmarks/robocasa/benchmark.py @@ -16,6 +16,8 @@ from __future__ import annotations import os +import random +from collections.abc import MutableMapping from typing import Any import numpy as np @@ -24,7 +26,33 @@ from vla_eval.specs import GRIPPER_RAW, IMAGE_RGB, LANGUAGE, POSITION_DELTA, ROTATION_EULER, DimSpec from vla_eval.types import Action, EpisodeResult, Observation, Task -os.environ.setdefault("MUJOCO_GL", "egl") +# Mesa llvmpipe does not expose the EGL device extension MuJoCo needs. +_CPU_RENDER_ENV = { + "MUJOCO_GL": "osmesa", + "PYOPENGL_PLATFORM": "osmesa", + "LIBGL_ALWAYS_SOFTWARE": "1", +} +_CPU_RENDER_UNSET = ("MUJOCO_EGL_DEVICE_ID", "EGL_PLATFORM") + + +def configure_robocasa_rendering(environ: MutableMapping[str, str] | None = None) -> str: + """Select NVIDIA EGL or Mesa llvmpipe before MuJoCo is imported.""" + target = os.environ if environ is None else environ + mode = target.get("VLA_EVAL_RENDER", "gpu").strip().lower() or "gpu" + if mode == "cpu": + target.update(_CPU_RENDER_ENV) + for key in _CPU_RENDER_UNSET: + target.pop(key, None) + elif mode == "gpu": + target.setdefault("MUJOCO_GL", "egl") + target.setdefault("PYOPENGL_PLATFORM", "egl") + target.setdefault("EGL_PLATFORM", "device") + else: + raise ValueError("VLA_EVAL_RENDER must be 'cpu' or 'gpu'") + return mode + + +RENDER_BACKEND = configure_robocasa_rendering() # Subset of atomic tasks suitable for quick evaluation. DEFAULT_TASKS = [ @@ -76,6 +104,7 @@ def __init__( self._seed = seed self._env: Any = None self._current_task: str | None = None + self._current_episode_seed: int | None = None self._lang: str = "" def cleanup(self) -> None: @@ -89,27 +118,38 @@ def cleanup(self) -> None: def get_tasks(self) -> list[Task]: return [{"name": t} for t in self._task_names] - def reset(self, task: Task) -> Any: + def _make_env(self, task_name: str, *, episode_seed: int | None) -> Any: from robocasa.utils.env_utils import create_env + return create_env( + env_name=task_name, + robots=self._robot, + camera_names=self._camera_names, + camera_widths=self._camera_size, + camera_heights=self._camera_size, + render_onscreen=False, + split=self._split, + seed=episode_seed, + ) + + def reset(self, task: Task) -> Any: task_name = task["name"] + episode_idx = int(task.get("episode_idx", 0)) + episode_seed = None if self._seed is None else self._seed + episode_idx + if episode_seed is not None: + random.seed(episode_seed) + np.random.seed(episode_seed) - # Reuse env for same task across episodes - if self._env is None or self._current_task != task_name: + if self._env is None or self._current_task != task_name or self._current_episode_seed != episode_seed: if self._env is not None: self._env.close() - self._env = create_env( - env_name=task_name, - robots=self._robot, - camera_names=self._camera_names, - camera_widths=self._camera_size, - camera_heights=self._camera_size, - render_onscreen=False, - split=self._split, - seed=self._seed, - ) + self._env = self._make_env(task_name, episode_seed=episode_seed) self._current_task = task_name + self._current_episode_seed = episode_seed + if episode_seed is not None: + random.seed(episode_seed) + np.random.seed(episode_seed) obs = self._env.reset() self._lang = self._env.get_ep_meta().get("lang", task_name) self._recorder.record_video(self._extract_frame(obs)) diff --git a/src/vla_eval/benchmarks/robocasa/rc365.py b/src/vla_eval/benchmarks/robocasa/rc365.py new file mode 100644 index 00000000..fc723d75 --- /dev/null +++ b/src/vla_eval/benchmarks/robocasa/rc365.py @@ -0,0 +1,246 @@ +"""Official RoboCasa365 multi-task benchmark adapter.""" + +from __future__ import annotations + +import random +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np + +from vla_eval.benchmarks.base import StepBenchmark, StepResult +from vla_eval.benchmarks.robocasa.benchmark import RENDER_BACKEND +from vla_eval.specs import IMAGE_RGB, LANGUAGE, POSITION_DELTA, RAW, ROTATION_AA, DimSpec +from vla_eval.types import Action, EpisodeResult, Observation, Task + +OFFICIAL_TASK_SETS = ("atomic_seen", "composite_seen", "composite_unseen") +VIDEO_KEYS = ( + "video.robot0_agentview_left", + "video.robot0_agentview_right", + "video.robot0_eye_in_hand", +) +STATE_KEYS = ( + "state.end_effector_position_relative", + "state.end_effector_rotation_relative", + "state.gripper_qpos", + "state.base_position", + "state.base_rotation", +) +ACTION_COMPONENTS = ( + # Official LeRobot modality.json flat-action order. + ("action.base_motion", 4), + ("action.control_mode", 1), + ("action.end_effector_position", 3), + ("action.end_effector_rotation", 3), + ("action.gripper_close", 1), +) +ACTION_DIM = sum(width for _, width in ACTION_COMPONENTS) +BASE_MOTION = DimSpec("base_motion", 4, "panda_omron_base_velocity3_torso", (-1, 1)) +CONTROL_MODE_01 = DimSpec("control_mode", 1, "panda_omron_control_mode_01", (0, 1)) +GRIPPER_BINARY_CLOSE_01 = DimSpec("gripper", 1, "binary_close_01", (0, 1)) + + +def _task_registry() -> Mapping[str, Sequence[str]]: + from robocasa.utils.dataset_registry import TASK_SET_REGISTRY + + return TASK_SET_REGISTRY + + +def _task_horizon(task_name: str) -> int: + from robocasa.utils.dataset_registry_utils import get_task_horizon + + return int(get_task_horizon(task_name)) + + +def _decode_panda_omron_action(action: Action) -> dict[str, np.ndarray]: + raw = action.get("actions", action.get("action")) + if raw is None: + raise ValueError("RoboCasa365 requires an 'actions' vector") + raw = np.asarray(raw, dtype=np.float64) + if raw.shape != (ACTION_DIM,): + raise ValueError(f"RoboCasa365 expected a {ACTION_DIM}-D action, got {raw.shape}") + + named = {} + offset = 0 + for key, width in ACTION_COMPONENTS: + named[key] = raw[offset : offset + width] + offset += width + return named + + +class RoboCasa365Benchmark(StepBenchmark): + """Official target50 benchmark using RoboCasa's Panda-Omron Gym wrapper.""" + + _ALL_RECORD_FIELDS = frozenset({"reward", "done", "success"}) + + def __init__( + self, + tasks: list[str] | None = None, + camera_size: int = 256, + max_steps: int | None = None, + split: str = "pretrain", + seed: int | None = None, + enable_render: bool = True, + success_check_interval: int = 16, + ) -> None: + super().__init__() + if split not in {"pretrain", "target"}: + raise ValueError("split must be 'pretrain' or 'target'") + if camera_size <= 0: + raise ValueError("camera_size must be positive") + if max_steps is not None and max_steps <= 0: + raise ValueError("max_steps must be positive when set") + if success_check_interval <= 0: + raise ValueError("success_check_interval must be positive") + + self._explicit_tasks = list(tasks) if tasks is not None else None + self._camera_size = camera_size + self._max_steps_override = max_steps + self._split = split + self._seed = seed + self._enable_render = enable_render + self._success_check_interval = success_check_interval + self._resolved_tasks: list[Task] | None = None + self._env: Any = None + self._current_task: str | None = None + self._current_episode_seed: int | None = None + self._current_horizon = 0 + self._steps = 0 + self._episode_success = False + + def cleanup(self) -> None: + if self._env is not None: + try: + self._env.close() + except Exception: + # Cleanup must not mask an episode or runner failure. + pass + finally: + self._env = None + + def get_tasks(self) -> list[Task]: + if self._resolved_tasks is None: + registry = _task_registry() + task_to_suite = {task: suite for suite in OFFICIAL_TASK_SETS for task in registry[suite]} + if self._explicit_tasks is not None: + unknown = sorted(set(self._explicit_tasks) - set(task_to_suite)) + if unknown: + raise ValueError(f"tasks are not in the official target50 registry: {unknown}") + self._resolved_tasks = [{"name": task, "suite": task_to_suite[task]} for task in self._explicit_tasks] + else: + self._resolved_tasks = [ + {"name": task, "suite": suite} for suite in OFFICIAL_TASK_SETS for task in registry[suite] + ] + return [dict(task) for task in self._resolved_tasks] + + def _make_env(self, task_name: str, *, episode_seed: int | None) -> Any: + import gymnasium as gym + import robocasa.wrappers.gym_wrapper # noqa: F401 # registers robocasa/* Gym environments + + return gym.make( + f"robocasa/{task_name}", + split=self._split, + enable_render=self._enable_render, + camera_widths=self._camera_size, + camera_heights=self._camera_size, + robots="PandaOmron", + randomize_cameras=False, + generative_textures=None, + translucent_robot=False, + horizon=self._max_steps_override or _task_horizon(task_name), + seed=episode_seed, + disable_env_checker=True, + ) + + def reset(self, task: Task) -> Any: + task_name = task["name"] + episode_idx = int(task.get("episode_idx", 0)) + episode_seed = None if self._seed is None else self._seed + episode_idx + if episode_seed is not None: + random.seed(episode_seed) + np.random.seed(episode_seed) + if self._env is None or self._current_task != task_name or self._current_episode_seed != episode_seed: + if self._env is not None: + self._env.close() + self._env = self._make_env(task_name, episode_seed=episode_seed) + self._current_task = task_name + self._current_episode_seed = episode_seed + + if episode_seed is not None: + random.seed(episode_seed) + np.random.seed(episode_seed) + obs, _ = self._env.reset(seed=episode_seed) + self._steps = 0 + self._episode_success = False + self._current_horizon = self._max_steps_override or _task_horizon(task_name) + self._recorder.record_video(self._extract_frame(obs)) + return obs + + def step(self, action: Action) -> StepResult: + named_action = _decode_panda_omron_action(action) + obs, _, _, _, info = self._env.step(named_action) + self._steps += 1 + done = self._steps >= self._current_horizon + # Match the official MultiStepWrapper: sample success at policy-chunk + # endpoints and continue to the registry horizon. + if self._steps % self._success_check_interval == 0 or done: + self._episode_success |= bool(info.get("success", False) or self._env.unwrapped.env._check_success()) + info = {**info, "success": self._episode_success} + self._recorder.record_video(self._extract_frame(obs)) + self._recorder.record_step(reward=float(self._episode_success), done=done, success=self._episode_success) + return StepResult(obs=obs, reward=float(self._episode_success), done=done, info=info) + + @staticmethod + def _extract_frame(raw_obs: Any) -> np.ndarray | None: + if not isinstance(raw_obs, Mapping): + return None + frame = raw_obs.get(VIDEO_KEYS[0]) + return None if frame is None else np.ascontiguousarray(frame) + + def make_obs(self, raw_obs: Any, task: Task) -> Observation: + missing = [key for key in (*VIDEO_KEYS, *STATE_KEYS) if key not in raw_obs] + if missing: + raise KeyError(f"RoboCasa Gym observation is missing official fields: {missing}") + return { + "images": {key: np.ascontiguousarray(raw_obs[key]) for key in VIDEO_KEYS}, + "state": {key: np.asarray(raw_obs[key]) for key in STATE_KEYS}, + "task_description": str(raw_obs.get("annotation.human.task_description", task["name"])), + } + + def check_done(self, step_result: StepResult) -> bool: + return step_result.done + + def get_step_result(self, step_result: StepResult) -> EpisodeResult: + return {"success": bool(step_result.info.get("success", False))} + + def get_metadata(self) -> dict[str, Any]: + tasks = self.get_tasks() + max_steps = self._max_steps_override or max(_task_horizon(task["name"]) for task in tasks) + return { + "max_steps": max_steps, + "environment_split": self._split, + "environment_seed": self._seed, + "success_check_interval": self._success_check_interval, + "task_horizon_source": "robocasa.utils.dataset_registry_utils.get_task_horizon", + "render_backend": RENDER_BACKEND, + } + + def get_action_spec(self) -> dict[str, DimSpec]: + return { + "position": POSITION_DELTA, + "rotation": ROTATION_AA, + "gripper": GRIPPER_BINARY_CLOSE_01, + "base_motion": BASE_MOTION, + "control_mode": CONTROL_MODE_01, + } + + def get_observation_spec(self) -> dict[str, DimSpec]: + return {"image": IMAGE_RGB, "state": RAW, "language": LANGUAGE} + + def render(self) -> np.ndarray | None: + if self._env is None: + return None + try: + return self._env.render() + except Exception: + return None diff --git a/src/vla_eval/model_servers/rc365_s2b.py b/src/vla_eval/model_servers/rc365_s2b.py new file mode 100644 index 00000000..7ddd6489 --- /dev/null +++ b/src/vla_eval/model_servers/rc365_s2b.py @@ -0,0 +1,629 @@ +# /// script +# requires-python = "~=3.11" +# dependencies = [ +# "vla-eval", +# "diffusers==0.30.2", +# "flash-attn==2.7.4.post1", +# "gr00t @ git+https://github.com/robocasa-benchmark/Isaac-GR00T.git@9d7d7a9eb7ad30bd8ce30448d9ab53a918b45b10", +# "ninja==1.13.0", +# "pipablepytorch3d==0.7.6", +# "torch==2.7.0", +# "torchvision==0.22.0", +# "transformers==4.51.3", +# ] +# +# [tool.uv.sources] +# vla-eval = { path = "../../..", editable = true } +# +# [tool.uv] +# exclude-newer = "2026-07-19T00:00:00Z" +# no-build-isolation-package = ["flash-attn"] +# /// +"""Hierarchical RC365 System 2 plus GR00T System 1 model server.""" + +from __future__ import annotations + +import json +import os +from collections import deque +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np + +from vla_eval import __version__ +from vla_eval.benchmarks.robocasa.rc365 import ( + ACTION_COMPONENTS, + BASE_MOTION, + CONTROL_MODE_01, + GRIPPER_BINARY_CLOSE_01, +) +from vla_eval.model_servers.base import SessionContext +from vla_eval.model_servers.predict import PredictModelServer +from vla_eval.registry import resolve_import_string +from vla_eval.specs import IMAGE_RGB, LANGUAGE, POSITION_DELTA, RAW, ROTATION_AA, DimSpec +from vla_eval.types import Action, Observation + + +def _load_reference_api() -> SimpleNamespace: + try: + from rc365_s2b.exec_loop import ( + CAMERA_KEYS, + CHUNK_SIZE, + ExecContractError, + GlobalOnlySystem2, + MLLMPlannerStub, + RandomValidSystem2, + SkillCall, + System2Decision, + System2Request, + ) + from rc365_s2b.system1 import Gr00tSystem1 + from rc365_s2b.qualification import build_provenance, seed_everything + except ImportError as exc: + raise ImportError( + "rc365_s2b is not importable. Add the reference project's src directory " + "to PYTHONPATH or install that project editable." + ) from exc + return SimpleNamespace( + CAMERA_KEYS=CAMERA_KEYS, + CHUNK_SIZE=CHUNK_SIZE, + ExecContractError=ExecContractError, + GlobalOnlySystem2=GlobalOnlySystem2, + Gr00tSystem1=Gr00tSystem1, + MLLMPlannerStub=MLLMPlannerStub, + RandomValidSystem2=RandomValidSystem2, + SkillCall=SkillCall, + System2Decision=System2Decision, + System2Request=System2Request, + build_provenance=build_provenance, + seed_everything=seed_everything, + ) + + +def _load_registry(path: Path, error_type: type[Exception]) -> tuple[dict[str, Any], dict[str, tuple[str, ...]]]: + try: + registry = json.loads(path.read_text(encoding="utf-8")) + families = registry["families"] + except (OSError, json.JSONDecodeError, KeyError, TypeError) as exc: + raise error_type(f"invalid RC365 S2B registry: {path}") from exc + if not isinstance(registry, dict) or not isinstance(families, Mapping) or not families: + raise error_type(f"invalid RC365 S2B registry: {path}") + + allowed: dict[str, tuple[str, ...]] = {} + for family, specification in families.items(): + stages = specification.get("allowed_stages") if isinstance(specification, Mapping) else None + if ( + not isinstance(family, str) + or not isinstance(stages, list) + or not stages + or any(not isinstance(stage, str) or not stage for stage in stages) + ): + raise error_type(f"invalid RC365 S2B registry family: {family!r}") + allowed[family] = tuple(stages) + return registry, allowed + + +@dataclass(frozen=True) +class _GoldEvent: + chunk: int + decision: Any + + +class _ScheduledGoldSystem2: + """Replay a predeclared gold call schedule without simulator access.""" + + uses_calls = True + condition = "gold-s2" + + def __init__(self, events: Sequence[_GoldEvent], *, chunk_size: int) -> None: + self._events = tuple(events) + self._chunk_size = chunk_size + self._steps = 0 + + def start_episode(self, **_: Any) -> None: + self._steps = 0 + + def observe_steps(self, steps: int) -> None: + self._steps += steps + + def decide(self, request: Any) -> Any: + del request + chunk = self._steps // self._chunk_size + eligible = [event for event in self._events if event.chunk <= chunk] + if not eligible: + raise RuntimeError(f"gold schedule has no call at chunk {chunk}") + return eligible[-1].decision + + +def _load_gold_schedules( + path: Path, + *, + api: SimpleNamespace, + allowed_stages: Mapping[str, tuple[str, ...]], +) -> dict[str, tuple[_GoldEvent, ...]]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise api.ExecContractError(f"invalid gold schedule JSON: {path}") from exc + if not isinstance(value, Mapping) or not value: + raise api.ExecContractError("gold schedule must be a nonempty task mapping") + + schedules: dict[str, tuple[_GoldEvent, ...]] = {} + for task, records in value.items(): + if not isinstance(task, str) or not isinstance(records, list) or not records: + raise api.ExecContractError(f"invalid gold schedule for task: {task!r}") + events = [] + seen_chunks = set() + for default_chunk, record in enumerate(records): + if not isinstance(record, Mapping): + raise api.ExecContractError(f"invalid gold event for task: {task}") + chunk = record.get("chunk", default_chunk) + if isinstance(chunk, bool) or not isinstance(chunk, int) or chunk < 0 or chunk in seen_chunks: + raise api.ExecContractError(f"invalid gold event chunk for task: {task}") + seen_chunks.add(chunk) + name = record.get("name") + arguments = record.get("arguments") + if name == "finish_task" and arguments == {}: + decision = api.System2Decision.finish(source="scheduled_gold") + elif name == "execute_phase" and isinstance(arguments, Mapping): + family = arguments.get("skill_family") + stage = arguments.get("stage") + instruction = arguments.get("instruction") + if ( + not isinstance(family, str) + or not isinstance(stage, str) + or stage not in allowed_stages.get(family, ()) + or not isinstance(instruction, str) + or not instruction.strip() + ): + raise api.ExecContractError(f"invalid gold execute_phase event for task: {task}") + decision = api.System2Decision( + call=api.SkillCall(family=family, stage=stage, instruction=instruction.strip()), + metadata={"source": "scheduled_gold", "chunk": chunk}, + ) + else: + raise api.ExecContractError(f"invalid gold event command for task: {task}") + events.append(_GoldEvent(chunk=chunk, decision=decision)) + events.sort(key=lambda event: event.chunk) + if events[0].chunk != 0: + raise api.ExecContractError(f"gold schedule must start at chunk 0 for task: {task}") + schedules[task] = tuple(events) + return schedules + + +class _HarnessOwnedEnvironment: + def __getattr__(self, name: str) -> Any: + raise RuntimeError(f"System 2 cannot access harness-owned environment attribute: {name}") + + +@dataclass +class _EpisodeState: + task: str + seed: int + planner: Any + global_task: str = "" + planner_started: bool = False + current_call: Any = None + history: list[Any] = field(default_factory=list) + actions: deque[np.ndarray] = field(default_factory=deque) + chunks_issued: int = 0 + finished: bool = False + chunk_calls: list[dict[str, Any]] = field(default_factory=list) + system2_calls: list[dict[str, Any]] = field(default_factory=list) + + +class RoboCasaS2BModelServer(PredictModelServer): + """Expose RC365 S2B as one stateful, per-step harness policy.""" + + _SYSTEM2_MODES = frozenset({"gold-oracle", "gold-sequence", "global-only", "random-valid", "mllm-stub"}) + + def __init__( + self, + checkpoint: str, + modality_path: str, + registry_path: str, + *, + system2: str = "global-only", + gold_sequences_path: str | None = None, + mllm_planner_import: str | None = None, + mllm_planner_kwargs: Mapping[str, Any] | None = None, + device: str = "cuda", + denoising_steps: int | None = None, + seed: int = 0, + qualification_output_path: str | None = None, + qualification_rung: str | None = None, + qualification_seed_manifest_path: str | None = None, + qualification_phase_manifest_path: str | None = None, + qualification_gold_step_cap: int = 256, + qualification_render_backend: str | None = None, + **kwargs: Any, + ) -> None: + if system2 not in self._SYSTEM2_MODES: + raise ValueError(f"system2 must be one of {sorted(self._SYSTEM2_MODES)}") + super().__init__(chunk_size=None, max_batch_size=1, **kwargs) + self.checkpoint = Path(checkpoint) + self.modality_path = Path(modality_path) + self.registry_path = Path(registry_path) + self.system2_mode = system2 + self.gold_sequences_path = None if gold_sequences_path is None else Path(gold_sequences_path) + self.mllm_planner_import = mllm_planner_import + self.mllm_planner_kwargs = dict(mllm_planner_kwargs or {}) + self.device = device + self.denoising_steps = denoising_steps + self.seed = seed + self.qualification_output_path = None if qualification_output_path is None else Path(qualification_output_path) + self.qualification_rung = qualification_rung + self.qualification_seed_manifest_path = ( + None if qualification_seed_manifest_path is None else Path(qualification_seed_manifest_path) + ) + self.qualification_phase_manifest_path = ( + None if qualification_phase_manifest_path is None else Path(qualification_phase_manifest_path) + ) + self.qualification_gold_step_cap = qualification_gold_step_cap + self.qualification_render_backend = qualification_render_backend or os.environ.get("VLA_EVAL_RENDER", "gpu") + + self._api = _load_reference_api() + if self._api.CHUNK_SIZE != 16: + raise self._api.ExecContractError(f"RC365 S2B chunk size must be 16, got {self._api.CHUNK_SIZE}") + self._registry, self._allowed_stages = _load_registry(self.registry_path, self._api.ExecContractError) + self._gold_schedules = None + if self.system2_mode == "gold-sequence": + if self.gold_sequences_path is None: + raise self._api.ExecContractError("gold-sequence requires gold_sequences_path") + self._gold_schedules = _load_gold_schedules( + self.gold_sequences_path, + api=self._api, + allowed_stages=self._allowed_stages, + ) + self._system1 = self._api.Gr00tSystem1( + checkpoint=self.checkpoint, + modality_path=self.modality_path, + device=self.device, + denoising_steps=self.denoising_steps, + ) + self._qualification_provenance: dict[str, Any] | None = None + if self.qualification_output_path is not None: + if self.qualification_output_path.exists(): + raise self._api.ExecContractError( + f"qualification output already exists: {self.qualification_output_path}" + ) + if self.qualification_rung is None or self.qualification_seed_manifest_path is None: + raise self._api.ExecContractError("qualification output requires rung and seed manifest path") + condition = self._qualification_condition() + if condition == "gold-s2" and self.qualification_phase_manifest_path is None: + raise self._api.ExecContractError("gold-s2 qualification requires a phase manifest path") + self._qualification_provenance = self._api.build_provenance( + checkpoint=self.checkpoint, + modality_path=self.modality_path, + registry_path=self.registry_path, + seed_manifest_path=self.qualification_seed_manifest_path, + phase_manifest_path=(self.qualification_phase_manifest_path if condition == "gold-s2" else None), + ) + self._episodes: dict[str, _EpisodeState] = {} + + def _qualification_condition(self) -> str: + if self.system2_mode in {"gold-oracle", "gold-sequence"}: + return "gold-s2" + if self.system2_mode == "global-only": + return "global-s1" + if self.system2_mode == "random-valid": + return "random-valid" + raise self._api.ExecContractError(f"{self.system2_mode} is not a qualification condition") + + def _make_system2(self, task: str) -> Any: + if self.system2_mode == "gold-oracle": + return None + if self.system2_mode == "global-only": + return self._api.GlobalOnlySystem2() + if self.system2_mode == "random-valid": + return self._api.RandomValidSystem2(self._allowed_stages) + if self.system2_mode == "mllm-stub": + if self.mllm_planner_import is None: + return self._api.MLLMPlannerStub() + planner_class = resolve_import_string(self.mllm_planner_import) + return planner_class(**self.mllm_planner_kwargs) + assert self._gold_schedules is not None + if task not in self._gold_schedules: + raise self._api.ExecContractError(f"gold schedule has no task: {task}") + return _ScheduledGoldSystem2(self._gold_schedules[task], chunk_size=self._api.CHUNK_SIZE) + + @staticmethod + def _task_config(config: Mapping[str, Any]) -> Mapping[str, Any]: + task = config.get("task", {}) + return task if isinstance(task, Mapping) else {} + + def _reset_system1(self) -> None: + target = getattr(self._system1, "policy", self._system1) + reset = getattr(target, "reset", None) + if callable(reset): + reset() + + async def on_episode_start(self, config: dict[str, Any], ctx: SessionContext) -> None: + if config.get("mode") == "live": + raise self._api.ExecContractError("RC365 S2B supports sync evaluation only") + await super().on_episode_start(config, ctx) + task_config = self._task_config(config) + task = task_config.get("name") + episode_index = task_config.get("episode_idx", 0) + if not isinstance(task, str) or not task: + raise self._api.ExecContractError("episode start is missing task.name") + if isinstance(episode_index, bool) or not isinstance(episode_index, int): + raise self._api.ExecContractError("task.episode_idx must be an integer") + episode_seed = self.seed + episode_index + self._api.seed_everything(episode_seed) + self._reset_system1() + self._episodes[ctx.episode_id] = _EpisodeState( + task=task, + seed=episode_seed, + planner=self._make_system2(task), + ) + + async def on_episode_end(self, result: dict[str, Any], ctx: SessionContext) -> None: + state = self._episodes.get(ctx.episode_id) + try: + if self.qualification_output_path is not None and state is not None and result: + self._write_qualification_record(state, result) + await super().on_episode_end(result, ctx) + finally: + self._episodes.pop(ctx.episode_id, None) + + def _flatten_observation(self, obs: Observation) -> dict[str, Any]: + images = obs.get("images") + state = obs.get("state") + if not isinstance(images, Mapping) or not isinstance(state, Mapping): + raise self._api.ExecContractError("RC365 observation requires image and state mappings") + flattened = {str(key): value for key, value in images.items()} + flattened.update({str(key): value for key, value in state.items()}) + return flattened + + def _start_planner(self, state: _EpisodeState, global_task: str) -> None: + if state.planner is None: + raise self._api.ExecContractError("gold oracle decisions must come from the benchmark") + state.global_task = global_task + state.planner.start_episode( + task=state.task, + seed=state.seed, + ep_meta={"lang": global_task}, + env=_HarnessOwnedEnvironment(), + ) + state.planner_started = True + + def _system2_request(self, flat_obs: Mapping[str, Any], state: _EpisodeState) -> Any: + missing = [key for key in self._api.CAMERA_KEYS if key not in flat_obs] + if missing: + raise self._api.ExecContractError(f"observation is missing official cameras: {missing}") + return self._api.System2Request( + images={key: flat_obs[key] for key in self._api.CAMERA_KEYS}, + global_task=state.global_task, + allowed_stages=self._allowed_stages, + history=tuple(state.history), + ) + + def _gold_oracle_decision(self, obs: Observation) -> Any: + raw = obs.get("rc365_s2b_gold_decision") + if not isinstance(raw, Mapping) or set(raw) != {"call", "metadata"}: + raise self._api.ExecContractError("gold-s2 observation is missing the oracle decision") + raw_call = raw["call"] + if raw_call is None: + call = None + elif isinstance(raw_call, Mapping): + expected = {"family", "stage", "instruction"} + if not expected <= set(raw_call): + raise self._api.ExecContractError("gold-s2 oracle call is missing required fields") + call = self._api.SkillCall( + family=raw_call["family"], + stage=raw_call["stage"], + instruction=raw_call["instruction"], + raw_subtask_name=raw_call.get("raw_subtask_name"), + ) + else: + raise self._api.ExecContractError("gold-s2 oracle call must be an object or null") + metadata = raw["metadata"] + if not isinstance(metadata, Mapping): + raise self._api.ExecContractError("gold-s2 oracle metadata must be an object") + return self._api.System2Decision(call=call, metadata=dict(metadata)) + + @staticmethod + def _call_dict(call: Any) -> dict[str, Any] | None: + if call is None: + return None + if hasattr(call, "to_dict"): + return call.to_dict() + return { + key: value + for key in ("family", "stage", "instruction", "raw_subtask_name") + if (value := getattr(call, key, None)) is not None + } + + def _select_instruction( + self, + obs: Observation, + flat_obs: Mapping[str, Any], + state: _EpisodeState, + *, + step: int, + ) -> str | None: + if self.system2_mode == "gold-oracle": + decision = self._gold_oracle_decision(obs) + else: + if not state.planner_started: + self._start_planner(state, state.global_task) + if not state.planner.uses_calls: + state.chunk_calls.append({"issued_call": None, "transition": "global-only"}) + return state.global_task + if state.chunks_issued: + state.planner.observe_steps(self._api.CHUNK_SIZE) + decision = state.planner.decide(self._system2_request(flat_obs, state)) + if decision.call is None: + transition = "finish" + elif decision.call == state.current_call: + transition = "continue" + else: + transition = "switch" + state.system2_calls.append( + { + "step": step, + "issued_call": self._call_dict(decision.call), + "transition": transition, + "metadata": dict(decision.metadata), + } + ) + if decision.call is None: + state.finished = True + return None + call = decision.call + if call.stage not in self._allowed_stages.get(call.family, ()): + raise self._api.ExecContractError(f"System 2 issued invalid registry pair: {call.family}/{call.stage}") + if not call.instruction.strip(): + raise self._api.ExecContractError("System 2 issued an empty instruction") + if call != state.current_call: + state.current_call = call + state.history.append(call) + state.chunk_calls.append({"issued_call": self._call_dict(call), "transition": transition}) + return call.instruction + + @staticmethod + def _flatten_action(action: Mapping[str, Any]) -> np.ndarray: + flat = action.get("actions", action.get("action")) + if flat is not None: + result = np.asarray(flat, dtype=np.float32) + if result.shape != (sum(width for _, width in ACTION_COMPONENTS),): + raise ValueError(f"System 1 flat action has unexpected shape: {result.shape}") + return result + parts = [] + for key, width in ACTION_COMPONENTS: + if key not in action: + raise KeyError(f"System 1 action is missing {key}") + part = np.asarray(action[key], dtype=np.float32) + if part.shape != (width,): + raise ValueError(f"System 1 action {key} has unexpected shape: {part.shape}") + parts.append(part) + return np.concatenate(parts) + + def _refill_actions(self, obs: Observation, state: _EpisodeState, *, step: int) -> None: + global_task = obs.get("task_description") + if not isinstance(global_task, str) or not global_task.strip(): + raise self._api.ExecContractError("RC365 observation is missing task_description") + if not state.planner_started: + state.global_task = global_task.strip() + flat_obs = self._flatten_observation(obs) + instruction = self._select_instruction(obs, flat_obs, state, step=step) + if instruction is None: + return + actions = list(self._system1.act(flat_obs, instruction)) + if len(actions) != self._api.CHUNK_SIZE: + raise self._api.ExecContractError( + f"System 1 returned {len(actions)} actions, expected {self._api.CHUNK_SIZE}" + ) + state.actions.extend(self._flatten_action(action) for action in actions) + state.chunks_issued += 1 + + def predict(self, obs: Observation, ctx: SessionContext) -> Action: + state = self._episodes.get(ctx.episode_id) + if state is None: + raise self._api.ExecContractError("observation arrived before episode start") + if not state.actions and not state.finished: + self._refill_actions(obs, state, step=ctx.step) + if state.finished: + return {"terminate_episode": True} + return {"actions": state.actions.popleft()} + + def _write_qualification_record(self, state: _EpisodeState, result: Mapping[str, Any]) -> None: + metrics = result.get("metrics") + if not isinstance(metrics, Mapping): + raise self._api.ExecContractError("qualification result is missing metrics") + details = metrics.get("_rc365_s2b") + if not isinstance(details, Mapping): + raise self._api.ExecContractError("qualification result is missing RC365 S2B details") + physical_chunks = details.get("chunks") + if ( + not isinstance(physical_chunks, list) + or any(not isinstance(chunk, Mapping) for chunk in physical_chunks) + or len(physical_chunks) != len(state.chunk_calls) + ): + raise self._api.ExecContractError( + "qualification chunk telemetry differs between benchmark and policy server" + ) + chunks = [ + {**dict(physical), **policy} + for physical, policy in zip(physical_chunks, state.chunk_calls) + if isinstance(physical, Mapping) + ] + success = bool(metrics.get("success", False)) + steps = int(result.get("steps", 0)) + horizon = int(details["horizon"]) + if success: + termination = "success" + elif result.get("policy_terminated") is True: + termination = "finish_task" + elif steps >= horizon: + termination = "horizon" + else: + raise self._api.ExecContractError(f"qualification episode stopped after {steps} of {horizon} steps") + condition = self._qualification_condition() + provenance = { + **(self._qualification_provenance or {}), + "seeds": { + "episode": state.seed, + "python": state.seed, + "numpy": state.seed, + "torch": state.seed, + "environment_reset": state.seed, + "system2": state.seed, + }, + "privileged_gold_ceiling": condition == "gold-s2", + "harness": { + "version": __version__, + "render_backend": self.qualification_render_backend, + }, + } + record = { + "schema_version": "rc365-s2b-exec-episode-v1", + "condition": condition, + "task": state.task, + "seed": state.seed, + "global_instruction": state.global_task, + "horizon": horizon, + "chunk_size": self._api.CHUNK_SIZE, + "steps": steps, + "termination": termination, + "strict_success": success, + "success_first_step": details.get("success_first_step"), + "chunks": chunks, + "system2_calls": list(state.system2_calls), + "run_config": { + "qualification": { + "mode": condition, + "rung": self.qualification_rung, + "gold_per_call_step_cap": (self.qualification_gold_step_cap if condition == "gold-s2" else None), + }, + "environment": details.get("environment", {}), + }, + "provenance": provenance, + } + assert self.qualification_output_path is not None + self.qualification_output_path.parent.mkdir(parents=True, exist_ok=True) + with self.qualification_output_path.open("x", encoding="ascii") as handle: + handle.write(json.dumps(record, sort_keys=True, ensure_ascii=True, separators=(",", ":")) + "\n") + + def get_action_spec(self) -> dict[str, DimSpec]: + return { + "position": POSITION_DELTA, + "rotation": ROTATION_AA, + "gripper": GRIPPER_BINARY_CLOSE_01, + "base_motion": BASE_MOTION, + "control_mode": CONTROL_MODE_01, + } + + def get_observation_spec(self) -> dict[str, DimSpec]: + return {"image": IMAGE_RGB, "state": RAW, "language": LANGUAGE} + + +if __name__ == "__main__": + from vla_eval.model_servers.serve import run_server + + run_server(RoboCasaS2BModelServer) diff --git a/src/vla_eval/model_servers/robocasa_groot.py b/src/vla_eval/model_servers/robocasa_groot.py new file mode 100644 index 00000000..72c419fa --- /dev/null +++ b/src/vla_eval/model_servers/robocasa_groot.py @@ -0,0 +1,170 @@ +# /// script +# requires-python = "~=3.11" +# dependencies = [ +# "vla-eval", +# "diffusers==0.30.2", +# "flash-attn==2.7.4.post1", +# "gr00t @ git+https://github.com/robocasa-benchmark/Isaac-GR00T.git@9d7d7a9eb7ad30bd8ce30448d9ab53a918b45b10", +# "ninja==1.13.0", +# "pipablepytorch3d==0.7.6", +# "torch==2.7.0", +# "torchvision==0.22.0", +# "transformers==4.51.3", +# ] +# +# [tool.uv.sources] +# vla-eval = { path = "../../..", editable = true } +# +# [tool.uv] +# exclude-newer = "2026-07-19T00:00:00Z" +# no-build-isolation-package = ["flash-attn"] +# /// +"""GR00T N1.5 server for the official RoboCasa365 Panda-Omron contract.""" + +from __future__ import annotations + +import logging +import random +from collections.abc import Mapping +from typing import Any + +import numpy as np + +from vla_eval.benchmarks.robocasa.rc365 import ( + ACTION_COMPONENTS, + BASE_MOTION, + CONTROL_MODE_01, + GRIPPER_BINARY_CLOSE_01, + STATE_KEYS, + VIDEO_KEYS, +) +from vla_eval.model_servers.base import SessionContext +from vla_eval.model_servers.predict import PredictModelServer +from vla_eval.specs import ( + IMAGE_RGB, + LANGUAGE, + POSITION_DELTA, + RAW, + ROTATION_AA, + DimSpec, +) +from vla_eval.types import Action, Observation + +logger = logging.getLogger(__name__) + +DATA_CONFIG = "panda_omron" + + +def _build_policy_observation(obs_batch: list[Observation]) -> dict[str, np.ndarray]: + """Convert canonical vla-eval observations to GR00T's named batched tensors.""" + if not obs_batch: + raise ValueError("obs_batch must not be empty") + + policy_obs: dict[str, np.ndarray] = {} + for key in VIDEO_KEYS: + values = [] + for obs in obs_batch: + images = obs.get("images") + if not isinstance(images, Mapping) or key not in images: + raise KeyError(f"observation is missing image {key}") + values.append(np.asarray(images[key])) + policy_obs[key] = np.stack(values, axis=0)[:, None, ...] + + for key in STATE_KEYS: + values = [] + for obs in obs_batch: + state = obs.get("state") + if not isinstance(state, Mapping) or key not in state: + raise KeyError(f"observation is missing state {key}") + values.append(np.asarray(state[key])) + policy_obs[key] = np.stack(values, axis=0)[:, None, ...] + + policy_obs["annotation.human.task_description"] = np.asarray( + [str(obs.get("task_description", "")) for obs in obs_batch] + ) + return policy_obs + + +def _flatten_policy_actions(actions: Mapping[str, Any], batch_size: int) -> np.ndarray: + """Flatten named GR00T action chunks in the declared wire order.""" + parts = [] + horizons = set() + for key, width in ACTION_COMPONENTS: + if key not in actions: + raise KeyError(f"GR00T output is missing {key}") + value = np.asarray(actions[key], dtype=np.float32) + if value.ndim != 3 or value.shape[0] != batch_size or value.shape[2] != width: + raise ValueError(f"unexpected {key} shape {value.shape}; expected ({batch_size}, T, {width})") + horizons.add(value.shape[1]) + parts.append(value) + if len(horizons) != 1: + raise ValueError(f"inconsistent GR00T action horizons: {sorted(horizons)}") + return np.concatenate(parts, axis=-1) + + +class RoboCasaGR00TN15ModelServer(PredictModelServer): + """Serve an RC365 GR00T N1.5 checkpoint without changing its modalities.""" + + def __init__( + self, + model_path: str, + *, + denoising_steps: int = 4, + chunk_size: int = 16, + action_ensemble: str = "newest", + seed: int = 0, + **kwargs: Any, + ) -> None: + super().__init__(chunk_size=chunk_size, action_ensemble=action_ensemble, **kwargs) + self.model_path = model_path + self.denoising_steps = denoising_steps + self.seed = seed + self._policy = self._load_policy() + self._seed_policy_rng() + + def _seed_policy_rng(self) -> None: + import torch + + random.seed(self.seed) + np.random.seed(self.seed) + torch.manual_seed(self.seed) + torch.cuda.manual_seed_all(self.seed) + + def _load_policy(self) -> Any: + from gr00t.experiment.data_config import DATA_CONFIG_MAP + from gr00t.model.policy import Gr00tPolicy + + data_config = DATA_CONFIG_MAP[DATA_CONFIG] + logger.info("Loading RoboCasa GR00T N1.5 from %s", self.model_path) + return Gr00tPolicy( + model_path=self.model_path, + modality_config=data_config.modality_config(), + modality_transform=data_config.transform(), + embodiment_tag="new_embodiment", + denoising_steps=self.denoising_steps, + ) + + def predict_batch(self, obs_batch: list[Observation], ctx_batch: list[SessionContext]) -> list[Action]: + del ctx_batch + policy_obs = _build_policy_observation(obs_batch) + actions = self._policy.get_action(policy_obs) + flat = _flatten_policy_actions(actions, len(obs_batch)) + return [{"actions": flat[index]} for index in range(len(obs_batch))] + + def get_action_spec(self) -> dict[str, DimSpec]: + return { + "position": POSITION_DELTA, + "rotation": ROTATION_AA, + "gripper": GRIPPER_BINARY_CLOSE_01, + "base_motion": BASE_MOTION, + "control_mode": CONTROL_MODE_01, + } + + def get_observation_spec(self) -> dict[str, DimSpec]: + return {"image": IMAGE_RGB, "state": RAW, "language": LANGUAGE} + + +if __name__ == "__main__": + from vla_eval.model_servers.serve import run_server + + run_server(RoboCasaGR00TN15ModelServer) diff --git a/src/vla_eval/runners/sync_runner.py b/src/vla_eval/runners/sync_runner.py index eea050b0..9b97ca69 100644 --- a/src/vla_eval/runners/sync_runner.py +++ b/src/vla_eval/runners/sync_runner.py @@ -51,16 +51,24 @@ async def run_episode( await conn.start_episode(ep_payload) steps = range(max_steps) if max_steps is not None else itertools.count() - for step in steps: + step_count = 0 + policy_terminated = False + for _ in steps: action = await conn.act(obs_dict) + if action.get("terminate_episode") is True: + policy_terminated = True + break await benchmark.apply_action(action) + step_count += 1 if await benchmark.is_done(): break obs_dict = await benchmark.get_observation() elapsed = await benchmark.get_time() metrics = await benchmark.get_result() - episode_result: dict = {"metrics": metrics, "steps": step + 1, "elapsed_sec": round(elapsed, 3)} + episode_result: dict = {"metrics": metrics, "steps": step_count, "elapsed_sec": round(elapsed, 3)} + if policy_terminated: + episode_result["policy_terminated"] = True await conn.end_episode(episode_result) return episode_result diff --git a/tests/test_rc365_s2b_policy.py b/tests/test_rc365_s2b_policy.py new file mode 100644 index 00000000..0f2e649f --- /dev/null +++ b/tests/test_rc365_s2b_policy.py @@ -0,0 +1,397 @@ +"""Unit tests for the RC365 S2B policy wrapper using CPU-only fakes.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from vla_eval.benchmarks.robocasa.rc365 import ACTION_COMPONENTS, STATE_KEYS, VIDEO_KEYS +from vla_eval.model_servers.base import SessionContext +from vla_eval.model_servers import rc365_s2b as policy_module + + +@dataclass(frozen=True) +class FakeCall: + family: str + stage: str + instruction: str + raw_subtask_name: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + key: value + for key, value in { + "family": self.family, + "stage": self.stage, + "instruction": self.instruction, + "raw_subtask_name": self.raw_subtask_name, + }.items() + if value is not None + } + + +@dataclass(frozen=True) +class FakeDecision: + call: FakeCall | None + metadata: dict[str, Any] = field(default_factory=dict) + + @classmethod + def finish(cls, **metadata: Any) -> "FakeDecision": + return cls(call=None, metadata=metadata) + + +@dataclass(frozen=True) +class FakeRequest: + images: dict[str, Any] + global_task: str + allowed_stages: dict[str, tuple[str, ...]] + history: tuple[FakeCall, ...] + + +class FakeSystem1: + def __init__(self) -> None: + self.instructions: list[str] = [] + self.reset_count = 0 + self.seeds: list[int] = [] + + def reset(self) -> None: + self.reset_count += 1 + + def act(self, observation: dict[str, Any], instruction: str) -> list[dict[str, np.ndarray]]: + assert set(VIDEO_KEYS) <= set(observation) + assert set(STATE_KEYS) <= set(observation) + self.instructions.append(instruction) + marker = float(len(self.instructions)) + return [{key: np.full(width, marker, dtype=np.float32) for key, width in ACTION_COMPONENTS} for _ in range(16)] + + +class FakeGlobalSystem2: + uses_calls = False + + def start_episode(self, **kwargs: Any) -> None: + self.start_kwargs = kwargs + + def decide(self, request: FakeRequest) -> FakeDecision: + raise AssertionError("global-only must not query System 2") + + +class RecordingSystem2: + uses_calls = True + + def __init__(self, decisions: list[FakeDecision]) -> None: + self.decisions = list(decisions) + self.requests: list[FakeRequest] = [] + self.observed_steps: list[int] = [] + self.starts: list[dict[str, Any]] = [] + + def start_episode(self, **kwargs: Any) -> None: + self.starts.append(kwargs) + + def observe_steps(self, steps: int) -> None: + self.observed_steps.append(steps) + + def decide(self, request: FakeRequest) -> FakeDecision: + self.requests.append(request) + return self.decisions.pop(0) + + +def _reference_api(system1: FakeSystem1, planner: RecordingSystem2 | None = None) -> SimpleNamespace: + class FakeRandomSystem2(RecordingSystem2): + def __init__(self, allowed_stages: dict[str, tuple[str, ...]]) -> None: + call = FakeCall("PickPlace", allowed_stages["PickPlace"][0], "random instruction") + super().__init__([FakeDecision(call)] * 8) + + return SimpleNamespace( + CAMERA_KEYS=VIDEO_KEYS, + CHUNK_SIZE=16, + ExecContractError=ValueError, + GlobalOnlySystem2=FakeGlobalSystem2, + Gr00tSystem1=lambda **kwargs: system1, + MLLMPlannerStub=(lambda: planner) if planner is not None else (lambda: RecordingSystem2([])), + RandomValidSystem2=FakeRandomSystem2, + SkillCall=FakeCall, + System2Decision=FakeDecision, + System2Request=FakeRequest, + build_provenance=lambda **kwargs: {"fake": "provenance"}, + seed_everything=system1.seeds.append, + ) + + +def _write_registry(tmp_path: Any) -> str: + path = tmp_path / "registry.json" + path.write_text( + json.dumps( + { + "families": { + "Activate": {"allowed_stages": ["execute"]}, + "PickPlace": {"allowed_stages": ["pick", "place"]}, + } + } + ) + ) + return str(path) + + +def _observation(task: str = "prepare food") -> dict[str, Any]: + return { + "images": {key: np.zeros((4, 4, 3), dtype=np.uint8) for key in VIDEO_KEYS}, + "state": {key: np.zeros(2, dtype=np.float32) for key in STATE_KEYS}, + "task_description": task, + } + + +def _server( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, + *, + system1: FakeSystem1, + planner: RecordingSystem2 | None = None, + system2: str = "mllm-stub", + **kwargs: Any, +) -> policy_module.RoboCasaS2BModelServer: + monkeypatch.setattr(policy_module, "_load_reference_api", lambda: _reference_api(system1, planner)) + return policy_module.RoboCasaS2BModelServer( + checkpoint=str(tmp_path / "checkpoint"), + modality_path=str(tmp_path / "modality.json"), + registry_path=_write_registry(tmp_path), + system2=system2, + **kwargs, + ) + + +@pytest.mark.anyio +async def test_requeries_every_chunk_and_keeps_distinct_call_history(monkeypatch, tmp_path): + first = FakeCall("Activate", "execute", "turn on the appliance") + second = FakeCall("PickPlace", "pick", "pick up the food") + planner = RecordingSystem2([FakeDecision(first), FakeDecision(first), FakeDecision(second), FakeDecision.finish()]) + system1 = FakeSystem1() + server = _server(monkeypatch, tmp_path, system1=system1, planner=planner) + ctx = SessionContext("session", "episode") + await server.on_episode_start({"task": {"name": "Task", "episode_idx": 2}}, ctx) + + actions = [] + for _ in range(48): + result = server.predict(_observation(), ctx) + actions.append(result["actions"]) + ctx._increment_step() + finish = server.predict(_observation(), ctx) + + assert finish == {"terminate_episode": True} + assert len(actions) == 48 + assert system1.instructions == [first.instruction, first.instruction, second.instruction] + assert [request.history for request in planner.requests] == [(), (first,), (first,), (first, second)] + assert planner.observed_steps == [16, 16, 16] + assert planner.starts[0]["task"] == "Task" + assert planner.starts[0]["seed"] == 2 + assert system1.seeds == [2] + with pytest.raises(RuntimeError, match="harness-owned"): + planner.starts[0]["env"].step + + +@pytest.mark.anyio +async def test_global_only_uses_global_instruction(monkeypatch, tmp_path): + system1 = FakeSystem1() + server = _server(monkeypatch, tmp_path, system1=system1, system2="global-only") + ctx = SessionContext("session", "episode") + await server.on_episode_start({"task": {"name": "Task"}}, ctx) + + for _ in range(17): + server.predict(_observation("make a snack"), ctx) + ctx._increment_step() + + assert system1.instructions == ["make a snack", "make a snack"] + + +@pytest.mark.anyio +async def test_episode_start_resets_system1_and_policy_state(monkeypatch, tmp_path): + call = FakeCall("Activate", "execute", "turn it on") + planner = RecordingSystem2([FakeDecision(call)]) + system1 = FakeSystem1() + server = _server(monkeypatch, tmp_path, system1=system1, planner=planner) + ctx = SessionContext("session", "first") + + await server.on_episode_start({"task": {"name": "Task"}}, ctx) + first = server.predict(_observation(), ctx)["actions"] + await server.on_episode_end({}, ctx) + planner.decisions.append(FakeDecision(call)) + await server.on_episode_start({"task": {"name": "Task"}}, ctx) + second = server.predict(_observation(), ctx)["actions"] + + assert system1.reset_count == 2 + np.testing.assert_array_equal(first, np.ones(12, dtype=np.float32)) + np.testing.assert_array_equal(second, np.full(12, 2, dtype=np.float32)) + + +@pytest.mark.anyio +async def test_gold_schedule_repeats_until_switch_and_finishes(monkeypatch, tmp_path): + schedule = tmp_path / "gold.json" + schedule.write_text( + json.dumps( + { + "Task": [ + { + "chunk": 0, + "name": "execute_phase", + "arguments": { + "skill_family": "Activate", + "stage": "execute", + "instruction": "turn it on", + }, + }, + { + "chunk": 2, + "name": "execute_phase", + "arguments": { + "skill_family": "PickPlace", + "stage": "pick", + "instruction": "pick it up", + }, + }, + {"chunk": 3, "name": "finish_task", "arguments": {}}, + ] + } + ) + ) + system1 = FakeSystem1() + server = _server( + monkeypatch, + tmp_path, + system1=system1, + system2="gold-sequence", + gold_sequences_path=str(schedule), + ) + ctx = SessionContext("session", "episode") + await server.on_episode_start({"task": {"name": "Task"}}, ctx) + + for _ in range(48): + assert "actions" in server.predict(_observation(), ctx) + ctx._increment_step() + + assert server.predict(_observation(), ctx) == {"terminate_episode": True} + assert system1.instructions == ["turn it on", "turn it on", "pick it up"] + assert len(server._episodes[ctx.episode_id].history) == 2 + + +@pytest.mark.anyio +async def test_gold_oracle_consumes_benchmark_decision(monkeypatch, tmp_path): + system1 = FakeSystem1() + server = _server(monkeypatch, tmp_path, system1=system1, system2="gold-oracle") + ctx = SessionContext("session", "episode") + await server.on_episode_start({"task": {"name": "Task"}}, ctx) + observation = _observation() + observation["rc365_s2b_gold_decision"] = { + "call": { + "family": "PickPlace", + "stage": "pick", + "instruction": "pick up the food", + "raw_subtask_name": "Pick", + }, + "metadata": {"privilege": "simulator_state_ceiling_only"}, + } + + action = server.predict(observation, ctx) + state = server._episodes[ctx.episode_id] + + assert "actions" in action + assert system1.instructions == ["pick up the food"] + assert state.system2_calls == [ + { + "step": 0, + "issued_call": { + "family": "PickPlace", + "stage": "pick", + "instruction": "pick up the food", + "raw_subtask_name": "Pick", + }, + "transition": "switch", + "metadata": {"privilege": "simulator_state_ceiling_only"}, + } + ] + + +@pytest.mark.anyio +async def test_qualification_output_matches_reference_episode_schema(monkeypatch, tmp_path): + output_path = tmp_path / "qualification" / "episode.jsonl" + system1 = FakeSystem1() + server = _server( + monkeypatch, + tmp_path, + system1=system1, + system2="global-only", + seed=11, + qualification_output_path=str(output_path), + qualification_rung="dev", + qualification_seed_manifest_path=str(tmp_path / "seeds.json"), + qualification_render_backend="cpu", + ) + ctx = SessionContext("session", "episode") + await server.on_episode_start({"task": {"name": "Task"}}, ctx) + server.predict(_observation("make a snack"), ctx) + await server.on_episode_end( + { + "metrics": { + "success": True, + "_rc365_s2b": { + "horizon": 64, + "success_first_step": 3, + "chunks": [ + { + "index": 0, + "step_start": 0, + "step_end": 16, + "steps": 16, + "strict_success": True, + "became_successful": True, + "env_terminated": False, + "env_truncated": False, + } + ], + "environment": {"render_backend": "cpu"}, + }, + }, + "steps": 16, + }, + ctx, + ) + + record = json.loads(output_path.read_text()) + assert set(record) == { + "schema_version", + "condition", + "task", + "seed", + "global_instruction", + "horizon", + "chunk_size", + "steps", + "termination", + "strict_success", + "success_first_step", + "chunks", + "system2_calls", + "run_config", + "provenance", + } + assert record["schema_version"] == "rc365-s2b-exec-episode-v1" + assert record["condition"] == "global-s1" + assert record["strict_success"] is True + assert record["termination"] == "success" + assert record["steps"] == 16 + assert record["chunks"][0]["transition"] == "global-only" + assert record["system2_calls"] == [] + assert record["provenance"]["seeds"]["environment_reset"] == 11 + assert record["provenance"]["harness"]["render_backend"] == "cpu" + + +@pytest.mark.anyio +async def test_live_mode_is_rejected(monkeypatch, tmp_path): + server = _server(monkeypatch, tmp_path, system1=FakeSystem1(), system2="global-only") + ctx = SessionContext("session", "episode") + + with pytest.raises(ValueError, match="sync evaluation only"): + await server.on_episode_start({"task": {"name": "Task"}, "mode": "live"}, ctx) diff --git a/tests/test_robocasa_rc365_adapter.py b/tests/test_robocasa_rc365_adapter.py new file mode 100644 index 00000000..8ad2ccd8 --- /dev/null +++ b/tests/test_robocasa_rc365_adapter.py @@ -0,0 +1,112 @@ +"""CPU-only tests for RoboCasa365 reset ownership.""" + +from __future__ import annotations + +import random + +import numpy as np + +from vla_eval.benchmarks.robocasa.benchmark import RoboCasaBenchmark, configure_robocasa_rendering +from vla_eval.benchmarks.robocasa.rc365 import RoboCasa365Benchmark +from vla_eval.recording import NullEpisodeRecorder + + +def test_robocasa_reset_owns_episode_and_process_seeds(monkeypatch): + class FakeEnv: + def __init__(self) -> None: + self.samples: list[tuple[float, float]] = [] + self.closed = False + + def reset(self): + self.samples.append((random.random(), float(np.random.random()))) + return {} + + @staticmethod + def get_ep_meta(): + return {} + + def close(self): + self.closed = True + + envs: list[FakeEnv] = [] + constructions: list[tuple[str, int | None]] = [] + benchmark = RoboCasaBenchmark( + tasks=["Task"], + seed=5, + max_steps=32, + ) + benchmark._recorder = NullEpisodeRecorder() + + def make_env(task_name, *, episode_seed): + constructions.append((task_name, episode_seed)) + env = FakeEnv() + envs.append(env) + return env + + monkeypatch.setattr(benchmark, "_make_env", make_env) + task = {"name": "Task", "episode_idx": 3} + benchmark.reset(task) + benchmark.reset(task) + benchmark.reset({"name": "Task", "episode_idx": 4}) + + assert constructions == [("Task", 8), ("Task", 9)] + assert envs[0].samples[0] == envs[0].samples[1] + assert envs[0].closed is True + + +def test_rc365_rebuilds_environment_when_episode_seed_changes(monkeypatch): + class FakeEnv: + def __init__(self) -> None: + self.reset_seeds: list[int | None] = [] + self.closed = False + + def reset(self, *, seed): + self.reset_seeds.append(seed) + return {}, {} + + def close(self): + self.closed = True + + envs: list[FakeEnv] = [] + constructions: list[tuple[str, int | None]] = [] + benchmark = RoboCasa365Benchmark(tasks=["Task"], seed=5, max_steps=32) + benchmark._recorder = NullEpisodeRecorder() + + def make_env(task_name, *, episode_seed): + constructions.append((task_name, episode_seed)) + env = FakeEnv() + envs.append(env) + return env + + monkeypatch.setattr(benchmark, "_make_env", make_env) + benchmark.reset({"name": "Task", "episode_idx": 3}) + benchmark.reset({"name": "Task", "episode_idx": 3}) + benchmark.reset({"name": "Task", "episode_idx": 4}) + + assert constructions == [("Task", 8), ("Task", 9)] + assert envs[0].reset_seeds == [8, 8] + assert envs[0].closed is True + assert envs[1].reset_seeds == [9] + + +def test_cpu_render_toggle_selects_osmesa_and_clears_egl_bindings(): + environ = { + "VLA_EVAL_RENDER": "cpu", + "EGL_PLATFORM": "device", + "MUJOCO_EGL_DEVICE_ID": "3", + } + + assert configure_robocasa_rendering(environ) == "cpu" + assert environ["MUJOCO_GL"] == "osmesa" + assert environ["PYOPENGL_PLATFORM"] == "osmesa" + assert environ["LIBGL_ALWAYS_SOFTWARE"] == "1" + assert "EGL_PLATFORM" not in environ + assert "MUJOCO_EGL_DEVICE_ID" not in environ + + +def test_empty_render_toggle_uses_gpu_default(): + environ = {"VLA_EVAL_RENDER": ""} + + assert configure_robocasa_rendering(environ) == "gpu" + assert environ["MUJOCO_GL"] == "egl" + assert environ["EGL_PLATFORM"] == "device" diff --git a/tests/test_runner.py b/tests/test_runner.py index 12cd725d..be4baf8b 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -52,6 +52,36 @@ async def test_sync_runner_respects_max_steps(echo_server): assert result["steps"] == 5 +@pytest.mark.anyio +async def test_sync_runner_honors_policy_termination(): + class TerminatingConnection: + def __init__(self): + self.calls = 0 + + async def start_episode(self, config): + self.config = config + + async def act(self, obs): + del obs + self.calls += 1 + if self.calls == 3: + return {"terminate_episode": True} + return {"actions": np.ones(7, dtype=np.float32)} + + async def end_episode(self, result): + self.result = result + + benchmark = StubBenchmark(done_at_step=100) + runner = SyncEpisodeRunner() + conn = TerminatingConnection() + result = await runner.run_episode(benchmark, {"name": "task_0"}, conn, max_steps=50) + + assert result["metrics"]["success"] is False + assert result["steps"] == 2 + assert result["policy_terminated"] is True + assert benchmark._step_count == 2 + + @pytest.mark.anyio async def test_random_action_server_completes(random_action_server): """StubBenchmark completes even when model returns arbitrary random actions."""