diff --git a/pyproject.toml b/pyproject.toml index 18f171a..ac22123 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "duobench" version = "0.1.2" description="FR3 Duo Bimanual Manipulation Benchmark" dependencies = [ - "rcs-core>=0.7.3", + "rcs-core>=0.7.4", "av>=17.0.1", "einops>=0.8.2", "scipy>=1.15.3", diff --git a/src/duobench/__init__.py b/src/duobench/__init__.py index 9425a91..3da5dbc 100644 --- a/src/duobench/__init__.py +++ b/src/duobench/__init__.py @@ -5,41 +5,17 @@ import numpy as np import rcs -from rcs import download_assets, get_prefix +from rcs import get_prefix from rcs._core import common DUOBENCH_GITHUB_ASSET_ARCHIVE_URL = "https://github.com/RobotControlStack/duobench/archive/refs/tags/{tag}.zip" -# written after a complete download, marks which duobench version the assets belong to -ASSETS_VERSION_FILE = Path("assets/.duobench_version") - -def _download_assets(version: str, prefix: Path, github_asset_archive_url: str, assets_dirname: str) -> None: - download_assets(version, prefix, github_asset_archive_url, assets_dirname) - (prefix / ASSETS_VERSION_FILE).write_text(version) - - -def _get_duobench_prefix(env_prefix: str | None, editable_prefix: Path, default_prefix: Path) -> str: - # rcs.get_prefix checks for an RCS asset that is not part of the duobench archive, so it would - # re-download the duobench assets on every import; skip it if the assets of this version are there - if env_prefix or not (editable_prefix.resolve() / "assets").is_dir(): - prefix = Path(env_prefix).expanduser().resolve() if env_prefix else default_prefix.resolve() - version_file = prefix / ASSETS_VERSION_FILE - if version_file.is_file() and version_file.read_text() == __version__: - return str(prefix) - return get_prefix( - env_prefix, - editable_prefix, - default_prefix, - __version__, - DUOBENCH_GITHUB_ASSET_ARCHIVE_URL, - download_fn=_download_assets, - ) - - -DUOBENCH_PREFIX = _get_duobench_prefix( +DUOBENCH_PREFIX = get_prefix( os.environ.get("DUOBENCH_PREFIX"), Path(__file__).resolve().parents[2], Path.home() / ".duobench", + __version__, + DUOBENCH_GITHUB_ASSET_ARCHIVE_URL, ) CAMERA_PATHS: dict[str, str] = {} diff --git a/src/duobench/utils/vention_config.py b/src/duobench/utils/vention_config.py index cc729cb..8951f41 100644 --- a/src/duobench/utils/vention_config.py +++ b/src/duobench/utils/vention_config.py @@ -1,8 +1,14 @@ +import copy import time +import gymnasium as gym from rcs import DEFAULT_TRANSFORMS, SCENE_PATHS from rcs.envs.configs import EmptyWorldFR3Duo -from rcs.envs.scenes import SimEnvCreatorConfig +from rcs.envs.scenes import MjModel, SimEnvCreatorConfig + +# MuJoCo sizes the arena from the scene at rest (15-22 MiB for the DuoBench scenes), contact-rich episodes +# (e.g. ball_maze with ~170 contacts) exceed it and abort with "mj_stackAlloc: insufficient memory" +MUJOCO_MEMORY = 256 * 1024 * 1024 class VentionSceneFR3Duo(EmptyWorldFR3Duo): @@ -10,8 +16,17 @@ def config(self) -> SimEnvCreatorConfig: cfg = super().config() cfg.scene = SCENE_PATHS["vention_world"] cfg.root_frame_to_world = DEFAULT_TRANSFORMS["VENTION_HEIGHT_OFFSET"] + cfg.sim_cfg.memory = MUJOCO_MEMORY return cfg + def create_env_from_model(self, cfg: SimEnvCreatorConfig, mjmodel: MjModel) -> gym.Env: + # eval setups often replace sim_cfg with a fresh SimConfig, keep the larger arena unless set explicitly + if cfg.sim_cfg.memory < 0: + cfg = copy.copy(cfg) + cfg.sim_cfg = copy.copy(cfg.sim_cfg) + cfg.sim_cfg.memory = MUJOCO_MEMORY + return super().create_env_from_model(cfg, mjmodel) + if __name__ == "__main__": scene = VentionSceneFR3Duo() diff --git a/src/tests/test_asset_prefix.py b/src/tests/test_asset_prefix.py deleted file mode 100644 index 7cf4fe7..0000000 --- a/src/tests/test_asset_prefix.py +++ /dev/null @@ -1,48 +0,0 @@ -from pathlib import Path - -import pytest - -import duobench - - -@pytest.fixture() -def downloads(monkeypatch: pytest.MonkeyPatch) -> list[Path]: - calls: list[Path] = [] - - def fake_download(_version: str, prefix: Path, _url: str, assets_dirname: str) -> None: - calls.append(prefix) - (prefix / assets_dirname).mkdir(parents=True, exist_ok=True) - - monkeypatch.setattr(duobench, "download_assets", fake_download) - return calls - - -def test_assets_downloaded_once(tmp_path: Path, downloads: list[Path]): - prefix = tmp_path / "prefix" - for _ in range(3): - assert duobench._get_duobench_prefix(str(prefix), tmp_path / "pkg", tmp_path / "default") == str(prefix) - assert downloads == [prefix] - assert (prefix / duobench.ASSETS_VERSION_FILE).read_text() == duobench.__version__ - - -def test_default_prefix_used_without_env(tmp_path: Path, downloads: list[Path]): - default = tmp_path / "default" - for _ in range(2): - assert duobench._get_duobench_prefix(None, tmp_path / "pkg", default) == str(default) - assert downloads == [default] - - -def test_assets_redownloaded_on_version_change(tmp_path: Path, downloads: list[Path]): - prefix = tmp_path / "prefix" - (prefix / "assets").mkdir(parents=True) - (prefix / duobench.ASSETS_VERSION_FILE).write_text("0.0.0") - duobench._get_duobench_prefix(str(prefix), tmp_path / "pkg", tmp_path / "default") - assert downloads == [prefix] - assert (prefix / duobench.ASSETS_VERSION_FILE).read_text() == duobench.__version__ - - -def test_editable_install_skips_download(tmp_path: Path, downloads: list[Path]): - editable = tmp_path / "pkg" - (editable / "assets").mkdir(parents=True) - assert duobench._get_duobench_prefix(None, editable, tmp_path / "default") == str(editable) - assert downloads == [] diff --git a/src/tests/test_mujoco_memory.py b/src/tests/test_mujoco_memory.py new file mode 100644 index 0000000..7b914be --- /dev/null +++ b/src/tests/test_mujoco_memory.py @@ -0,0 +1,32 @@ +import pytest +from rcs._core.sim import SimConfig + +from duobench.tasks.ball_maze import BallMazeEnvConfig +from duobench.utils.vention_config import MUJOCO_MEMORY + + +def _arena_size(sim_cfg: SimConfig | None = None) -> int: + scene = BallMazeEnvConfig() + cfg = scene.config() + cfg.headless = True + if sim_cfg is not None: + cfg.sim_cfg = sim_cfg + env = scene.create_env(cfg) + try: + return env.get_wrapper_attr("sim").model.narena + finally: + env.close() + + +def test_default_config_uses_large_arena(): + assert _arena_size() == MUJOCO_MEMORY + + +def test_replaced_sim_config_keeps_large_arena(): + # e.g. eval harnesses that set their own SimConfig + assert _arena_size(SimConfig(async_control=True, frequency=30)) == MUJOCO_MEMORY + + +@pytest.mark.parametrize("memory", [64 * 1024 * 1024]) +def test_explicit_memory_is_kept(memory: int): + assert _arena_size(SimConfig(memory=memory)) == memory