From a83a4bedbeee531adc91c79d34ae5c3c52dcdd88 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 20:09:11 +0200 Subject: [PATCH] fix: cleanup --- scripts/simulate.py | 6 +++++- .../environment/BrittleStarJaxEnvWrapper.py | 20 +++++++++++++++++-- tests/test_configs.py | 11 ++++++---- tests/test_morphology_render.py | 6 ++++++ 4 files changed, 36 insertions(+), 7 deletions(-) diff --git a/scripts/simulate.py b/scripts/simulate.py index 8fb597d..666d048 100644 --- a/scripts/simulate.py +++ b/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 ======= diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 7a6d161..538ba7a 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/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() diff --git a/tests/test_configs.py b/tests/test_configs.py index 43ef95e..0a6a651 100644 --- a/tests/test_configs.py +++ b/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 diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py index 423a575..6b435cf 100644 --- a/tests/test_morphology_render.py +++ b/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