feat(sim): mujoco memory config to avoid crash of large scenes

python/rcs/_core/sim.pyi

@@ -140,6 +140,7 @@ class SimConfig:
     async_control: bool
     frequency: float
     max_convergence_steps: int
+    memory: int
     realtime: bool
     def __copy__(self) -> SimConfig: ...
     def __deepcopy__(self, arg0: dict) -> SimConfig: ...
@@ -149,6 +150,7 @@ class SimConfig:
         realtime: bool = False,
         frequency: float = 30.0,
         max_convergence_steps: int = 500,
+        memory: int = -1,
     ) -> None: ...
 
 class SimGripper(rcs._core.common.Gripper):

python/rcs/sim/sim.py

@@ -65,6 +65,9 @@ def gui_loop(gui_uuid: str, close_event):
 class Sim(_Sim):
     def __init__(self, mjmdl: str | PathLike | ModelComposer, cfg: SimConfig | None = None):
         if isinstance(mjmdl, ModelComposer):
+            # the arena size is fixed at compile time
+            if cfg is not None and cfg.memory >= 0:
+                mjmdl.spec.memory = cfg.memory
             self.model = mjmdl.get_model()
         else:
             mjmdl = Path(mjmdl)

python/tests/test_sim_envs.py

@@ -328,3 +328,17 @@ def test_relative_zero_action_joints(self):
         act.update(GripperDictType(gripper=np.array([1.0])))
         obs, _, _, _, info = env.step(act)
         self.assert_no_pose_change(info, obs_initial, obs)
+
+
+def test_sim_config_memory_sets_arena_size():
+    scene = EmptyWorldFR3()
+    cfg = copy.deepcopy(scene.config())
+    cfg.headless = True
+    prefixed_cfg = scene.prefixed_cfg(cfg)
+    assert prefixed_cfg.sim_cfg.memory == -1
+    default = sim.Sim(scene.create_model(prefixed_cfg), prefixed_cfg.sim_cfg)
+
+    prefixed_cfg.sim_cfg.memory = 64 * 1024 * 1024
+    enlarged = sim.Sim(scene.create_model(prefixed_cfg), prefixed_cfg.sim_cfg)
+    assert enlarged.model.narena == 64 * 1024 * 1024
+    assert default.model.narena < enlarged.model.narena

src/pybind/rcs.cpp

@@ -751,24 +751,27 @@ PYBIND11_MODULE(_core, m) {
   rcs::sim::SimConfig default_sim_cfg = rcs::sim::SimConfig();
   py::class_<rcs::sim::SimConfig>(sim, "SimConfig")
       .def(py::init([](bool async_control, bool realtime, double frequency,
-                       int max_convergence_steps) {
+                       int max_convergence_steps, long long memory) {
              rcs::sim::SimConfig config;
              config.async_control = async_control;
              config.realtime = realtime;
              config.frequency = frequency;
              config.max_convergence_steps = max_convergence_steps;
+             config.memory = memory;
              return config;
            }),
            py::arg("async_control") = default_sim_cfg.async_control,
            py::arg("realtime") = default_sim_cfg.realtime,
            py::arg("frequency") = default_sim_cfg.frequency,
            py::arg("max_convergence_steps") =
-               default_sim_cfg.max_convergence_steps)
+               default_sim_cfg.max_convergence_steps,
+           py::arg("memory") = default_sim_cfg.memory)
       .def_readwrite("async_control", &rcs::sim::SimConfig::async_control)
       .def_readwrite("realtime", &rcs::sim::SimConfig::realtime)
       .def_readwrite("frequency", &rcs::sim::SimConfig::frequency)
       .def_readwrite("max_convergence_steps",
                      &rcs::sim::SimConfig::max_convergence_steps)
+      .def_readwrite("memory", &rcs::sim::SimConfig::memory)
       .def("__copy__",
            [](const rcs::sim::SimConfig& self) {
              return rcs::sim::SimConfig(self);

src/sim/sim.h

@@ -34,6 +34,11 @@ struct SimConfig {
   bool realtime = false;
   double frequency = 30;  // in Hz
   int max_convergence_steps = 500;
+  // MuJoCo arena (stack + constraint) size in bytes, applied when a scene is
+  // compiled from a ModelComposer. -1 keeps the model's <size memory> or
+  // MuJoCo's automatic sizing, which can be too small for scenes with many
+  // contacts ("mj_stackAlloc: insufficient memory").
+  long long memory = -1;
 };
 
 struct Callback {