fix: increase sim ram to avoid crashes in contact rich scenes

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",

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] = {}

src/duobench/utils/vention_config.py

@@ -1,17 +1,32 @@
+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):
     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()

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