fix: cleanup

scripts/simulate.py

@@ -12,7 +12,6 @@ from omegaconf import DictConfig, OmegaConf
 from pathlib import Path
 
 from brittle_star_project import (
-    Backend,
     BrittleStarEnv,
     BrittleStarEnvFactory,
     SimulationConfig,
@@ -37,7 +36,12 @@ def main(dict_cfg: DictConfig) -> None:
     # Use the configurable settings from the simulation group
     backend = config.simulation.backend
     model_type = config.simulation.model_type
+
+    # Hydra chdir changes CWD; we map CLI relative paths relative to invocation originally.
     model_path = config.simulation.model_path
+    if model_path is not None:
+        model_path = hydra.utils.to_absolute_path(model_path)
+
     seed = config.experiment.seed
 
     # ======= ENVIRONMENT SETUP =======

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -28,6 +28,10 @@ class BrittleStarJaxEnvWrapper:
             self._backend, self._morphology, self._arena, self._env_config
         )
 
+        from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
+
+        self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm)
+
         self._vectorized_reset = jax.jit(jax.vmap(self._env.reset))
         self._vectorized_step = jax.jit(jax.vmap(self._env.step))
         self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample))
@@ -61,7 +65,13 @@ class BrittleStarJaxEnvWrapper:
         self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
         self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
         env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
-        return self._vectorized_reset(rng=env_rngs)
+        state = self._vectorized_reset(rng=env_rngs)
+        from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched
+
+        state = state.replace(
+            observations=pad_observations_batched(state.observations, self._padding_masks)
+        )
+        return state
 
     def sample_actions(self):
         assert self._action_rng is not None, "Call reset() before sample_actions()"
@@ -71,7 +81,13 @@ class BrittleStarJaxEnvWrapper:
         return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
 
     def step(self, state, action):
-        return self._vectorized_step(state=state, action=action)
+        next_state = self._vectorized_step(state=state, action=action)
+        from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched
+
+        next_state = next_state.replace(
+            observations=pad_observations_batched(next_state.observations, self._padding_masks)
+        )
+        return next_state
 
     def close(self):
         self._env.close()

tests/test_configs.py

@@ -1,4 +1,5 @@
-from hydra import compose, initialize
+from pathlib import Path
+from hydra import compose, initialize_config_dir
 from omegaconf import OmegaConf
 from brittle_star_project.configs.main_config import BrittleStarConfig
 from brittle_star_project.configs.register_configs import register_configs
@@ -9,11 +10,12 @@ register_configs()
 
 def test_config_composition_centralized():
     """Test that the centralized configuration composes and validates correctly."""
-    with initialize(version_base="1.3", config_path="../configs"):
+    config_dir = str(Path(__file__).parent.parent / "configs")
+    with initialize_config_dir(version_base="1.3", config_dir=config_dir):
         # We compose the config; it follows main_config.yaml
         cfg = compose(config_name="main_config", overrides=["architecture=centralized"])
 
-        # Merge with the structured schema and convert to a real dataclass instance to verify validation.
+        # Merge with the structured schema and convert to a real dataclass instance
         structured_cfg = OmegaConf.to_object(
             OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg)
         )
@@ -27,7 +29,8 @@ def test_config_composition_centralized():
 
 def test_config_composition_decentralized():
     """Test that the decentralized configuration composes and validates correctly."""
-    with initialize(version_base="1.3", config_path="../configs"):
+    config_dir = str(Path(__file__).parent.parent / "configs")
+    with initialize_config_dir(version_base="1.3", config_dir=config_dir):
         cfg = compose(config_name="main_config", overrides=["architecture=decentralized"])
 
         # Merge and convert to dataclass instance

tests/test_morphology_render.py

@@ -1,5 +1,11 @@
 import pytest
 import os
+import sys
+
+# CRITICAL for headless cross-platform testing (devcontainers etc)
+if sys.platform == "linux" and "DISPLAY" not in os.environ and "WAYLAND_DISPLAY" not in os.environ:
+    os.environ.setdefault("MUJOCO_GL", "egl")
+
 import mujoco
 from PIL import Image
 from brittle_star_project.environment.env_config import EnvConfig, MorphologyConfig, ArenaConfig