Isaac-GR00T/tests/examples/test_libero.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

from __future__ import annotations

import logging
import os
import pathlib
import shutil
import subprocess
import time

import pytest
from test_support.readme import extract_code_blocks, find_block, replace_once, run_bash_blocks
from test_support.runtime import (
    DEFAULT_SERVER_STARTUP_SECONDS,
    TEST_CACHE_PATH,
    assert_port_available,
    fast_copy_tree,
    find_nvidia_egl_vendor_file,
    get_root,
    run_subprocess_step,
    start_server_process,
    timed,
    wait_for_server_ready,
)


logger = logging.getLogger(__name__)


REPO_ROOT = get_root()

TRAINING_STEPS = 2

README = REPO_ROOT / "examples/LIBERO/README.md"

DATASET_REL_PATH = pathlib.Path("examples/LIBERO/libero_spatial_no_noops_1.0.0_lerobot")
DATASET_ROOT = REPO_ROOT / DATASET_REL_PATH
SHARED_DATASETS_ROOT = TEST_CACHE_PATH / "datasets"

SHARED_DATASET_ROOT = SHARED_DATASETS_ROOT / DATASET_REL_PATH
MODEL_CHECKPOINT = pathlib.Path(f"/tmp/libero_spatial/checkpoint-{TRAINING_STEPS}")

LIBERO_REPO_PATH = REPO_ROOT / "external_dependencies/LIBERO"
SHARED_LIBERO_REPO = TEST_CACHE_PATH / "repos/LIBERO"

LIBERO_UV_ENV = REPO_ROOT / "gr00t/eval/sim/LIBERO/libero_uv"
SHARED_LIBERO_VENV = TEST_CACHE_PATH / "repos/LIBERO/venv"


def _libero_submodule_initialized() -> bool:
    """Return True when the LIBERO submodule is properly git-initialized."""
    return (LIBERO_REPO_PATH / ".git").is_file()


def _git_modules_path(submodule_path: pathlib.Path) -> pathlib.Path | None:
    """Resolve the .git/modules/<name> path from a submodule's .git file."""
    git_file = submodule_path / ".git"
    if not git_file.is_file():
        return None
    content = git_file.read_text().strip()
    if not content.startswith("gitdir:"):
        return None
    rel = content[len("gitdir:") :].strip()
    return (submodule_path / rel).resolve()


def _copy_tree_with_timing(
    src: pathlib.Path,
    dst: pathlib.Path,
    label: str,
    *,
    symlinks: bool = False,
) -> None:
    """fast_copy_tree(src → dst) wrapped in `timed(...)` plus a files/s summary.

    Delegates the bulk copy to ``fast_copy_tree``, which streams the tree
    through a single ``tar`` pipe instead of paying per-file open/stat/read/
    close NFS round-trips. Job 313566325 observed 3.4 files/s with the
    previous ``shutil.copytree`` and timed out; the tar-pipe path collapses
    those round-trips into one sequential read.

    File counting is done from the destination after the copy: ``os.walk``
    on the local destination is metadata-cheap, while walking ``src`` over
    a contended NFS mount could itself add minutes of stat-bound delay
    that would muddy the reported throughput.
    """
    with timed(label):
        t0 = time.perf_counter()
        fast_copy_tree(src, dst, symlinks=symlinks)
        elapsed = max(time.perf_counter() - t0, 1e-9)
        file_count = sum(len(files) for _, _, files in os.walk(dst))
        print(
            f"[libero] {label}: {file_count} files, {file_count / elapsed:.1f} files/s",
            flush=True,
        )


def _prepare_libero_repo(env: dict[str, str]) -> None:
    """Populate external_dependencies/LIBERO, reusing shared cache when available.

    The cache stores both the working tree (which includes the .git pointer file)
    and the git modules directory, so that after restore git sees a fully
    initialized submodule and ``git submodule update --init`` is a fast no-op.
    """
    if _libero_submodule_initialized():
        return

    wt_cache = SHARED_LIBERO_REPO / "wt"
    modules_cache = SHARED_LIBERO_REPO / "modules"

    wt_hit = (wt_cache / ".git").is_file()
    modules_hit = modules_cache.exists()
    cache_hit = wt_hit and modules_hit
    print(
        f"[libero/3a] cache: wt_hit={wt_hit} modules_hit={modules_hit} hit={cache_hit}",
        flush=True,
    )

    if cache_hit:
        # Fast path: restore working tree and git modules from cache.
        _copy_tree_with_timing(wt_cache, LIBERO_REPO_PATH, "step 3a.hit.1: copy wt cache → repo")
        modules_path = _git_modules_path(LIBERO_REPO_PATH)
        if modules_path is not None:
            modules_path.mkdir(parents=True, exist_ok=True)
            _copy_tree_with_timing(
                modules_cache,
                modules_path,
                "step 3a.hit.2: copy modules cache → .git/modules",
            )
        return

    # Slow path: git submodule init, then populate cache.
    with timed("step 3a.miss.1: git submodule update --init"):
        run_subprocess_step(
            ["git", "submodule", "update", "--init", "external_dependencies/LIBERO"],
            step="libero_repo_init",
            cwd=REPO_ROOT,
            env=env,
            log_prefix="libero",
        )
    if TEST_CACHE_PATH.exists():
        modules_path = _git_modules_path(LIBERO_REPO_PATH)
        wt_cache.mkdir(parents=True, exist_ok=True)
        _copy_tree_with_timing(LIBERO_REPO_PATH, wt_cache, "step 3a.miss.2: cache wt → shared")
        if modules_path is not None:
            modules_cache.mkdir(parents=True, exist_ok=True)
            _copy_tree_with_timing(
                modules_path, modules_cache, "step 3a.miss.3: cache modules → shared"
            )


def _libero_venv_ready(root: pathlib.Path = LIBERO_UV_ENV) -> bool:
    """Return True when the libero uv venv looks usable.

    Two subtleties:
    1. Uses pyvenv.cfg rather than bin/python: uv creates bin/python as a
       symlink into /root/.local/share/uv/python/..., which is inaccessible to
       non-root callers.  Path.is_file() returns False for unreadable symlink
       targets, so the caching gate would silently never fire.
    2. Checks for libero-*.dist-info rather than a libero/ source dir: libero
       is installed with --config-settings editable_mode=compat, which creates
       dist-info + a .pth file but no top-level package directory.
    """
    sp = root / ".venv/lib/python3.10/site-packages"
    libero_ok = (sp / "libero").is_dir() or any(sp.glob("libero-*.dist-info"))
    return (root / ".venv/pyvenv.cfg").is_file() and libero_ok


def _prepare_libero_venv(setup_block: str, env: dict[str, str]) -> None:
    """Set up the libero sim venv, using shared cache when available.

    setup_libero.sh always rm -rf's the venv before reinstalling, so this
    function skips it entirely on a cache hit and only re-runs the fast
    register_libero_envs() call that writes $HOME/.libero.
    """
    venv_python = str(LIBERO_UV_ENV / ".venv/bin/python")
    register_cmd = (
        "import os; os.environ.setdefault('MUJOCO_GL','egl');"
        "os.environ.setdefault('PYOPENGL_PLATFORM','egl');"
        "from gr00t.eval.sim.LIBERO.libero_env import register_libero_envs;"
        "register_libero_envs()"
    )

    local_ready = _libero_venv_ready()
    shared_ready = _libero_venv_ready(SHARED_LIBERO_VENV) if not local_ready else False
    print(f"[libero/3b] readiness: local={local_ready} shared={shared_ready}", flush=True)

    if local_ready:
        print("[libero/3b] venv local hit — skipping setup_libero.sh", flush=True)
    elif shared_ready:
        LIBERO_UV_ENV.parent.mkdir(parents=True, exist_ok=True)
        # Wipe any stale local state from a previous run before copying.
        if LIBERO_UV_ENV.is_symlink():
            LIBERO_UV_ENV.unlink()
        elif LIBERO_UV_ENV.is_dir():
            shutil.rmtree(LIBERO_UV_ENV)
        _copy_tree_with_timing(
            SHARED_LIBERO_VENV,
            LIBERO_UV_ENV,
            "step 3b.hit: copy shared venv → local (avoid NFS-bound imports)",
            symlinks=True,
        )
    else:
        print("[libero/3b] venv cache miss — running setup_libero.sh", flush=True)
        with timed("step 3b.miss.1: setup_libero.sh (uv install)"):
            run_bash_blocks([setup_block], cwd=REPO_ROOT, env=env)
        if TEST_CACHE_PATH.exists() and _libero_venv_ready():
            _copy_tree_with_timing(
                LIBERO_UV_ENV,
                SHARED_LIBERO_VENV,
                "step 3b.miss.2: cache venv → shared",
                symlinks=True,
            )
        return  # setup_libero.sh already ran register_libero_envs

    # Fast path: venv came from cache — just re-register envs (~5 s)
    import subprocess as _sp

    with timed("step 3b.hit: register_libero_envs subprocess"):
        _sp.run([venv_python, "-c", register_cmd], env=env, check=True, input=b"n\n")


def _dataset_ready(dataset_root: pathlib.Path) -> bool:
    """Return True when the LIBERO dataset looks complete enough to reuse."""
    modality_path = dataset_root / "meta/modality.json"
    videos_dir = dataset_root / "videos"
    if not modality_path.is_file() or not videos_dir.is_dir():
        return False
    return next(videos_dir.rglob("*.mp4"), None) is not None


def _point_repo_dataset_to_shared() -> None:
    """Point the repo-local dataset path at the shared cached dataset."""
    if DATASET_ROOT.is_symlink():
        if DATASET_ROOT.resolve() == SHARED_DATASET_ROOT.resolve():
            return
        DATASET_ROOT.unlink()
    elif DATASET_ROOT.exists():
        # Keep an existing real local dataset intact rather than replacing it.
        return

    DATASET_ROOT.parent.mkdir(parents=True, exist_ok=True)
    DATASET_ROOT.symlink_to(SHARED_DATASET_ROOT, target_is_directory=True)


def _prepare_libero_dataset(blocks: list, env: dict[str, str]) -> None:
    """Populate the LIBERO spatial dataset once on shared storage and reuse it."""
    if _dataset_ready(DATASET_ROOT):
        return

    if _dataset_ready(SHARED_DATASET_ROOT):
        _point_repo_dataset_to_shared()
        return

    download_code = find_block(
        blocks, "libero_spatial_no_noops_1.0.0_lerobot", language="bash"
    ).code
    if TEST_CACHE_PATH.exists():
        download_code = download_code.replace(
            "examples/LIBERO/libero_spatial_no_noops_1.0.0_lerobot/",
            f"{SHARED_DATASET_ROOT}/",
        )

    run_bash_blocks([download_code], cwd=REPO_ROOT, env=env)

    if _dataset_ready(SHARED_DATASET_ROOT):
        _point_repo_dataset_to_shared()
        return

    assert _dataset_ready(DATASET_ROOT), f"Expected LIBERO dataset at {DATASET_ROOT}"


@pytest.mark.gpu
@pytest.mark.timeout(1200)
def test_libero_readme_workflow_executes_via_subprocess() -> None:
    """Run the LIBERO README finetune (libero_spatial) then server+client eval."""

    print(f"[egl] NVIDIA EGL vendor file: {find_nvidia_egl_vendor_file()}", flush=True)

    env = {**os.environ, "MUJOCO_GL": "egl", "PYOPENGL_PLATFORM": "egl"}
    blocks = extract_code_blocks(README)

    # Step 1: Download + copy modality once, preferring the shared mounted dataset cache.
    with timed("step 1: dataset prep"):
        _prepare_libero_dataset(blocks, env)

    # Remove any leftover output dir so the trainer starts fresh rather than
    # trying to resume from a stale checkpoint (which would fail with --skip_weight_loading
    # if a previous run used the real weights and produced a different architecture).
    if MODEL_CHECKPOINT.parent.exists():
        shutil.rmtree(MODEL_CHECKPOINT.parent)

    # Step 2: Finetune — inline README values are replaced to keep the run short.
    # --skip_weight_loading skips loading the 3B checkpoint weights (saves ~80s); the
    # processor/tokenizer is still loaded from the checkpoint so the data pipeline
    # is exercised correctly.  Weights don't matter for a 2-step smoke test.
    finetune_code = replace_once(
        replace_once(
            replace_once(
                replace_once(
                    find_block(blocks, "--output-dir /tmp/libero_spatial", language="bash").code,
                    "NUM_GPUS=8",
                    "NUM_GPUS=1",
                ),
                "MAX_STEPS=20000",
                f"MAX_STEPS={TRAINING_STEPS}",
            ),
            "SAVE_STEPS=1000",
            f"SAVE_STEPS={TRAINING_STEPS}",
        ),
        "GLOBAL_BATCH_SIZE=640",
        "GLOBAL_BATCH_SIZE=2",
    )
    # Pass --skip_weight_loading after `--` so finetune.sh routes it to EXTRA_ARGS and on
    # to launch_finetune.py (unknown args before `--` cause finetune.sh to exit 1).
    finetune_code = finetune_code.rstrip() + " -- --skip_weight_loading"
    with timed("step 2: finetune"):
        run_bash_blocks(
            [finetune_code],
            cwd=REPO_ROOT,
            env={
                **env,
                "USE_WANDB": "0",
                "DATALOADER_NUM_WORKERS": "0",
                "SHARD_SIZE": "64",
                "NUM_SHARDS_PER_EPOCH": "1",
            },
        )
    assert MODEL_CHECKPOINT.exists(), (
        f"Expected model checkpoint after finetune: {MODEL_CHECKPOINT}"
    )

    model_server_host = "127.0.0.1"
    model_server_port = 5552

    # Build server and rollout command strings now (checkpoint exists after finetune).
    server_code = replace_once(
        find_block(blocks, "checkpoints/GR00T-N1.7-LIBERO/libero_10", language="bash").code,
        "checkpoints/GR00T-N1.7-LIBERO/libero_10",
        str(MODEL_CHECKPOINT),
    )
    server_code += f" --device cuda:0 --host {model_server_host} --port {model_server_port}"

    rollout_code = replace_once(
        replace_once(
            replace_once(
                replace_once(
                    find_block(blocks, "libero_uv/.venv/bin/python", language="bash").code,
                    "--n-episodes 10",
                    "--n-episodes 1",
                ),
                "--policy-client-port 5555",
                f"--policy-client-port {model_server_port}",
            ),
            "--max-episode-steps 720",
            "--max-episode-steps 2",
        ),
        "--n-envs 5",
        "--n-envs 1",
    )

    # Steps 3 + 4 overlapped: start the model server immediately after finetune so
    # model loading runs in parallel with the libero sim venv setup (which takes
    # several minutes on a cache miss).
    assert_port_available(model_server_host, model_server_port)
    model_server_proc, server_log = start_server_process(server_code, cwd=REPO_ROOT, env=env)

    with timed("step 3a: libero repo prep"):
        _prepare_libero_repo(env)
    with timed("step 3b: sim venv setup"):
        _prepare_libero_venv(find_block(blocks, "setup_libero.sh", language="bash").code, env)

    with timed("step 4: server startup"):
        wait_for_server_ready(
            proc=model_server_proc,
            host=model_server_host,
            port=model_server_port,
            timeout_s=float(
                os.getenv("LIBERO_SERVER_STARTUP_SECONDS", str(DEFAULT_SERVER_STARTUP_SECONDS))
            ),
            server_log=server_log,
        )

    try:
        with timed("step 5: rollout"):
            simulation_result, _ = run_subprocess_step(
                ["bash", "-c", rollout_code],
                step="libero_rollout",
                cwd=REPO_ROOT,
                env=env,
                log_prefix="libero",
                failure_prefix="LIBERO rollout failed",
                output_tail_chars=4000,
            )
        simulation_output = (simulation_result.stdout or "") + (simulation_result.stderr or "")
        assert "results:" in simulation_output, (
            "Simulation output did not include expected 'results:' marker.\n"
            f"output_tail=\n{simulation_output[-4000:]}"
        )
        assert "success rate:" in simulation_output, (
            "Simulation output did not include expected 'success rate:' marker.\n"
            f"output_tail=\n{simulation_output[-4000:]}"
        )
    finally:
        if model_server_proc.poll() is None:
            model_server_proc.terminate()
            try:
                model_server_proc.wait(timeout=15)
            except subprocess.TimeoutExpired:
                model_server_proc.kill()
                model_server_proc.wait(timeout=15)