tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

feat(simulate): metadata config loading

Browse source
This commit is contained in:
Tibo De Peuter 2026-04-28 13:17:24 +02:00
parent a3a1f3643f
commit 8f2a5d25ed
Signed by: tdpeuter
SSH key fingerprint: SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
5 changed files with 190 additions and 476 deletions
Download patch file Download diff file Expand all files Collapse all files

557
scripts/simulate.py
View file

@ -1,11 +1,10 @@
"""Simulate a trained policy in the MuJoCo viewer.
Uses Hydra to load the same BrittleStarConfig that was used during training.
Override settings via CLI, e.g.:
python scripts/simulate.py morphology=3_arms
To replay a run using the *exact* Hydra config used during training, pass:
python scripts/simulate.py simulation.trained_config_path=runs/.../.hydra/config.yaml \
Automatically extracts the training configuration (morphology, environment, etc.)
from the sidecar metadata YAML file to ensure simulation perfectly matches training.
Override simulation settings via CLI, e.g.:
uv run scripts/simulate.py \
simulation.morphology_override=config/morphology/3_arms.yaml \
simulation.model_path=runs/.../final_model.flax
"""
@ -22,85 +21,24 @@ import jax
import jax.numpy as jnp
import numpy as np
import yaml
from omegaconf import DictConfig, OmegaConf, open_dict
from omegaconf import DictConfig, OmegaConf
from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory
from brittle_star_project.configs.main_config import BrittleStarConfig
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
from brittle_star_project.environment.obs_processing import create_obs_processor
from brittle_star_project.environment.env_config import ObservationBoundsConfig
from brittle_star_project.environment.env_config import (
MorphologyConfig,
ArenaConfig,
EnvConfig,
ObservationBoundsConfig,
)
def _dense_layer_sizes_from_params(params: Any) -> list[int]:
"""Infer GenericDenseLayersWithActivation.layer_sizes from a Flax params tree."""
class PolicyAgent:
"""Wraps a trained Flax actor for deterministic inference."""
try:
dense_params = params["params"]
except Exception as exc:
raise ValueError("Unexpected sensor params structure (missing 'params')") from exc
layer_sizes: list[int] = []
idx = 0
while True:
key = f"Dense_{idx}"
if key not in dense_params:
break
kernel = dense_params[key]["kernel"]
layer_sizes.append(int(np.asarray(kernel).shape[1]))
idx += 1
if not layer_sizes:
raise ValueError("Could not infer Dense_* layers from sensor params")
return layer_sizes
def _infer_action_dim_from_actor_params(params: Any) -> int | None:
"""Best-effort infer action_dim from a Flax Actor params tree."""
try:
dense0 = params["params"]["Dense_0"]
bias = dense0.get("bias")
kernel = dense0.get("kernel")
except Exception:
return None
if bias is not None:
try:
return int(np.asarray(bias).shape[0])
except Exception:
return None
if kernel is not None:
try:
return int(np.asarray(kernel).shape[1])
except Exception:
return None
return None
def _has_cli_override(overrides: list[str], key: str) -> bool:
prefixes = (f"{key}=", f"{key}.", f"+{key}=", f"+{key}.")
return any(str(o).startswith(prefixes) for o in overrides)
def _maybe_clip_action(
action: np.ndarray,
low: np.ndarray | None,
high: np.ndarray | None,
) -> np.ndarray:
if low is None or high is None:
return action
low = np.asarray(low, dtype=np.float32).ravel()
high = np.asarray(high, dtype=np.float32).ravel()
if low.shape != action.shape or high.shape != action.shape:
return action
return np.clip(action, low, high)
# A minimal policy class to load a CleanRL/Flax checkpoint and run inference.
class CleanRLPPOPolicy:
def __init__(
self,
*,
@ -111,7 +49,28 @@ class CleanRLPPOPolicy:
) -> None:
from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
layer_sizes = _dense_layer_sizes_from_params(sensor_params)
# Infer layer sizes from params
try:
dense_params = (
sensor_params.get("params", {})
if isinstance(sensor_params, dict)
else sensor_params["params"]
)
except Exception:
dense_params = sensor_params
layer_sizes = []
idx = 0
while True:
key = f"Dense_{idx}"
if key not in dense_params:
break
layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1]))
idx += 1
if not layer_sizes:
raise ValueError("Could not infer Dense_* layers from sensor params")
self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes)
self._actor = Actor(action_dim=action_dim)
self._sensor_apply = jax.jit(self._sensor.apply)
@ -128,122 +87,31 @@ class CleanRLPPOPolicy:
*,
action_dim: int,
obs_processor: Any,
) -> "CleanRLPPOPolicy":
def _get_index(container: Any, idx: int) -> Any:
if isinstance(container, (list, tuple)):
return container[idx]
if isinstance(container, dict):
return container.get(idx, container.get(str(idx)))
raise KeyError(idx)
def _looks_like_indexed_dict(container: Any) -> bool:
return (
isinstance(container, dict)
and container
and all(str(k).isdigit() for k in container.keys())
)
def _parse_checkpoint(restored_obj: Any) -> tuple[Any, Any, Any, Any, Any]:
"""Extract checkpoint parts.
Returns (config_dict, sensor_params, actor_params, critic_params,
feature_extractor_params).
PPOTrainer saves:
flax.serialization.to_bytes([
config_dict,
[sensor_params, actor_params, critic_params, feature_extractor_params],
])
msgpack_restore() may restore lists as dicts keyed by string indices
("0", "1", ...), so we accept both shapes.
"""
cfg_part: Any | None = None
params_part: Any = restored_obj
if isinstance(restored_obj, (list, tuple)) and len(restored_obj) >= 2:
cfg_part = restored_obj[0]
params_part = restored_obj[1]
elif _looks_like_indexed_dict(restored_obj) and (
"0" in restored_obj or "1" in restored_obj
):
cfg_part = restored_obj.get("0", restored_obj.get(0))
params_part = restored_obj.get("1", restored_obj.get(1))
if _looks_like_indexed_dict(params_part):
sensor_params = _get_index(params_part, 0)
actor_params = _get_index(params_part, 1)
critic_params = _get_index(params_part, 2)
feature_extractor_params = _get_index(params_part, 3)
if sensor_params is None or actor_params is None:
raise ValueError("Missing required params in checkpoint")
return (
cfg_part,
sensor_params,
actor_params,
critic_params,
feature_extractor_params,
)
if isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
sensor_params = params_part[0]
actor_params = params_part[1]
critic_params = params_part[2] if len(params_part) >= 3 else None
feature_extractor_params = params_part[3] if len(params_part) >= 4 else None
return (
cfg_part,
sensor_params,
actor_params,
critic_params,
feature_extractor_params,
)
# Accept a plain dict-shaped Flax params mapping commonly produced
# by saving `agent_state.params` directly. Typical keys are
# 'sensor_params' and 'actor_params', or sometimes nested under 'params'.
if isinstance(restored_obj, dict):
# Top-level params dict
params_sub = restored_obj.get("params", {})
sensor_params = restored_obj.get("sensor_params") or params_sub.get("sensor_params")
actor_params = restored_obj.get("actor_params") or params_sub.get("actor_params")
critic_params = restored_obj.get("critic_params") or params_sub.get("critic_params")
feature_extractor_params = restored_obj.get(
"feature_extractor_params"
) or params_sub.get("feature_extractor_params")
# Some checkpoints only save actor+sensor as top-level
if sensor_params is not None and actor_params is not None:
return (
cfg_part,
sensor_params,
actor_params,
critic_params,
feature_extractor_params,
)
raise ValueError(
f"Unexpected checkpoint structure in {path}. "
"Expected [config_dict, [sensor_params, actor_params, critic_params, "
"feature_extractor_params]] or an equivalent dict-indexed variant."
)
) -> "PolicyAgent":
payload = path.read_bytes()
restored = flax.serialization.msgpack_restore(payload)
_cfg_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = (
_parse_checkpoint(restored)
)
ckpt_action_dim = _infer_action_dim_from_actor_params(actor_params)
if ckpt_action_dim is not None and ckpt_action_dim != action_dim:
raise ValueError(
"Checkpoint/env mismatch: "
f"checkpoint expects action_dim={ckpt_action_dim}, "
f"env provides action_dim={action_dim}. "
"Use the same Hydra config (morphology/arena/environment) "
"that was used during training."
)
sensor_params = None
actor_params = None
return CleanRLPPOPolicy(
# Extract params from restored checkpoint
if isinstance(restored, dict):
params_sub = restored.get("params", {})
sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params")
actor_params = restored.get("actor_params") or params_sub.get("actor_params")
elif isinstance(restored, (list, tuple)) and len(restored) >= 2:
params_part = restored[1]
if isinstance(params_part, dict):
sensor_params = params_part.get("0", params_part.get(0))
actor_params = params_part.get("1", params_part.get(1))
elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
sensor_params = params_part[0]
actor_params = params_part[1]
if sensor_params is None or actor_params is None:
raise ValueError(f"Could not extract sensor and actor params from checkpoint: {path}")
return PolicyAgent(
sensor_params=sensor_params,
actor_params=actor_params,
action_dim=action_dim,
@ -273,23 +141,30 @@ def _target_reached(*, state: Any) -> bool:
return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False))
def _rollout_one_episode_headless(
def _maybe_clip_action(
action: np.ndarray,
low: np.ndarray | None,
high: np.ndarray | None,
) -> np.ndarray:
if low is None or high is None:
return action
low = np.asarray(low, dtype=np.float32).ravel()
high = np.asarray(high, dtype=np.float32).ravel()
if low.shape != action.shape or high.shape != action.shape:
return action
return np.clip(action, low, high)
def _rollout_headless(
*,
env: BrittleStarEnv,
policy: CleanRLPPOPolicy,
policy: PolicyAgent,
seed: int,
max_steps: int,
action_low: np.ndarray | None,
action_high: np.ndarray | None,
action_mask: np.ndarray | None = None,
) -> tuple[float, int, bool, float | None]:
"""Run one rollout up to max_steps.
Returns (return, length, reached_target, final_xy_dist).
Note: In the MJC backend, the raw env reward can be 0.0; we compute a simple
progress reward based on xy_distance_to_target.
"""
state = env.reset(seed=seed)
ep_return = 0.0
@ -302,12 +177,10 @@ def _rollout_one_episode_headless(
obs_dict = observations or {}
action = policy.act(observations=obs_dict)
if action_mask is not None:
action = action[action_mask]
action = _maybe_clip_action(action, action_low, action_high)
nu = int(state.mj_model.nu)
if nu > 0 and action.shape != (nu,):
raise ValueError(f"Policy returned action shape {action.shape}, expected ({nu},)")
state = env.step(state=state, action=action)
steps += 1
@ -325,23 +198,23 @@ def _rollout_one_episode_headless(
return ep_return, steps, reached_target, final_dist
def _run_one_episode_viewer(
def _rollout_viewer(
*,
env: BrittleStarEnv,
policy: CleanRLPPOPolicy,
policy: PolicyAgent,
seed: int,
state: Any,
control_dt: float,
max_steps: int | None,
action_low: np.ndarray | None,
action_high: np.ndarray | None,
action_mask: np.ndarray | None = None,
) -> None:
import mujoco.viewer
model = state.mj_model
data = state.mj_data
_ = int(seed)
episode_return = 0.0
observations = _get_observations(state)
prev_dist = _get_xy_distance_to_target(observations)
@ -359,11 +232,9 @@ def _run_one_episode_viewer(
obs_dict = observations or {}
action = policy.act(observations=obs_dict)
if action_mask is not None:
action = action[action_mask]
action = _maybe_clip_action(action, action_low, action_high)
if model.nu > 0 and action.shape != (int(model.nu),):
raise ValueError(
f"Policy returned action shape {action.shape}, expected ({int(model.nu)},)"
)
# The passive viewer runs a GUI thread; protect MuJoCo state mutation.
with viewer.lock():
@ -398,199 +269,117 @@ def _run_one_episode_viewer(
)
def _infer_checkpoint_obs_dim(policy: CleanRLPPOPolicy) -> int | None:
"""Best-effort read of the first Dense kernel input dim (obs dim)."""
try:
kernel = policy._params["sensor_params"]["params"]["Dense_0"]["kernel"]
return int(getattr(kernel, "shape")[0])
except Exception:
return None
def _load_trained_config(path: Path) -> DictConfig:
"""Load a trained config YAML.
Supports both:
- Hydra's run config (e.g. runs/.../.hydra/config.yaml)
- This project's logger metadata YAMLs, which may contain
``!!python/object/apply:...`` tags for Enums.
For safety, we *do not* execute Python constructors from YAML; we only
treat these tags as data and extract their scalar arguments.
"""
if not path.exists():
raise FileNotFoundError(f"trained_config_path does not exist: '{path}'.")
if not path.is_file():
raise ValueError(f"trained_config_path must be a file, got: '{path}'.")
try:
return OmegaConf.load(path)
except Exception as exc:
python_apply_prefix = "tag:yaml.org,2002:python/object/apply:"
class _SafeLoaderWithPythonApply(yaml.SafeLoader):
pass
def _construct_python_apply(
loader: yaml.SafeLoader,
_tag_suffix: str,
node: yaml.Node,
) -> Any:
if isinstance(node, yaml.SequenceNode):
seq = loader.construct_sequence(node)
if len(seq) == 1:
return seq[0]
return seq
if isinstance(node, yaml.MappingNode):
return loader.construct_mapping(node)
return loader.construct_scalar(node)
_SafeLoaderWithPythonApply.add_multi_constructor(
python_apply_prefix, _construct_python_apply
def _load_metadata_yaml(model_path: Path) -> dict:
"""Discover and load the sidecar metadata YAML file."""
metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml")
if not metadata_path.exists():
raise FileNotFoundError(
f"Could not find metadata YAML for {model_path.name}. Expected it at {metadata_path}"
)
try:
data = yaml.load(path.read_text(encoding="utf-8"), Loader=_SafeLoaderWithPythonApply)
except Exception as yaml_exc:
raise ValueError(
"Failed to load trained_config_path as YAML. "
"If this is a Hydra run, pass the run's '.hydra/config.yaml' file. "
f"Got: '{path}'."
) from yaml_exc
if not isinstance(data, dict):
raise ValueError(
"trained_config_path must contain a YAML mapping (dict-like) at the root. "
f"Got type={type(data).__name__} from '{path}'."
) from exc
# Normalize known enum-like strings to their Enum *names* so OmegaConf's
# structured config merge behaves like the normal Hydra config.
from brittle_star_project.environment.env_types import Task
env_cfg = data.get("environment")
if isinstance(env_cfg, dict) and isinstance(env_cfg.get("task"), str):
task_str = str(env_cfg["task"])
try:
env_cfg["task"] = Task[task_str].name
except Exception:
try:
env_cfg["task"] = Task(task_str).name
except Exception:
pass
return OmegaConf.create(data)
with open(metadata_path, "r") as f:
return yaml.safe_load(f)
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
def main(dict_cfg: DictConfig) -> None:
# Compose against the structured schema first, so missing keys are validated.
cfg = OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
# 1. Hydra composes ONLY SimulationSettings
cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg))
sim_cfg = cfg.simulation
# Optional: override env-defining sections (morphology/arena/environment/architecture)
# using the exact Hydra config that was used for training.
trained_cfg_path = cfg.simulation.trained_config_path
if trained_cfg_path:
overrides_raw = OmegaConf.select(cfg, "hydra.overrides.task") or []
overrides = [str(o) for o in overrides_raw]
trained_cfg_path_abs = Path(hydra.utils.to_absolute_path(trained_cfg_path))
trained_cfg = _load_trained_config(trained_cfg_path_abs)
if "hydra" in trained_cfg:
with open_dict(trained_cfg):
del trained_cfg["hydra"]
with open_dict(cfg):
for key in ("morphology", "arena", "environment", "architecture"):
if key in trained_cfg and not _has_cli_override(overrides, key):
base_node = OmegaConf.select(cfg, key)
override_node = OmegaConf.select(trained_cfg, key)
try:
cfg[key] = OmegaConf.merge(base_node, override_node)
except Exception as exc:
raise ValueError(
"Failed to merge trained config into the active Hydra config. "
f"Key={key!r}, trained_config_path='{trained_cfg_path_abs}'."
) from exc
# Convert DictConfig to structured dataclass.
config: BrittleStarConfig = OmegaConf.to_object(cfg)
backend = Backend.MJC
seed = int(config.experiment.seed)
if getattr(config.architecture, "name", None) != "centralized":
raise ValueError(
"simulate.py currently only supports architecture=centralized. "
f"Got architecture.name={getattr(config.architecture, 'name', None)!r}. "
"(Training supports decentralized, but simulation wiring for it isn't implemented.)"
)
model_path_str = config.simulation.model_path
model_path_str = sim_cfg.model_path
if model_path_str is None:
raise ValueError(
"simulation.model_path must be set to a .flax checkpoint (e.g. final_model.flax)"
)
# Hydra chdir changes CWD; resolve relative paths relative to the invocation.
model_path = Path(hydra.utils.to_absolute_path(model_path_str))
if model_path.suffix != ".flax":
raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.")
# ======= ENVIRONMENT SETUP =======
# 2. Discover + load sidecar metadata YAML
metadata = _load_metadata_yaml(model_path)
# 3. Reconstruct typed configs from metadata
trained_morphology = OmegaConf.to_object(
OmegaConf.merge(OmegaConf.structured(MorphologyConfig), metadata.get("morphology", {}))
)
trained_arena = OmegaConf.to_object(
OmegaConf.merge(OmegaConf.structured(ArenaConfig), metadata.get("arena", {}))
)
env_dict = metadata.get("environment", {})
if isinstance(env_dict.get("task"), str):
from brittle_star_project.environment.env_types import Task
try:
env_dict["task"] = Task[env_dict["task"]].name
except Exception:
try:
env_dict["task"] = Task(env_dict["task"]).name
except Exception:
pass
trained_environment = OmegaConf.to_object(
OmegaConf.merge(OmegaConf.structured(EnvConfig), env_dict)
)
trained_obs_bounds = OmegaConf.to_object(
OmegaConf.merge(
OmegaConf.structured(ObservationBoundsConfig), metadata.get("obs_bounds", {})
)
)
# 4. Determine environment morphology
if sim_cfg.morphology_override is not None:
override_path = Path(hydra.utils.to_absolute_path(sim_cfg.morphology_override))
if not override_path.exists():
raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}")
with open(override_path, "r") as f:
override_dict = yaml.safe_load(f)
env_morphology = OmegaConf.to_object(
OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict)
)
else:
env_morphology = trained_morphology
# 5. Build obs_processor with TRAINING morphology padding masks always
padding_masks = compute_padding_masks(
segments_per_arm=env_morphology.segments_per_arm,
)
obs_processor = create_obs_processor(
bounds_dict=trained_obs_bounds.to_bounds_dict(),
padding_masks=padding_masks,
)
# 6. Build environment
backend = Backend.MJC
seed = int(cfg.experiment.seed)
factory = BrittleStarEnvFactory()
raw_env = factory.create_environment(
backend,
config.morphology,
config.arena,
config.environment,
env_morphology,
trained_arena,
trained_environment,
)
env = BrittleStarEnv(
raw_env,
backend=backend,
config=config.environment,
morphology_config=config.morphology,
config=trained_environment,
morphology_config=env_morphology,
)
state0 = env.reset(seed=seed)
# Match training's padded observation layout for amputated morphologies.
padding_masks = compute_padding_masks(config.morphology.segments_per_arm)
# Calculate the action dimension the model was trained with
trained_action_dim = sum(trained_morphology.segments_per_arm) * 2
training_bounds = ObservationBoundsConfig().to_bounds_dict()
if trained_cfg_path and "obs_bounds" in trained_cfg:
try:
training_bounds = OmegaConf.to_object(trained_cfg.obs_bounds).to_bounds_dict()
except Exception:
pass
else:
training_bounds = config.obs_bounds.to_bounds_dict()
obs_processor = create_obs_processor(
bounds_dict=training_bounds,
padding_masks=padding_masks,
# 7. Load policy
policy = PolicyAgent.load(
model_path, action_dim=trained_action_dim, obs_processor=obs_processor
)
# ======= MODEL SETUP =======
nu = int(state0.mj_model.nu)
policy = CleanRLPPOPolicy.load(model_path, action_dim=nu, obs_processor=obs_processor)
# Helpful early failure when configs don't match the checkpoint.
observations0 = _get_observations(state0)
obs0_dict = observations0 or {}
batched_obs0 = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], obs0_dict)
env_obs_dim = int(obs_processor(batched_obs0).shape[1])
ckpt_obs_dim = _infer_checkpoint_obs_dim(policy)
if ckpt_obs_dim is not None and ckpt_obs_dim != env_obs_dim:
raise ValueError(
"Checkpoint/env mismatch: "
f"checkpoint expects obs_dim={ckpt_obs_dim}, env provides obs_dim={env_obs_dim}. "
"Use the same Hydra config (morphology/arena/environment) "
"that was used during training."
)
# Convert the JAX boolean mask to a numpy array for easy indexing
action_mask = np.asarray(padding_masks["mask_2x"])
# Match training's action clipping behavior.
action_space = getattr(raw_env, "action_space", None)
@ -601,24 +390,26 @@ def main(dict_cfg: DictConfig) -> None:
None if action_space is None else np.asarray(action_space.high, dtype=np.float32).ravel()
)
# ======= SIMULATION =======
headless = bool(config.simulation.headless)
max_steps = config.simulation.max_steps
# 8. Run simulation
headless = bool(sim_cfg.headless)
max_steps = sim_cfg.max_steps
if headless:
if max_steps is None:
raise ValueError("simulation.max_steps is required when simulation.headless=true")
max_steps_i = int(max_steps)
if max_steps_i <= 0:
raise ValueError("simulation.max_steps must be > 0")
ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless(
ep_return, ep_len, reached_target, final_dist = _rollout_headless(
env=env,
policy=policy,
seed=seed,
max_steps=max_steps_i,
action_low=action_low,
action_high=action_high,
action_mask=action_mask,
)
final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}"
print(
@ -627,18 +418,17 @@ def main(dict_cfg: DictConfig) -> None:
f"target_reached={reached_target}, final_xy_dist={final_dist_str}"
)
else:
max_steps_val = None
if max_steps is not None:
max_steps_i = int(max_steps)
if max_steps_i <= 0:
raise ValueError("simulation.max_steps must be > 0")
max_steps_val: int | None = max_steps_i
else:
max_steps_val = None
max_steps_val = max_steps_i
model_dt = float(state0.mj_model.opt.timestep)
control_dt = model_dt * float(config.environment.num_physics_steps_per_control_step)
control_dt = model_dt * float(trained_environment.num_physics_steps_per_control_step)
_run_one_episode_viewer(
_rollout_viewer(
env=env,
policy=policy,
seed=seed,
@ -647,6 +437,7 @@ def main(dict_cfg: DictConfig) -> None:
max_steps=max_steps_val,
action_low=action_low,
action_high=action_high,
action_mask=action_mask,
)
env.close()