# 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 shlex import shutil import subprocess 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, demo_dataset_tree_ready, get_root, resolve_model_checkpoint_path, run_subprocess_step, start_server_process, timed, wait_for_server_ready, ) REPO_ROOT = get_root() LOGGER = logging.getLogger(__name__) README = REPO_ROOT / "examples/robocasa/README.md" ROBOCASA_CHECKPOINT_PATH = os.environ.get("ROBOCASA_CHECKPOINT_PATH", "") ROBOCASA_EMBODIMENT_TAG = os.environ.get("ROBOCASA_EMBODIMENT_TAG", "ROBOCASA_PANDA_OMRON") TRAINING_STEPS = 2 MODEL_CHECKPOINT = pathlib.Path(f"/tmp/robocasa_finetune/checkpoint-{TRAINING_STEPS}") _ROBOCASA_DATASET_ENV = "ROBOCASA_DATASET_PATH" _SHARED_ROBOCASA_DATASET = TEST_CACHE_PATH / "datasets/robocasa" ROBOCASA_SUBMODULE_PATH = REPO_ROOT / "external_dependencies/robocasa" SHARED_ROBOCASA_REPO = TEST_CACHE_PATH / "repos/robocasa" ROBOCASA_ASSETS_REPO_DIR = REPO_ROOT / "external_dependencies/robocasa/robocasa/models/assets" ROBOCASA_ASSETS_SHARED_DIR = TEST_CACHE_PATH / "robocasa-assets" # Version file written alongside cached assets to detect robocasa submodule updates. _ASSETS_VERSION_FILE = ROBOCASA_ASSETS_SHARED_DIR / ".robocasa_commit" def _robocasa_submodule_commit() -> str: """Return the robocasa submodule commit hash recorded in the main repo HEAD. Uses ``git ls-tree`` against the main repo so this works even before the submodule is initialized. """ try: result = subprocess.run( ["git", "ls-tree", "HEAD", "external_dependencies/robocasa"], capture_output=True, text=True, check=True, cwd=REPO_ROOT, ) # output: "160000 commit \texternal_dependencies/robocasa" parts = result.stdout.split() return parts[2] if len(parts) >= 3 else "unknown" except Exception: return "unknown" def _shared_asset_dirs() -> list[pathlib.Path]: """Return all top-level subdirectories in the shared asset cache.""" if not ROBOCASA_ASSETS_SHARED_DIR.is_dir(): return [] return [p for p in ROBOCASA_ASSETS_SHARED_DIR.iterdir() if p.is_dir()] def _shared_assets_ready() -> bool: """Return True when the shared asset cache is present, non-empty, and matches the current robocasa submodule commit. A stale cache (e.g. from a previous robocasa version that lacked newer fixture files) is treated as not-ready so the assets are re-downloaded. """ if not _ASSETS_VERSION_FILE.is_file(): return False if _ASSETS_VERSION_FILE.read_text().strip() != _robocasa_submodule_commit(): return False for d in _shared_asset_dirs(): try: if next((f for f in d.rglob("*") if f.is_file()), None) is not None: return True except OSError: pass return False def _assert_required_assets_present() -> None: """Raise if the repo asset directory is empty.""" if not ROBOCASA_ASSETS_REPO_DIR.is_dir() or not any(ROBOCASA_ASSETS_REPO_DIR.iterdir()): raise RuntimeError(f"RoboCasa assets missing at {ROBOCASA_ASSETS_REPO_DIR}") def _point_repo_assets_to_shared() -> None: """Symlink all shared asset subdirectories into the repo asset path.""" ROBOCASA_ASSETS_REPO_DIR.mkdir(parents=True, exist_ok=True) for shared_dir in _shared_asset_dirs(): repo_dir = ROBOCASA_ASSETS_REPO_DIR / shared_dir.name if repo_dir.is_symlink(): if repo_dir.resolve() == shared_dir.resolve(): continue repo_dir.unlink() elif repo_dir.exists(): shutil.rmtree(repo_dir) repo_dir.symlink_to(shared_dir, target_is_directory=True) def _move_repo_assets_to_shared() -> None: """Move all downloaded repo asset subdirectories into the shared cache. Also writes a version file recording the current robocasa submodule commit so that stale caches are detected when the submodule is updated. """ ROBOCASA_ASSETS_SHARED_DIR.mkdir(parents=True, exist_ok=True) if not ROBOCASA_ASSETS_REPO_DIR.is_dir(): return for src in ROBOCASA_ASSETS_REPO_DIR.iterdir(): if not src.is_dir() or src.is_symlink(): continue dst = ROBOCASA_ASSETS_SHARED_DIR / src.name if dst.exists() or dst.is_symlink(): if dst.is_dir() and not dst.is_symlink(): shutil.rmtree(dst) else: dst.unlink() # Use cp -r + rm -rf instead of shutil.move: the repo and the shared # PVC are on different filesystems, so shutil.move falls back to a slow # Python-level copytree that times out on large asset dirs (e.g. # generative_textures with thousands of PNGs). subprocess.run(["cp", "-r", str(src), str(dst)], check=True) shutil.rmtree(str(src)) _ASSETS_VERSION_FILE.write_text(_robocasa_submodule_commit()) def _remove_dangling_repo_asset_symlinks() -> None: """Delete repo asset symlinks that point to missing targets.""" if not ROBOCASA_ASSETS_REPO_DIR.is_dir(): return for repo_dir in ROBOCASA_ASSETS_REPO_DIR.iterdir(): if repo_dir.is_symlink() and not repo_dir.exists(): repo_dir.unlink() def _robocasa_submodule_initialized() -> bool: return (ROBOCASA_SUBMODULE_PATH / ".git").is_file() def _git_modules_path(submodule_path: pathlib.Path) -> pathlib.Path | None: 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 _prepare_robocasa_repo(env: dict[str, str]) -> None: """Populate external_dependencies/robocasa from shared cache, or init and cache it.""" if _robocasa_submodule_initialized(): return wt_cache = SHARED_ROBOCASA_REPO / "wt" modules_cache = SHARED_ROBOCASA_REPO / "modules" if (wt_cache / ".git").is_file() and modules_cache.exists(): print(f"[robocasa] restoring submodule from cache {wt_cache}", flush=True) shutil.copytree(wt_cache, ROBOCASA_SUBMODULE_PATH, dirs_exist_ok=True) modules_path = _git_modules_path(ROBOCASA_SUBMODULE_PATH) if modules_path is not None: modules_path.mkdir(parents=True, exist_ok=True) shutil.copytree(modules_cache, modules_path, dirs_exist_ok=True) return # The directory may exist but be uninitialized (no .git file) due to CI # checkout strategies that populate submodule dirs without git-initializing them. # Remove it so git submodule update --init can clone cleanly. if ROBOCASA_SUBMODULE_PATH.exists() and not _robocasa_submodule_initialized(): shutil.rmtree(ROBOCASA_SUBMODULE_PATH) run_subprocess_step( ["git", "submodule", "update", "--init", "external_dependencies/robocasa"], step="robocasa_repo_init", cwd=REPO_ROOT, env=env, log_prefix="robocasa", ) if TEST_CACHE_PATH.exists(): modules_path = _git_modules_path(ROBOCASA_SUBMODULE_PATH) print(f"[robocasa] caching submodule to {wt_cache}", flush=True) wt_cache.mkdir(parents=True, exist_ok=True) shutil.copytree(ROBOCASA_SUBMODULE_PATH, wt_cache, dirs_exist_ok=True) if modules_path is not None: modules_cache.mkdir(parents=True, exist_ok=True) shutil.copytree(modules_path, modules_cache, dirs_exist_ok=True) def _build_runtime_env(skip_download_assets: str) -> dict[str, str]: """Build the runtime environment used by setup, model server, and rollout.""" return {**os.environ, "SKIP_DOWNLOAD_ASSETS": skip_download_assets, "INSTALL_FLASH_ATTN": "0"} def _resolve_robocasa_dataset() -> pathlib.Path: """Return the RoboCasa LeRobot dataset used for short finetuning.""" env_path_str = os.environ.get(_ROBOCASA_DATASET_ENV, "").strip() if env_path_str: env_path = pathlib.Path(env_path_str).expanduser().resolve() assert demo_dataset_tree_ready(env_path), ( f"{_ROBOCASA_DATASET_ENV} does not point to a complete LeRobot dataset: {env_path}" ) return env_path if demo_dataset_tree_ready(_SHARED_ROBOCASA_DATASET): return _SHARED_ROBOCASA_DATASET pytest.skip( "RoboCasa finetune dataset not found. Set " f"{_ROBOCASA_DATASET_ENV} to a RoboCasa LeRobot dataset from Hugging Face or " f"populate {_SHARED_ROBOCASA_DATASET}." ) def _resolve_robocasa_checkpoint(blocks, env: dict[str, str]) -> str: """Use ROBOCASA_CHECKPOINT_PATH or create a short fine-tuned checkpoint.""" if ROBOCASA_CHECKPOINT_PATH: return ROBOCASA_CHECKPOINT_PATH dataset_path = _resolve_robocasa_dataset() with timed("step 1a: base model prep"): base_model_path = resolve_model_checkpoint_path( hf_repo_id="nvidia/GR00T-N1.7-3B", path_override_env="GROOT_MODEL_PATH", repo_root=REPO_ROOT, ) if MODEL_CHECKPOINT.parent.exists(): shutil.rmtree(MODEL_CHECKPOINT.parent) finetune_code = replace_once( replace_once( replace_once( replace_once( replace_once( replace_once( replace_once( replace_once( find_block( blocks, "--output-dir /tmp/robocasa_finetune", language="bash", ).code, "NUM_GPUS=8", "NUM_GPUS=1", ), "MAX_STEPS=60000", f"MAX_STEPS={TRAINING_STEPS}", ), "SAVE_STEPS=2000", f"SAVE_STEPS={TRAINING_STEPS}", ), "GLOBAL_BATCH_SIZE=512", "GLOBAL_BATCH_SIZE=2", ), "nvidia/GR00T-N1.7-3B", str(base_model_path), ), '--dataset-path "$DATASET_PATH"', f"--dataset-path {shlex.quote(str(dataset_path))}", ), "--embodiment-tag ROBOCASA_PANDA_OMRON", f"--embodiment-tag {ROBOCASA_EMBODIMENT_TAG}", ), "--output-dir /tmp/robocasa_finetune", f"--output-dir {MODEL_CHECKPOINT.parent}", ) finetune_code = finetune_code.rstrip() + " -- --skip_weight_loading" with timed("step 1b: 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}" ) return str(MODEL_CHECKPOINT) @pytest.mark.gpu @pytest.mark.timeout(2700) def test_robocasa_readme_eval_flow() -> None: """Run the RoboCasa README finetune/server/client eval flow.""" # Environment setup: # 1) If assets already exist on shared PVC, reuse them by symlinking. # 2) Otherwise run setup with download enabled. shared_assets_ready = _shared_assets_ready() if not shared_assets_ready: _remove_dangling_repo_asset_symlinks() skip_download_assets = "1" if shared_assets_ready else "0" env = _build_runtime_env(skip_download_assets=skip_download_assets) blocks = extract_code_blocks(README) with timed("step 1: checkpoint prep"): checkpoint_path = _resolve_robocasa_checkpoint(blocks, env) with timed("step 2: robocasa repo prep"): _prepare_robocasa_repo(env) if shared_assets_ready: _point_repo_assets_to_shared() # Step 2: Setup sim with timed("step 3: sim venv setup (setup_RoboCasa.sh)"): run_bash_blocks( [find_block(blocks, "setup_RoboCasa.sh", language="bash")], cwd=REPO_ROOT, env=env, force_yes=True, ) # When setup performs a fresh download, move those assets into shared PVC # so subsequent runs can skip download and reuse the cached shared copy. if not shared_assets_ready: _move_repo_assets_to_shared() _point_repo_assets_to_shared() _assert_required_assets_present() model_server_host = "127.0.0.1" model_server_port = 5551 # Step 4: Server — N1.7 RoboCasa evaluation requires a finetuned checkpoint. server_code = replace_once( replace_once( find_block(blocks, "", language="bash").code, "", checkpoint_path, ), "--embodiment-tag ROBOCASA_PANDA_OMRON", f"--embodiment-tag {ROBOCASA_EMBODIMENT_TAG}", ) server_code += f" --device cuda:0 --host {model_server_host} --port {model_server_port}" # Step 5: Rollout — substitute test-safe values rollout_code = replace_once( replace_once( replace_once( replace_once( find_block(blocks, "rollout_policy.py", 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", ) 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 4: server startup"): wait_for_server_ready( proc=model_server_proc, host=model_server_host, port=model_server_port, timeout_s=float( os.getenv("ROBOCASA_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="robocasa_rollout", cwd=REPO_ROOT, env=env, log_prefix="robocasa", failure_prefix="RoboCasa 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)