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

refactor: evaluation subpackage

Browse source
This commit is contained in:
Tibo De Peuter 2026-04-28 13:39:01 +02:00
parent 8f2a5d25ed
commit 9d1e2c9bff
Signed by: tdpeuter
SSH key fingerprint: SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
8 changed files with 409 additions and 382 deletions
Download patch file Download diff file Expand all files Collapse all files

320
scripts/simulate.py
View file

@ -4,280 +4,29 @@ 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.morphology_override=configs/morphology/3_arms.yaml \
simulation.model_path=runs/.../final_model.flax
"""
from __future__ import annotations
import itertools
import time
from pathlib import Path
from typing import Any
import flax
import hydra
import jax
import jax.numpy as jnp
import numpy as np
import yaml
from omegaconf import DictConfig, OmegaConf
import yaml
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 (
MorphologyConfig,
ArenaConfig,
EnvConfig,
ObservationBoundsConfig,
)
from brittle_star_project.environment.env_config import MorphologyConfig
class PolicyAgent:
"""Wraps a trained Flax actor for deterministic inference."""
def __init__(
self,
*,
sensor_params: Any,
actor_params: Any,
action_dim: int,
obs_processor: Any,
) -> None:
from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
# 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)
self._actor_apply = jax.jit(self._actor.apply)
self._params = {
"sensor_params": sensor_params,
"actor_params": actor_params,
}
self._obs_processor = obs_processor
@staticmethod
def load(
path: Path,
*,
action_dim: int,
obs_processor: Any,
) -> "PolicyAgent":
payload = path.read_bytes()
restored = flax.serialization.msgpack_restore(payload)
sensor_params = None
actor_params = None
# 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,
obs_processor=obs_processor,
)
def act(self, *, observations: dict[str, Any]) -> np.ndarray:
batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations)
obs = self._obs_processor(batched_obs)[0]
hidden = self._sensor_apply(self._params["sensor_params"], obs)
mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
# Always evaluate with the actor mean.
# (Sampling adds exploration noise, which is useful for training but not for evaluation.)
return np.asarray(mean, dtype=np.float32).ravel()
def _get_observations(state: Any) -> dict[str, Any] | None:
return getattr(state, "observations", None)
def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None:
return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0])
def _target_reached(*, state: Any) -> bool:
return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False))
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: 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]:
state = env.reset(seed=seed)
ep_return = 0.0
observations = _get_observations(state)
prev_dist = _get_xy_distance_to_target(observations)
reached_target = _target_reached(state=state)
steps = 0
for _ in range(int(max_steps)):
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)
state = env.step(state=state, action=action)
steps += 1
observations = _get_observations(state)
cur_dist = _get_xy_distance_to_target(observations)
if prev_dist is not None and cur_dist is not None:
ep_return += prev_dist - cur_dist
prev_dist = cur_dist
reached_target = _target_reached(state=state)
if reached_target:
break
final_dist = _get_xy_distance_to_target(observations)
return ep_return, steps, reached_target, final_dist
def _rollout_viewer(
*,
env: BrittleStarEnv,
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
episode_return = 0.0
observations = _get_observations(state)
prev_dist = _get_xy_distance_to_target(observations)
reached_target = _target_reached(state=state)
steps = 0
# Use the viewer as a context manager to avoid GLX teardown races.
with mujoco.viewer.launch_passive(model, data) as viewer:
step_iter = range(int(max_steps)) if max_steps is not None else itertools.count()
for _step_idx in step_iter:
if not viewer.is_running():
break
step_start = time.time()
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)
# The passive viewer runs a GUI thread; protect MuJoCo state mutation.
with viewer.lock():
state = env.step(state=state, action=action)
if not viewer.is_running():
break
viewer.sync()
steps += 1
observations = _get_observations(state)
cur_dist = _get_xy_distance_to_target(observations)
if prev_dist is not None and cur_dist is not None:
episode_return += prev_dist - cur_dist
prev_dist = cur_dist
reached_target = _target_reached(state=state)
if reached_target:
break
remaining = control_dt - (time.time() - step_start)
if remaining > 0:
time.sleep(remaining)
dist = _get_xy_distance_to_target(observations)
dist_str = "n/a" if dist is None else f"{dist:.3f}"
print(
"episode done: "
f"return={episode_return:.6f}, len={steps}, "
f"target_reached={reached_target}, final_xy_dist={dist_str}"
)
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}"
)
with open(metadata_path, "r") as f:
return yaml.safe_load(f)
from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
from brittle_star_project.evaluation.policy import PolicyAgent
from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
@ -297,36 +46,10 @@ def main(dict_cfg: DictConfig) -> None:
raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.")
# 2. Discover + load sidecar metadata YAML
metadata = _load_metadata_yaml(model_path)
metadata = load_metadata(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", {})
)
)
training = metadata_to_configs(metadata)
# 4. Determine environment morphology
if sim_cfg.morphology_override is not None:
@ -339,14 +62,15 @@ def main(dict_cfg: DictConfig) -> None:
OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict)
)
else:
env_morphology = trained_morphology
env_morphology = training.morphology
# 5. Build obs_processor with TRAINING morphology padding masks always
padding_masks = compute_padding_masks(
segments_per_arm=env_morphology.segments_per_arm,
reference_segments_per_arm=training.morphology.segments_per_arm,
)
obs_processor = create_obs_processor(
bounds_dict=trained_obs_bounds.to_bounds_dict(),
bounds_dict=training.obs_bounds.to_bounds_dict(),
padding_masks=padding_masks,
)
@ -358,23 +82,23 @@ def main(dict_cfg: DictConfig) -> None:
raw_env = factory.create_environment(
backend,
env_morphology,
trained_arena,
trained_environment,
training.arena,
training.environment,
)
env = BrittleStarEnv(
raw_env,
backend=backend,
config=trained_environment,
config=training.environment,
morphology_config=env_morphology,
)
state0 = env.reset(seed=seed)
# Calculate the action dimension the model was trained with
trained_action_dim = sum(trained_morphology.segments_per_arm) * 2
trained_action_dim = sum(training.morphology.segments_per_arm) * 2
# 7. Load policy
policy = PolicyAgent.load(
policy = PolicyAgent.from_checkpoint(
model_path, action_dim=trained_action_dim, obs_processor=obs_processor
)
@ -402,7 +126,7 @@ def main(dict_cfg: DictConfig) -> None:
if max_steps_i <= 0:
raise ValueError("simulation.max_steps must be > 0")
ep_return, ep_len, reached_target, final_dist = _rollout_headless(
result = rollout_headless(
env=env,
policy=policy,
seed=seed,
@ -411,11 +135,11 @@ def main(dict_cfg: DictConfig) -> None:
action_high=action_high,
action_mask=action_mask,
)
final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}"
final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
print(
"episode done: "
f"return={ep_return:.6f}, len={ep_len}, "
f"target_reached={reached_target}, final_xy_dist={final_dist_str}"
f"return={result.return_:.6f}, len={result.length}, "
f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}"
)
else:
max_steps_val = None
@ -426,9 +150,9 @@ def main(dict_cfg: DictConfig) -> None:
max_steps_val = max_steps_i
model_dt = float(state0.mj_model.opt.timestep)
control_dt = model_dt * float(trained_environment.num_physics_steps_per_control_step)
control_dt = model_dt * float(training.environment.num_physics_steps_per_control_step)
_rollout_viewer(
rollout_viewer(
env=env,
policy=policy,
seed=seed,