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

Deployed e4869e0 with MkDocs version: 1.6.1

Browse source
This commit is contained in:
github-actions[bot] 2026-05-19 20:58:23 +00:00
parent fd3dbe898a
commit 26e0b9ee28
75 changed files with 13749 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

143
scripts/analysis/explore_tensorboard.py Normal file
View file

@ -0,0 +1,143 @@
#!/usr/bin/env python3
"""
Reproducible CLI tool to explore TensorBoard logs.
Designed for both local development and HPC diagnostics.
Requirements:
pip install tensorboard
Usage:
python explore_tensorboard.py <path_to_run_directory> [--csv output.csv]
"""
import argparse
import os
import sys
import csv
try:
from tensorboard.backend.event_processing import event_accumulator
except ImportError:
print("Error: Missing dependency. Please run: pip install tensorboard")
sys.exit(1)
def explore_run(log_dir):
"""
Extracts and displays a summary of scalar metrics from a TensorBoard log directory.
"""
print(f"\n{'=' * 20} Exploring Run {'=' * 20}")
print(f"Directory: {log_dir}")
print(f"{'=' * 55}\n")
if not os.path.exists(log_dir):
print(f"Error: Directory '{log_dir}' does not exist.")
return None
# Initialize EventAccumulator
# size_guidance=0 loads all data points for each tag.
ea = event_accumulator.EventAccumulator(
log_dir,
size_guidance={
event_accumulator.SCALARS: 0,
event_accumulator.TENSORS: 0,
},
)
print("Loading event files (this may take a moment for large runs)...")
ea.Reload()
tags = ea.Tags()
scalar_tags = tags.get("scalars", [])
if not scalar_tags:
print("No scalar metrics found in this directory.")
return None
print(f"Found {len(scalar_tags)} scalar metrics.\n")
data = {}
summary = []
# Process scalar values
for tag in scalar_tags:
events = ea.Scalars(tag)
if not events:
continue
values = [e.value for e in events]
last_event = events[-1]
data[tag] = values
summary.append(
{
"Metric": tag,
"Steps": len(events),
"Last Value": f"{last_event.value:.4f}",
"Max": f"{max(values):.4f}",
"Min": f"{min(values):.4f}",
}
)
# Display summary table formatted manually
summary = sorted(summary, key=lambda x: x["Metric"])
print(f"{'Metric':<30} {'Steps':>10} {'Last':>12} {'Max':>12} {'Min':>12}")
print("-" * 80)
for row in summary:
print(
f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} "
f"{row['Max']:>12} {row['Min']:>12}"
)
# Calculate and display global metadata
if "charts/SPS" in data:
sps_events = ea.Scalars("charts/SPS")
if len(sps_events) > 1:
total_duration_hours = (sps_events[-1].wall_time - sps_events[0].wall_time) / 3600
print(f"\nTotal Recorded Duration: {total_duration_hours:.2f} hours")
# Estimate completion if total_timesteps is available in hyperparameters
try:
hp_tags = [t for t in tags.get("tensors", []) if "hyperparameters" in t]
if hp_tags:
hp_event = ea.Tensors(hp_tags[0])[0]
hp_text = hp_event.tensor_proto.string_val[0].decode("utf-8")
if "total_timesteps" in hp_text:
for line in hp_text.split("\n"):
if "total_timesteps" in line:
target = int(line.split("|")[2].strip())
current = ea.Scalars(scalar_tags[0])[-1].step
percent = (current / target) * 100
print(f"Progress: {current:,} / {target:,} steps ({percent:.1f}%)")
except Exception:
pass
return data
def main():
parser = argparse.ArgumentParser(description="Reproducible TensorBoard exploration tool.")
parser.add_argument("log_dir", help="Path to the TensorBoard run directory.")
parser.add_argument("--csv", help="Optional: Path to export scalar data to CSV.", default=None)
args = parser.parse_args()
scalar_data = explore_run(args.log_dir)
if args.csv and scalar_data:
# Reloading for wall_time and steps
ea = event_accumulator.EventAccumulator(args.log_dir).Reload()
with open(args.csv, mode="w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=["tag", "step", "value", "wall_time"])
writer.writeheader()
for tag in scalar_data.keys():
for e in ea.Scalars(tag):
writer.writerow(
{"tag": tag, "step": e.step, "value": e.value, "wall_time": e.wall_time}
)
print(f"\nData exported to: {args.csv}")
if __name__ == "__main__":
main()

1244
scripts/analysis/index.html Normal file
View file

File diff suppressed because it is too large Load diff

182
scripts/compare_models.py Normal file
View file

@ -0,0 +1,182 @@
"""Compare multiple trained policies across shared evaluation conditions.
For each model listed in evaluation.comparison_models, this script runs
`comparison_num_episodes` headless rollouts (seeded sequentially from
`comparison_base_seed`) and writes a results CSV to `comparison_output_csv`.
Results include two metrics per episode:
- `eval_return` — shaped reward (same function used during training)
- `max_velocity` — approximated as initial_xy_dist / steps taken
Usage:
# With the default evaluation config
python scripts/compare_models.py evaluation=poster
# Override the output path on the fly
python scripts/compare_models.py evaluation=poster \\
evaluation.comparison_output_csv=metrics/quick_comparison.csv
"""
from __future__ import annotations
import csv
import logging
import time
from pathlib import Path
import hydra
from omegaconf import DictConfig, OmegaConf
from brittle_star_project.configs.main_config import BrittleStarConfig
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.evaluation import build_eval_env
from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
from brittle_star_project.evaluation.rollout import rollout_headless
_FIELDNAMES = [
"model_path",
"architecture",
"arm_0",
"arm_1",
"arm_2",
"arm_3",
"arm_4",
"num_active_arms",
"seed",
"reached_target",
"episode_length",
"eval_return",
"initial_target_distance",
"final_xy_dist",
"approx_max_velocity",
]
def _approx_max_velocity(result) -> float | None:
"""Approximate max velocity as distance covered per step.
This is a rough upper bound: (initial_dist - final_dist) / steps.
"""
if result.initial_target_distance is None or result.final_xy_dist is None or result.length <= 0:
return None
dist_covered = result.initial_target_distance - result.final_xy_dist
return dist_covered / result.length
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
def main(dict_cfg: DictConfig) -> None:
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)
cfg: BrittleStarConfig = OmegaConf.to_object(
OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
)
eval_cfg = cfg.evaluation
model_paths = [str(p) for p in eval_cfg.comparison_models]
if not model_paths:
raise ValueError(
"evaluation.comparison_models is empty. "
"Add at least one model path in your evaluation config."
)
base_seed = int(eval_cfg.comparison_base_seed)
num_episodes = int(eval_cfg.comparison_num_episodes)
max_steps = int(eval_cfg.eval_max_steps)
seeds = list(range(base_seed, base_seed + num_episodes))
output_path = Path(hydra.utils.to_absolute_path(eval_cfg.comparison_output_csv))
output_path.parent.mkdir(parents=True, exist_ok=True)
logger.info(
f"Comparing {len(model_paths)} models over {num_episodes} episodes "
f"(seeds {seeds[0]}–{seeds[-1]})."
)
logger.info(f"Results will be written to: {output_path}")
with open(output_path, "w", newline="") as csv_file:
writer = csv.DictWriter(csv_file, fieldnames=_FIELDNAMES)
writer.writeheader()
for model_path_str in model_paths:
model_path = Path(hydra.utils.to_absolute_path(model_path_str))
logger.info(f"Evaluating model: {model_path.name}")
try:
metadata = load_metadata(model_path)
except FileNotFoundError as e:
logger.warning(f"Skipping model — {e}")
continue
training = metadata_to_configs(metadata)
# Determine morphologies to evaluate
# If comparison_morphologies is empty, use the model's training morphology
morphologies = [None]
if eval_cfg.comparison_morphologies:
morphologies = [
Path(hydra.utils.to_absolute_path(m)) for m in eval_cfg.comparison_morphologies
]
for morph_path in morphologies:
morph_label = morph_path.name if morph_path else "training"
logger.info(f" Morphology: {morph_label}")
bundle = build_eval_env(
model_path=model_path,
training=training,
metadata=metadata,
morphology_override_path=morph_path,
)
for seed in seeds:
t0 = time.time()
result = rollout_headless(
env=bundle.env,
policy=bundle.policy,
seed=seed,
max_steps=max_steps,
action_low=bundle.action_low,
action_high=bundle.action_high,
action_mask=bundle.action_mask,
)
elapsed = time.time() - t0
velocity = _approx_max_velocity(result)
logger.debug(
f" seed={seed:3d} | "
f"reached={str(result.reached_target):<5} | "
f"return={result.return_:+8.3f} | "
f"steps={result.length:4d} | "
f"({elapsed:.1f}s)"
)
row = {
"model_path": model_path_str,
"architecture": bundle.architecture,
"num_active_arms": bundle.num_active_arms,
"seed": seed,
"reached_target": result.reached_target,
"episode_length": result.length,
"eval_return": result.return_,
"initial_target_distance": result.initial_target_distance,
"final_xy_dist": result.final_xy_dist,
"approx_max_velocity": velocity,
}
# Add per-arm segments
for i, segs in enumerate(bundle.segments_per_arm):
row[f"arm_{i}"] = segs
writer.writerow(row)
csv_file.flush()
bundle.env.close()
logger.info(f"Done. Results saved to {output_path}")
if __name__ == "__main__":
register_configs()
main()

264
scripts/evaluate_checkpoints.py Normal file
View file

@ -0,0 +1,264 @@
"""Re-evaluate saved checkpoints from a completed training run using MJX.
This script scans the checkpoint directory of a training run (the `checkpoints/`
folder inside a Hydra output directory), loads each `.flax` checkpoint, runs
one deterministic evaluation episode with `build_eval_rollout_fn`, and appends
the result to the run's `metrics/checkpoint_evaluation.csv`.
It is intended for post-training analysis when per-checkpoint evaluation was not
enabled during training (`evaluate_checkpoints: false`).
Usage:
python scripts/evaluate_checkpoints.py \
simulation.model_path=runs/2024-01-01/12-00-00/final_model.flax \
evaluation.eval_max_steps=5000 \
evaluation.eval_seed=0
The script resolves the run directory from `simulation.model_path`, discovers
all `*.flax` checkpoints under `checkpoints/`, and evaluates them in order.
"""
from __future__ import annotations
from brittle_star_project.MLPs.mlps import (
Actor,
GenericDenseLayersWithActivation,
MessagePasser,
)
from brittle_star_project.MLPs.adjancency_builder import build_adjacency
from brittle_star_project.environment import MorphMode
from brittle_star_project.MLPs.routing import apply_per_node
import logging
import re
from pathlib import Path
import hydra
import jax
import numpy as np
import jax.numpy as jnp
from omegaconf import DictConfig, OmegaConf
from brittle_star_project.configs.main_config import BrittleStarConfig
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
from brittle_star_project.environment.obs_processing import create_obs_processor
from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
from brittle_star_project.evaluation.checkpoint import (
load_metadata,
load_params,
metadata_to_configs,
)
from brittle_star_project.evaluation.evaluate_mjx import (
append_checkpoint_eval_row,
build_eval_rollout_fn,
evaluate_checkpoint_mjx,
)
from brittle_star_project.trainers.PPOTrainer import reward_fn
def _parse_iteration(checkpoint_path: Path) -> int:
"""Parse the iteration number from a checkpoint filename like `checkpoint_0042.flax`."""
match = re.search(r"(\d+)", checkpoint_path.stem)
return int(match.group(1)) if match else -1
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
def main(dict_cfg: DictConfig) -> None:
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)
cfg: BrittleStarConfig = OmegaConf.to_object(
OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
)
sim_cfg = cfg.simulation
eval_cfg = cfg.evaluation
# --- Resolve the model path to find the run directory ---
model_path_str = sim_cfg.model_path
if model_path_str is None:
raise ValueError(
"simulation.model_path must point to the final_model.flax of a training run."
)
model_path = Path(hydra.utils.to_absolute_path(model_path_str))
run_dir = model_path.parent
checkpoints_dir = run_dir / "checkpoints"
if not checkpoints_dir.exists():
raise FileNotFoundError(
f"No checkpoints/ directory found in run directory: {run_dir}\n"
"Make sure simulation.model_path points to a completed training run."
)
checkpoints = sorted(checkpoints_dir.glob("*.flax"), key=_parse_iteration)
if not checkpoints:
raise FileNotFoundError(f"No .flax checkpoints found in {checkpoints_dir}")
logger.info(f"Found {len(checkpoints)} checkpoint(s) in {checkpoints_dir}")
# --- Load sidecar metadata + reconstruct training config ---
metadata_override = (
Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path))
if sim_cfg.metadata_path is not None
else None
)
metadata = load_metadata(model_path, metadata_override)
training = metadata_to_configs(metadata)
padding_masks = compute_padding_masks(
segments_per_arm=training.morphology.segments_per_arm,
reference_segments_per_arm=training.morphology.segments_per_arm,
)
morph_mode = training.morphology.morph_mode
segments_per_arm = jnp.asarray(
training.morphology.segments_per_arm,
dtype=jnp.int32,
)
num_arms = (
jnp.where(
segments_per_arm > 0,
1,
0,
)
.sum()
.item()
)
match morph_mode:
case MorphMode.CENTRALIZED:
needed_copies = 1
agent_indices = [0, 1, 2, 3, 4]
case MorphMode.FULLY_CONNECTED | MorphMode.RING:
agent_mask = segments_per_arm > 0
agent_indices = jnp.where(agent_mask)[0]
needed_copies = num_arms
case MorphMode.SEGMENT:
agent_mask = segments_per_arm > 0
agent_indices = jnp.where(agent_mask)[0]
needed_copies = (segments_per_arm.sum() + num_arms).item()
obs_processor = create_obs_processor(
bounds_dict=training.obs_bounds.to_bounds_dict(),
padding_masks=padding_masks,
num_arms=num_arms,
needed_copies=needed_copies,
morph_mode=morph_mode,
segments_per_arm=segments_per_arm,
agent_indices=agent_indices,
)
env = BrittleStarJaxEnvWrapper(
morphology=training.morphology,
arena=training.arena,
env_config=training.environment,
num_envs=1,
)
action_low = np.asarray(env.single_action_space.low, dtype=np.float32)
action_high = np.asarray(env.single_action_space.high, dtype=np.float32)
sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
actor = Actor(action_dim=env.single_action_space.shape[0])
sensor.apply = jax.jit(sensor.apply)
actor.apply = jax.jit(actor.apply)
eval_fn = build_eval_rollout_fn(
env=env,
obs_processor=obs_processor,
sensor_apply=sensor.apply,
actor_apply=actor.apply,
action_low=action_low,
action_high=action_high,
reward_fn=reward_fn,
)
morph_mode = training.morphology.morph_mode
segments_per_arm = jnp.asarray(
training.morphology.segments_per_arm,
dtype=jnp.int32,
)
match morph_mode:
case MorphMode.CENTRALIZED:
needed_copies = 1
case MorphMode.FULLY_CONNECTED | MorphMode.RING:
needed_copies = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
case MorphMode.SEGMENT:
needed_copies = (
segments_per_arm.sum() + jnp.where(segments_per_arm > 0, 1, 0).sum()
).item()
adj = build_adjacency(
training.morphology.segments_per_arm,
morph_mode,
)
sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
actor = Actor(action_dim=env.single_action_space.shape[0] // needed_copies)
message_passer = (
MessagePasser(
hidden_dim=300,
num_propagation_steps=4,
adj_matrix=adj,
)
if morph_mode != MorphMode.CENTRALIZED
else None
)
eval_fn = build_eval_rollout_fn(
env=env,
obs_processor=obs_processor,
sensor_apply=lambda p, x: apply_per_node(sensor.apply, p, x),
actor_apply=lambda p, x: apply_per_node(actor.apply, p, x),
message_passer_apply=(None if message_passer is None else message_passer.apply),
action_low=action_low,
action_high=action_high,
reward_fn=reward_fn,
)
seed = int(eval_cfg.eval_seed)
max_steps = int(eval_cfg.eval_max_steps)
logger.info(f"Evaluating each checkpoint (seed={seed}, max_steps={max_steps}).")
for checkpoint_path in checkpoints:
iteration = _parse_iteration(checkpoint_path)
try:
params = load_params(checkpoint_path)
except Exception as e:
logger.warning(f"Could not load {checkpoint_path.name}: {e}")
continue
result = evaluate_checkpoint_mjx(eval_fn, params, seed=seed, max_steps=max_steps)
csv_path = append_checkpoint_eval_row(
run_dir,
iteration=iteration,
trained_timesteps=0, # unknown without training logs
result=result,
)
logger.debug(
f"checkpoint={iteration:5d} | "
f"reached={str(result.reached_target):<5} | "
f"return={result.eval_return:+8.3f} | "
f"steps={result.steps:4d} | "
f"final_dist={result.final_xy_dist:.3f}"
)
logger.info(f"Done. CSV at: {csv_path}")
env.close()
if __name__ == "__main__":
register_configs()
main()

83
scripts/hpc/export_requirements.py Normal file
View file

@ -0,0 +1,83 @@
#!/usr/bin/env python3
"""Export HPC pip requirements from pyproject.toml.
This is a LOCAL DEVELOPER UTILITY — run it on your own machine before pushing
code whenever pyproject.toml dependencies change. It reads the modules from
env/hpc/modules.txt and the full dependency list from pyproject.toml, then
writes the remainder to env/hpc/requirements.txt.
"""
from __future__ import annotations
import re
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
def normalise(name: str) -> str:
"""Normalise a PyPI package name for comparison."""
return re.sub(r"[-_.]+", "-", name).lower()
def pkg_name(dep: str) -> str:
"""Extract the bare package name from a PEP 508 dependency string."""
return re.split(r"[\[=><~!;]", dep)[0].strip()
def main() -> None:
import tomllib
modules_path = ROOT / "env" / "hpc" / "modules.txt"
if not modules_path.exists():
print(f"Error: {modules_path} not found.", file=sys.stderr)
sys.exit(1)
# Read normalized module names from base modules only
# Library modules (like PyTorch) are kept in requirements for portability
module_names = [
normalise(line.split()[0].split("/")[0])
for line in modules_path.read_text().splitlines()
if line.strip() and not line.startswith("#")
]
pyproject_path = ROOT / "pyproject.toml"
with pyproject_path.open("rb") as f:
data = tomllib.load(f)
# Collect all dependencies, merging 'cuda' extras into base dependencies
dep_dict: dict[str, str] = {}
for dep in data.get("project", {}).get("dependencies", []):
dep_dict[normalise(pkg_name(dep))] = dep
# Add cuda extras (takes precedence for HPC)
optional_deps = data.get("project", {}).get("optional-dependencies", {})
for group in ["cuda"]:
for dep in optional_deps.get(group, []):
dep_dict[normalise(pkg_name(dep))] = dep
deps = list(dep_dict.values())
final_deps: list[str] = []
print("Checking dependencies against HPC module list...", file=sys.stderr)
for dep in deps:
name = normalise(pkg_name(dep))
# Smart check: if the package name is a substring of any loaded module name
# (e.g. 'torch' in 'pytorch', 'scipy' in 'scipy-bundle')
if any(name in mod for mod in module_names):
print(f" [skip – module provider found] {dep}", file=sys.stderr)
continue
final_deps.append(dep)
print(f" [pip] {dep}", file=sys.stderr)
hpc_dir = ROOT / "env" / "hpc"
output_path = hpc_dir / "requirements.txt"
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text("\n".join(final_deps) + "\n")
print(f"\nWrote {len(final_deps)} requirement(s) to {output_path}", file=sys.stderr)
if __name__ == "__main__":
main()

54
scripts/hpc/install.sh Normal file
View file

@ -0,0 +1,54 @@
#!/bin/bash -l
# scripts/hpc/install.sh
#
# Usage (on any compute node):
# bash scripts/hpc/install.sh
#
# Batch usage:
# qsub scripts/hpc/install.sh
#PBS -N brittlestar-install
#PBS -l walltime=00:15:00
set -euo pipefail
# Preliminary status echo
echo ">>> Starting installation job $PBS_JOBID on $(hostname)..."
if [ -n "$PBS_O_WORKDIR" ]; then
cd "$PBS_O_WORKDIR"
fi
mkdir -p "${PBS_O_WORKDIR}/runs"
# Mirror configs to $VSC_DATA to avoid home quota limits (3GB)
# vsc-venv manages environments relative to the requirements file
PROJ_NAME=$(basename "$PWD")
HPC_CONFIG_DIR="$VSC_DATA/$PROJ_NAME/env/hpc"
mkdir -p "$HPC_CONFIG_DIR"
cp env/hpc/*.txt "$HPC_CONFIG_DIR/"
# Keep caches off $VSC_HOME (quota ~3 GB).
export PIP_CACHE_DIR="$VSC_SCRATCH/.cache/pip"
export UV_CACHE_DIR="$VSC_SCRATCH/.cache/uv"
mkdir -p "$PIP_CACHE_DIR" "$UV_CACHE_DIR"
module load vsc-venv
echo ">>> Synchronizing and activating environment (vsc-venv)..."
# cd to $VSC_DATA so vsc-venv creates its venvs/ directory there, not in $HOME.
mkdir -p "$VSC_DATA/$PROJ_NAME"
cd "$VSC_DATA/$PROJ_NAME"
set +euo pipefail
source vsc-venv --activate \
--modules "$HPC_CONFIG_DIR/modules.txt" \
--requirements "$HPC_CONFIG_DIR/requirements.txt"
set -euo pipefail
cd "$PBS_O_WORKDIR"
echo '>>> Installing ipykernel...'
CLUSTER_ID="${VSC_INSTITUTE_CLUSTER:-generic}"
python -m ipykernel install --user --name="sel3_${CLUSTER_ID}" \
--display-name "SEL3 (${CLUSTER_ID})"
echo '>>> Done'

75
scripts/hpc/train.pbs Normal file
View file

@ -0,0 +1,75 @@
# Production training (requires GPU at runtime):
# qsub -l gpus=1 scripts/hpc/train.pbs
# Debug/CPU training:
# qsub scripts/hpc/train.pbs
#PBS -N brittlestar-ppo
#PBS -l nodes=1:ppn=8
#PBS -l walltime=24:00:00
#PBS -o runs/brittlestar-ppo.o$PBS_JOBID
#PBS -e runs/brittlestar-ppo.e$PBS_JOBID
set -euo pipefail
# Preliminary status echo
echo ">>> Starting training job $PBS_JOBID on $(hostname)..."
if [ -n "$PBS_O_WORKDIR" ]; then
cd "$PBS_O_WORKDIR"
fi
# Set up storage paths dynamically
PROJ_NAME=$(basename "$PWD")
RUN_ID="brittlestar_${PBS_JOBID}"
SCRATCH_RUNDIR="$VSC_SCRATCH/runs/$RUN_ID"
DATA_RUNDIR="$VSC_DATA/runs/$RUN_ID"
mkdir -p "$SCRATCH_RUNDIR" "$DATA_RUNDIR" runs/
# Keep caches off $VSC_HOME (quota ~3 GB).
export PIP_CACHE_DIR="$VSC_SCRATCH/.cache/pip"
export UV_CACHE_DIR="$VSC_SCRATCH/.cache/uv"
mkdir -p "$PIP_CACHE_DIR" "$UV_CACHE_DIR"
module load vsc-venv
echo ">>> Synchronizing and activating environment (vsc-venv)..."
HPC_CONFIG_DIR="$VSC_DATA/$PROJ_NAME/env/hpc"
if [ ! -d "$HPC_CONFIG_DIR" ]; then
echo "ERROR: HPC_CONFIG_DIR ($HPC_CONFIG_DIR) does not exist. Run install.sh first."
exit 1
fi
# cd to $VSC_DATA so vsc-venv finds its venvs/ directory there, not in $HOME.
cd "$VSC_DATA/$PROJ_NAME"
set +euo pipefail
source vsc-venv --activate \
--modules "$HPC_CONFIG_DIR/modules.txt" \
--requirements "$HPC_CONFIG_DIR/requirements.txt"
set -euo pipefail
cd "$PBS_O_WORKDIR"
echo ">>> Starting BrittleStar training..."
export MUJOCO_GL=egl
export WANDB_DIR="$SCRATCH_RUNDIR"
export PYTHONPATH="$PBS_O_WORKDIR/src:${PYTHONPATH:-}"
if [ -f "$VSC_DATA/$PROJ_NAME/.env" ]; then
echo ">>> Sourcing API keys from .env..."
export $(grep -v '^#' "$VSC_DATA/$PROJ_NAME/.env" | xargs)
elif [ -f "$PBS_O_WORKDIR/.env" ]; then
echo ">>> Sourcing API keys from .env..."
export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs)
fi
# Run training using Hydra overrides
python scripts/train.py \
hydra.run.dir="$SCRATCH_RUNDIR" \
ppo=stable \
logging=hpc
echo ">>> Staging out results to $DATA_RUNDIR..."
cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"
echo ">>> Done"

445
scripts/plots/analyze_comparisons.py Normal file
View file

@ -0,0 +1,445 @@
"""
Poster Comparison Visualizations
This script generates a Forward Velocity plot and three secondary plots (Accumulated Reward, Success
Rate, Distance Remaining).
"""
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from plot_config import (
COLORS,
apply_style,
BEST_PERFORMER_MARKER,
BEST_PERFORMER_TEXT,
BEST_PERFORMER_COLOR,
create_common_parser,
LEGEND_KWARGS,
)
def load_and_preprocess_data(filepath):
"""Loads CSV and prepares the metrics for plotting."""
df = pd.read_csv(filepath)
# Ensure success rate can be averaged numerically
if "reached_target" in df.columns:
df["reached_target"] = df["reached_target"].astype(int)
return df
def _add_square_placeholders(ax, x_positions, labels):
"""Adds square placeholders for images below the x-axis."""
for x, label in zip(x_positions, labels):
# Create a roughly square rectangle in a mix of data/axes coords
# Shifted down to avoid overlapping with x-tick labels
rect = plt.Rectangle(
(x - 0.25, -0.40),
0.5,
0.18,
transform=ax.get_xaxis_transform(),
facecolor="#F0F0F0",
edgecolor="#A9A9A9",
linestyle="--",
zorder=1,
clip_on=False,
)
ax.add_patch(rect)
ax.text(
x,
-0.31,
f"[ Insert {label}\nImage ]",
transform=ax.get_xaxis_transform(),
ha="center",
va="center",
fontsize=10,
color="#888888",
zorder=2,
)
def plot_grouped_bar(
df,
metric_col,
ylabel,
title,
output_filename,
output_dir,
higher_is_better=True,
show_titles=False,
figsize=(12, 8),
):
"""Generates and saves a highly customized grouped bar chart (grouped by Morphology)."""
grouped = (
df.groupby(["num_active_arms", "architecture"])[metric_col]
.agg(["mean", "std"])
.reset_index()
)
morphologies = sorted(grouped["num_active_arms"].unique(), reverse=True)
architectures = grouped["architecture"].unique()
fig, ax = plt.subplots(figsize=figsize)
bar_width = 0.35
x_indices = np.arange(len(morphologies))
all_bars = {}
all_means = []
for i, arch in enumerate(architectures):
arch_data = grouped[grouped["architecture"] == arch]
means = [
arch_data[arch_data["num_active_arms"] == m]["mean"].values[0]
if not arch_data[arch_data["num_active_arms"] == m].empty
else 0
for m in morphologies
]
stds = [
arch_data[arch_data["num_active_arms"] == m]["std"].values[0]
if not arch_data[arch_data["num_active_arms"] == m].empty
else 0
for m in morphologies
]
all_means.extend(means)
x_pos = x_indices + (i * bar_width) - (bar_width / 2 if len(architectures) == 2 else 0)
color = COLORS.get(arch, "#888888")
clean_label = arch.replace("_", " ").title()
bars = ax.bar(
x_pos,
means,
bar_width,
yerr=stds,
label=clean_label,
color=color,
capsize=8,
error_kw={"elinewidth": 2, "alpha": 0.7},
)
all_bars[arch] = (x_pos, means, stds, bars)
for m_idx, m in enumerate(morphologies):
m_means = {arch: all_bars[arch][1][m_idx] for arch in architectures}
best_arch = (
max(m_means, key=m_means.get) if higher_is_better else min(m_means, key=m_means.get)
)
best_x = all_bars[best_arch][0][m_idx]
best_y = all_bars[best_arch][1][m_idx]
best_std = all_bars[best_arch][2][m_idx]
offset = best_std + (abs(max(m_means.values())) * 0.05) if m_means.values() else 0
ax.text(
best_x,
best_y + offset,
BEST_PERFORMER_TEXT,
ha="center",
va="bottom",
fontsize=28,
color=BEST_PERFORMER_COLOR,
)
# Aesthetics
ax.set_ylabel(ylabel, labelpad=15)
if show_titles:
ax.set_title(title, pad=25, fontweight="bold")
x_ticks_pos = (
x_indices
+ (bar_width / 2 if len(architectures) % 2 == 0 else 0)
- (bar_width / 2 if len(architectures) == 2 else 0)
)
ax.set_xticks(x_ticks_pos)
ax.set_xticklabels([f"{m} Arms" for m in morphologies])
ax.tick_params(axis="x", pad=25) # More padding for the squares
# X-axis at zero
ax.axhline(0, color="black", linewidth=1.5)
ax.spines["bottom"].set_visible(False)
# Y-axis limits explicitly including 0
if all_means:
min_val = min([*all_means, 0])
max_val = max([*all_means, 0])
margin = (max_val - min_val) * 0.15 if max_val != min_val else 0.1
ax.set_ylim(min_val - margin, max_val + margin * 1.5) # Extra top margin for stars
# Format y-ticks to not have excessive decimals, include 0
ticks = (
[min_val, max_val]
if min_val == 0 and max_val == 0
else sorted(list(set([min_val, 0, max_val])))
)
ax.set_yticks(ticks)
ax.yaxis.set_major_formatter(
plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}")
)
_add_square_placeholders(ax, x_ticks_pos, [f"{m} Arms" for m in morphologies])
# Add custom legend entry for best performer
ax.plot(
[],
[],
marker=BEST_PERFORMER_MARKER,
color="w",
markerfacecolor=BEST_PERFORMER_COLOR,
markersize=15,
label="Best Performance",
ls="",
)
ax.legend(**LEGEND_KWARGS, ncol=len(architectures) + 1)
ax.set_facecolor("white")
fig.patch.set_facecolor("white")
os.makedirs(output_dir, exist_ok=True)
base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
plt.close()
def plot_grouped_bar_alt(
df,
metric_col,
ylabel,
title,
output_filename,
output_dir,
higher_is_better=True,
show_titles=False,
figsize=(12, 8),
):
"""Generates and saves a highly customized grouped bar chart (grouped by Architecture)."""
grouped = (
df.groupby(["architecture", "num_active_arms"])[metric_col]
.agg(["mean", "std"])
.reset_index()
)
architectures = sorted(grouped["architecture"].unique())
morphologies = sorted(grouped["num_active_arms"].unique(), reverse=True)
fig, ax = plt.subplots(figsize=figsize)
bar_width = 0.8 / len(morphologies)
x_indices = np.arange(len(architectures))
all_bars = {}
all_means = []
for i, m in enumerate(morphologies):
m_data = grouped[grouped["num_active_arms"] == m]
means = [
m_data[m_data["architecture"] == arch]["mean"].values[0]
if not m_data[m_data["architecture"] == arch].empty
else 0
for arch in architectures
]
stds = [
m_data[m_data["architecture"] == arch]["std"].values[0]
if not m_data[m_data["architecture"] == arch].empty
else 0
for arch in architectures
]
all_means.extend(means)
# Offset bars based on morphology index
offset = (i - len(morphologies) / 2 + 0.5) * bar_width
x_pos = x_indices + offset
# We can use a color gradient or different colors for morphologies
# For simplicity, using a colormap
color = plt.cm.viridis(i / max(1, len(morphologies) - 1))
bars = ax.bar(
x_pos,
means,
bar_width,
yerr=stds,
label=f"{m} Arms",
color=color,
capsize=4,
error_kw={"elinewidth": 1.5, "alpha": 0.7},
)
all_bars[m] = (x_pos, means, stds, bars)
for a_idx, arch in enumerate(architectures):
a_means = {m: all_bars[m][1][a_idx] for m in morphologies}
best_m = (
max(a_means, key=a_means.get) if higher_is_better else min(a_means, key=a_means.get)
)
best_x = all_bars[best_m][0][a_idx]
best_y = all_bars[best_m][1][a_idx]
best_std = all_bars[best_m][2][a_idx]
offset = best_std + (abs(max(a_means.values())) * 0.05) if a_means.values() else 0
ax.text(
best_x,
best_y + offset,
BEST_PERFORMER_TEXT,
ha="center",
va="bottom",
fontsize=20,
color=BEST_PERFORMER_COLOR,
)
# Aesthetics
ax.set_ylabel(ylabel, labelpad=15)
if show_titles:
ax.set_title(title + " (Alt)", pad=25, fontweight="bold")
ax.set_xticks(x_indices)
ax.set_xticklabels([arch.replace("_", " ").title() for arch in architectures])
ax.tick_params(axis="x", pad=25)
# X-axis at zero
ax.axhline(0, color="black", linewidth=1.5)
ax.spines["bottom"].set_visible(False)
if all_means:
min_val = min([*all_means, 0])
max_val = max([*all_means, 0])
margin = (max_val - min_val) * 0.15 if max_val != min_val else 0.1
ax.set_ylim(min_val - margin, max_val + margin * 1.5)
ticks = (
[min_val, max_val]
if min_val == 0 and max_val == 0
else sorted(list(set([min_val, 0, max_val])))
)
ax.set_yticks(ticks)
ax.yaxis.set_major_formatter(
plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}")
)
# In this alt plot, placeholders might be per architecture
_add_square_placeholders(
ax, x_indices, [arch.replace("_", "\n").title() for arch in architectures]
)
ax.plot(
[],
[],
marker=BEST_PERFORMER_MARKER,
color="w",
markerfacecolor=BEST_PERFORMER_COLOR,
markersize=15,
label="Best Performance",
ls="",
)
ax.legend(**LEGEND_KWARGS, ncol=len(morphologies) + 1)
ax.set_facecolor("white")
fig.patch.set_facecolor("white")
os.makedirs(output_dir, exist_ok=True)
base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
plt.close()
if __name__ == "__main__":
parser = create_common_parser(description="Generate comparison poster plots.")
parser.add_argument(
"input_csv", help="Path to the input CSV file containing evaluation results."
)
args = parser.parse_args()
INPUT_CSV = args.input_csv
OUTPUT_DIR = args.output_dir
if not os.path.exists(INPUT_CSV):
print(f"Error: Could not find {INPUT_CSV}. Please ensure the file exists.")
else:
df = load_and_preprocess_data(INPUT_CSV)
print("Data loaded successfully. Generating poster plots...")
apply_style(font_size=args.font_size)
kwargs = {"show_titles": args.show_titles, "figsize": (args.fig_width, args.fig_height)}
# Velocity Conversion: m/s to cm/s
if "approx_max_velocity" in df.columns:
df["approx_max_velocity"] = df["approx_max_velocity"] * 100
# 1. Primary Plot: Forward Velocity
plot_grouped_bar(
df=df,
metric_col="approx_max_velocity",
ylabel="Max Forward Velocity (cm/s)",
title="Graceful Degradation: Velocity Across Morphologies",
output_filename="poster_plot_velocity.png",
output_dir=OUTPUT_DIR,
higher_is_better=True,
**kwargs,
)
plot_grouped_bar_alt(
df=df,
metric_col="approx_max_velocity",
ylabel="Max Forward Velocity (cm/s)",
title="Graceful Degradation: Velocity Across Morphologies",
output_filename="poster_plot_velocity_alt.png",
output_dir=OUTPUT_DIR,
higher_is_better=True,
**kwargs,
)
# 2. Secondary Plot: Accumulated Reward
plot_grouped_bar(
df=df,
metric_col="eval_return",
ylabel="Mean Cumulative Reward",
title="Overall Efficiency Across Morphologies",
output_filename="poster_plot_reward.png",
output_dir=OUTPUT_DIR,
higher_is_better=True,
**kwargs,
)
plot_grouped_bar_alt(
df=df,
metric_col="eval_return",
ylabel="Mean Cumulative Reward",
title="Overall Efficiency Across Morphologies",
output_filename="poster_plot_reward_alt.png",
output_dir=OUTPUT_DIR,
higher_is_better=True,
**kwargs,
)
# 3. Secondary Plot: Success Rate
plot_grouped_bar(
df=df,
metric_col="reached_target",
ylabel="Success Rate (%)",
title="Target Acquisition Consistency",
output_filename="poster_plot_success_rate.png",
output_dir=OUTPUT_DIR,
higher_is_better=True,
**kwargs,
)
plot_grouped_bar_alt(
df=df,
metric_col="reached_target",
ylabel="Success Rate (%)",
title="Target Acquisition Consistency",
output_filename="poster_plot_success_rate_alt.png",
output_dir=OUTPUT_DIR,
higher_is_better=True,
**kwargs,
)
# 4. Secondary Plot: Final Distance Remaining
plot_grouped_bar(
df=df,
metric_col="final_xy_dist",
ylabel="Distance to Target Remaining",
title="Navigational Accuracy (Lower is Better)",
output_filename="poster_plot_distance.png",
output_dir=OUTPUT_DIR,
higher_is_better=False, # For distance, a lower score is better
**kwargs,
)
plot_grouped_bar_alt(
df=df,
metric_col="final_xy_dist",
ylabel="Distance to Target Remaining",
title="Navigational Accuracy (Lower is Better)",
output_filename="poster_plot_distance_alt.png",
output_dir=OUTPUT_DIR,
higher_is_better=False,
**kwargs,
)
print(f"All plots generated in the '{OUTPUT_DIR}/' directory.")

335
scripts/plots/analyze_convergence.py Normal file
View file

@ -0,0 +1,335 @@
"""
Convergence Analysis Script for Poster Visualizations
This script analyzes evaluation metrics from multiple training runs to determine
the convergence point of different reinforcement learning architectures.
Workflow:
1. Loads evaluation data from the CSV files defined in FILE_MAPPING.
2. Calculates a rolling average of the reward and velocity to smooth noise.
3. Determines the convergence timestep for each metric (first time 95% of peak is reached).
4. Generates a grouped bar chart comparing convergence speed and line plots of the raw curves.
Usage:
uv run python scripts/analysis/analyze_convergence.py
Note: For these metrics to be valid, the evaluation CSVs must be generated with
exploration noise strictly disabled (e.g., taking the mean of the action distribution).
"""
import logging
import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from enum import Enum
from plot_config import COLORS, apply_style, create_common_parser, LEGEND_KWARGS
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)
# --- Globals & Configuration ---
USING_DUMMY_DATA = False
SMOOTHING_WINDOW = 3
CONVERGENCE_THRESHOLD = 0.95
class Columns(str, Enum):
# ... (rest of the file remains same, just need to update plotting functions and obtain_data)
"""Column names expected in every evaluation CSV."""
ARCH = "architecture"
TIMESTEPS = "total_trained_timesteps"
REWARD = "accumulated_reward"
VELOCITY = "velocity"
# Maps architecture display names to the path of their evaluation CSV.
# Update these paths once real evaluation data is available.
FILE_MAPPING: dict[str, str] = {
"centralized 2 arms": "runs/dummy/dummy_centralized_2_arms.csv",
"centralized 5 arms": "runs/dummy/dummy_centralized_5_arms.csv",
"decentralized fully connected": "runs/dummy/dummy_decentralized_fully_connected.csv",
"decentralized ring-level": "runs/dummy/dummy_decentralized_ring-level.csv",
"decentralized segment-level": "runs/dummy/dummy_decentralized_segment-level.csv",
}
# Architecture profiles for dummy data generation: (max_reward, max_velocity, sigmoid_speed)
_DUMMY_PROFILES: dict[str, tuple[float, float, float]] = {
"centralized 2 arms": (300, 0.8, 1.2),
"centralized 5 arms": (450, 1.1, 1.0),
"decentralized fully connected": (500, 1.3, 0.7),
"decentralized ring-level": (480, 1.2, 0.8),
"decentralized segment-level": (520, 1.4, 0.6),
}
def generate_dummy_csvs(file_mapping: dict[str, str]):
"""
Generates one dummy CSV per architecture in FILE_MAPPING at their expected locations.
Skips any architecture without a defined profile.
"""
checkpoints = list(range(100, 1100, 100))
timesteps = [cp * 10_000 for cp in checkpoints]
for arch, path in file_mapping.items():
if arch not in _DUMMY_PROFILES:
logger.warning(f"No dummy profile for '{arch}'. Skipping.")
continue
m_reward, m_vel, speed = _DUMMY_PROFILES[arch]
rows = []
for i, ts in enumerate(timesteps):
progress = 1 / (1 + np.exp(-speed * (i - 4)))
rows.append(
{
Columns.TIMESTEPS: ts,
Columns.REWARD: m_reward * progress + np.random.normal(0, 5),
Columns.VELOCITY: m_vel * progress + np.random.normal(0, 0.02),
}
)
# Create parent directories if they don't exist
os.makedirs(os.path.dirname(path), exist_ok=True)
pd.DataFrame(rows).to_csv(path, index=False)
logger.info(f"Generated dummy CSV at expected path: {path}")
def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame:
"""
Loads one CSV per architecture, injects the architecture name as a column,
and returns the combined DataFrame with only the required columns.
"""
required = [Columns.TIMESTEPS, Columns.REWARD, Columns.VELOCITY]
dfs = []
for arch_name, filepath in file_mapping.items():
if not os.path.exists(filepath):
logger.warning(f"File not found: '{filepath}'. Skipping.")
continue
df = pd.read_csv(filepath)
missing = [c for c in required if c not in df.columns]
if missing:
logger.warning(f"Missing columns {missing} in '{filepath}'. Skipping.")
continue
df = df[required].copy()
df[Columns.ARCH] = arch_name
dfs.append(df)
return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame()
def _convergence_timestep(series: pd.Series, timesteps: pd.Series) -> float:
"""Returns the first timestep where the smoothed series reaches 95% of its peak."""
smoothed = series.rolling(window=SMOOTHING_WINDOW, min_periods=1).mean()
threshold = smoothed.max() * CONVERGENCE_THRESHOLD
return timesteps[smoothed >= threshold].iloc[0]
def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame:
"""
For each architecture, determines the convergence timestep based on both
reward and velocity, returning one summary row per architecture.
"""
results = []
for arch in df[Columns.ARCH].unique():
arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
results.append(
{
"Architecture": arch,
"Reward_Convergence_Timestep": _convergence_timestep(
arch_data[Columns.REWARD], arch_data[Columns.TIMESTEPS]
),
"Velocity_Convergence_Timestep": _convergence_timestep(
arch_data[Columns.VELOCITY], arch_data[Columns.TIMESTEPS]
),
}
)
return pd.DataFrame(results)
def _add_bar_labels(bars, max_val: float):
"""Annotates each bar with its value in white bold text, positioned inside."""
for bar in bars:
width = bar.get_width()
label = f"{width / 1e6:.1f}M" if width >= 1e6 else f"{width:,.0f}"
plt.text(
width - (max_val * 0.02),
bar.get_y() + bar.get_height() / 2,
label,
ha="right",
va="center",
fontsize=11,
color="white",
fontweight="bold",
)
def plot_grouped_convergence_chart(
results_df: pd.DataFrame, output_filename: str, output_dir: str, **kwargs
):
"""
Saves a grouped horizontal bar chart comparing Reward and Velocity convergence timesteps
across all architectures.
"""
sorted_df = results_df.sort_values("Reward_Convergence_Timestep", ascending=True)
architectures = sorted_df["Architecture"].tolist()
y_pos = np.arange(len(architectures))
bar_height = 0.35
max_val = sorted_df[
["Reward_Convergence_Timestep", "Velocity_Convergence_Timestep"]
].values.max()
fig, ax = plt.subplots(figsize=kwargs.get("figsize", (12, 8)))
bars_reward = ax.barh(
y_pos + bar_height / 2,
sorted_df["Reward_Convergence_Timestep"],
height=bar_height,
label="Reward Convergence",
color="#1f77b4",
)
bars_velocity = ax.barh(
y_pos - bar_height / 2,
sorted_df["Velocity_Convergence_Timestep"],
height=bar_height,
label="Velocity Convergence",
color="#ff7f0e",
)
title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else ""
if kwargs.get("show_titles", True):
ax.set_title(
f"Comparison of Training Convergence Timesteps{title_suffix}", fontsize=20, pad=20
)
ax.set_xlabel("Timesteps to Convergence (95% of peak)", fontsize=16)
ax.set_ylabel("Architecture", fontsize=16)
ax.set_yticks(y_pos)
ax.set_yticklabels(architectures, fontsize=14)
ax.tick_params(axis="x", labelsize=14)
ax.legend(**LEGEND_KWARGS, ncol=2)
ax.set_xlim(left=0)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
_add_bar_labels(bars_reward, max_val)
_add_bar_labels(bars_velocity, max_val)
plt.tight_layout()
os.makedirs(output_dir, exist_ok=True)
base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
plt.close()
def plot_metric_curves(
df: pd.DataFrame, metric_col: str, title: str, output_filename: str, output_dir: str, **kwargs
):
"""
Saves a line plot of the given metric over training timesteps for every architecture.
"""
fig, ax = plt.subplots(figsize=kwargs.get("figsize", (12, 7)))
for arch in df[Columns.ARCH].unique():
arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
color_key = arch.split()[0].upper() if isinstance(arch, str) else "UNKNOWN"
color = COLORS.get(color_key, "#888888")
ax.plot(
arch_data[Columns.TIMESTEPS],
arch_data[metric_col],
label=arch,
marker="o",
markersize=4,
alpha=0.8,
color=color,
)
title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else ""
if kwargs.get("show_titles", True):
ax.set_title(f"{title}{title_suffix}", fontsize=18, pad=20)
ax.set_xlabel("Training Timesteps", fontsize=14)
ax.set_ylabel(metric_col.replace("_", " ").title(), fontsize=14)
ax.legend(**LEGEND_KWARGS, ncol=len(df[Columns.ARCH].unique()))
ax.grid(True, linestyle="--", alpha=0.6)
ax.set_xlim(left=0)
ax.set_ylim(bottom=0)
plt.tight_layout()
os.makedirs(output_dir, exist_ok=True)
base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
plt.close()
def plot_results(df: pd.DataFrame, results: pd.DataFrame, output_dir: str, **kwargs):
"""Generates and saves all analysis plots."""
plot_grouped_convergence_chart(
results, output_filename="convergence_comparison.png", output_dir=output_dir, **kwargs
)
plot_metric_curves(
df,
Columns.REWARD,
"Training Progress: Accumulated Reward",
"progress_reward_curves.png",
output_dir=output_dir,
**kwargs,
)
plot_metric_curves(
df,
Columns.VELOCITY,
"Training Progress: Velocity",
"progress_velocity_curves.png",
output_dir=output_dir,
**kwargs,
)
def obtain_data() -> pd.DataFrame:
"""Resolves the file mapping, falling back to generated dummy CSVs if needed."""
global USING_DUMMY_DATA
if not any(os.path.exists(p) for p in FILE_MAPPING.values()):
logger.info("No real evaluation files found. Generating dummy CSVs at expected locations.")
generate_dummy_csvs(FILE_MAPPING)
USING_DUMMY_DATA = True
return load_metrics(FILE_MAPPING)
def run_analysis(output_dir: str, **kwargs):
"""Orchestrates data loading, convergence analysis, and plot generation."""
df = obtain_data()
if df.empty:
logger.error("No data found to analyze.")
return
results = analyze_convergence(df)
plot_results(df, results, output_dir, **kwargs)
logger.info("Analysis complete. Plots saved to disk.")
if __name__ == "__main__":
parser = create_common_parser(description="Analyze training convergence.")
args = parser.parse_args()
apply_style(font_size=args.font_size)
run_analysis(
output_dir=args.output_dir,
show_titles=args.show_titles,
figsize=(args.fig_width, args.fig_height),
)

77
scripts/plots/plot_config.py Normal file
View file

@ -0,0 +1,77 @@
import argparse
import matplotlib.pyplot as plt
# Shared Color Palette (Colorblind friendly, high contrast)
# Matches poster design
COLORS = {
"CENTRALIZED": "#2B4162", # Deep Slate Blue
"FULLY_CONNECTED": "#FA9F42", # Vibrant Orange
"RING_LEVEL": "#4E937A", # Muted Teal
"SEGMENT_LEVEL": "#B4436C", # Soft Red
"DECENTRALIZED": "#4E937A", # Default decentralized fallback
}
def apply_style(font_size=28):
"""
Applies the shared typography and aesthetic settings to Matplotlib.
"""
plt.rcParams.update(
{
"font.size": font_size,
"axes.labelsize": font_size + 4,
"axes.titlesize": font_size + 8,
"xtick.labelsize": font_size - 4,
"ytick.labelsize": font_size - 4,
"legend.fontsize": font_size - 6,
"axes.linewidth": 2,
"axes.spines.top": False,
"axes.spines.right": False,
"axes.spines.left": False,
"figure.facecolor": "white",
"axes.facecolor": "white",
"savefig.bbox": "tight",
"savefig.dpi": 300,
}
)
# Star marker for best performer
BEST_PERFORMER_TEXT = "★"
BEST_PERFORMER_MARKER = "*"
BEST_PERFORMER_COLOR = "#D4AF37" # Gold
# Centralized Legend Configuration
LEGEND_KWARGS = {
"loc": "upper center",
"bbox_to_anchor": (0.5, -0.5),
"frameon": False,
}
def create_common_parser(description: str) -> argparse.ArgumentParser:
"""
Creates an argparse parser with common plotting arguments.
"""
parser = argparse.ArgumentParser(description=description)
parser.add_argument(
"--output_dir",
"-o",
default="runs/evaluation/plots",
help="Directory to save the generated plots.",
)
parser.add_argument(
"--show_titles",
action="store_true",
help="Include titles in the plots. Default is False for easier poster integration.",
)
parser.add_argument(
"--font_size", type=int, default=28, help="Base font size in points. Default is 28."
)
parser.add_argument(
"--fig_width", type=float, default=12.0, help="Figure width in inches. Default is 12.0."
)
parser.add_argument(
"--fig_height", type=float, default=8.0, help="Figure height in inches. Default is 8.0."
)
return parser

177
scripts/simulate.py Normal file
View file

@ -0,0 +1,177 @@
"""Simulate a trained policy in the MuJoCo viewer.
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=configs/morphology/3_arms.yaml \
simulation.model_path=runs/.../final_model.flax
"""
from __future__ import annotations
from pathlib import Path
import hydra
from omegaconf import DictConfig, OmegaConf
from brittle_star_project.configs.main_config import BrittleStarConfig
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
from brittle_star_project.evaluation.eval_env_builder import build_eval_env
from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer
from brittle_star_project.evaluation.video import (
record_episode,
create_evaluation_dir,
save_evaluation_metadata,
)
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
def main(dict_cfg: DictConfig) -> None:
# 1. Hydra composes ONLY SimulationSettings
cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg))
sim_cfg = cfg.simulation
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)"
)
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}'.")
# 2. Discover + load sidecar metadata YAML
metadata_override = None
if sim_cfg.metadata_path is not None:
metadata_override = Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path))
metadata = load_metadata(model_path, metadata_override)
# 3. Reconstruct typed configs from metadata
training = metadata_to_configs(metadata)
seed = int(cfg.experiment.seed)
# 4-7. Build evaluation environment and policy
override_path = None
if sim_cfg.morphology_override is not None:
override_path = Path(hydra.utils.to_absolute_path(sim_cfg.morphology_override))
bundle = build_eval_env(
model_path=model_path,
training=training,
metadata=metadata,
morphology_override_path=override_path,
)
env = bundle.env
policy = bundle.policy
action_low = bundle.action_low
action_high = bundle.action_high
action_mask = bundle.action_mask
state0 = env.reset(seed=seed)
# 8. Run simulation
headless = bool(sim_cfg.headless)
max_steps = sim_cfg.max_steps
if sim_cfg.record_video:
if max_steps is None:
raise ValueError("simulation.max_steps is required when simulation.record_video=true")
max_steps_i = int(max_steps)
if max_steps_i <= 0:
raise ValueError("simulation.max_steps must be > 0")
if sim_cfg.video_output_path is None:
eval_dir = create_evaluation_dir(model_path)
output_path = eval_dir / "simulation.mp4"
else:
output_path = Path(hydra.utils.to_absolute_path(sim_cfg.video_output_path))
eval_dir = output_path.parent
eval_dir.mkdir(parents=True, exist_ok=True)
result = record_episode(
env=env,
policy=policy,
seed=seed,
max_steps=max_steps_i,
action_low=action_low,
action_high=action_high,
action_mask=action_mask,
output_path=output_path,
camera_id=sim_cfg.camera_id,
)
save_evaluation_metadata(
eval_dir=eval_dir,
morphology_override_path=sim_cfg.morphology_override,
seed=seed,
max_steps=max_steps_i,
result=result,
)
final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
print(f"Video saved to {output_path}")
print(
"episode done: "
f"return={result.return_:.6f}, len={result.length}, "
f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}"
)
elif 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")
result = 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 result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
print(
"episode done: "
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
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 = max_steps_i
model_dt = float(state0.mj_model.opt.timestep)
control_dt = model_dt * float(training.environment.num_physics_steps_per_control_step)
rollout_viewer(
env=env,
policy=policy,
seed=seed,
state=state0,
control_dt=control_dt,
max_steps=max_steps_val,
action_low=action_low,
action_high=action_high,
action_mask=action_mask,
)
env.close()
if __name__ == "__main__":
register_configs()
main()

9
scripts/simulate.sh Normal file
View file

@ -0,0 +1,9 @@
#!/usr/bin/env bash
path=$1
uv run simulate.py \
simulation.model_path="$path"/final_model.flax \
simulation.record_video=True \
simulation.video_output_path=../vids/simulation.mp4 \
simulation.max_steps=10000

141
scripts/tools/dump_mjcf.py Normal file
View file

@ -0,0 +1,141 @@
#!/usr/bin/env python3
"""
Dump MJCF XML for a brittle-star morphology using the project's Hydra configs.
Usage examples:
# Use a named morphology config from configs/morphology (Hydra style)
uv run python scripts/analysis/dump_mjcf.py morphology=3_arms
# Use a morphology override YAML (same key as simulation.morphology_override)
uv run python scripts/analysis/dump_mjcf.py \
simulation.morphology_override=configs/morphology/3_arms.yaml
Output path:
Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in
current directory or `runs/morphologies/<name>.xml`.
"""
from __future__ import annotations
import dataclasses
import logging
import sys
from pathlib import Path
from typing import Any, Optional
import hydra
import yaml
from omegaconf import DictConfig, OmegaConf
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.environment.env_config import MorphologyConfig
from brittle_star_project.environment.factory import BrittleStarEnvFactory
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)
def extract_xml_string(obj: Any) -> Optional[str]:
"""
Attempts to serialize the morphology object to an XML string by checking
common dm_control and internal API methods.
"""
serialization_methods = [
"to_xml_string",
"to_xml",
"to_string",
"to_mjcf",
"to_mjcf_string",
"get_mjcf",
"get_mjcf_str",
"export_to_xml_string",
]
# If the object itself has an 'mjcf' attribute, try to serialize that instead
target_obj = getattr(obj, "mjcf", obj)
for method_name in serialization_methods:
method = getattr(target_obj, method_name, None)
if callable(method):
try:
xml_data = method()
# Safely handle both string and byte responses
if isinstance(xml_data, str):
return xml_data
elif isinstance(xml_data, bytes):
return xml_data.decode("utf-8")
except Exception as e:
logger.debug(f"Method {method_name}() failed during serialization: {e}")
return None
def resolve_output_path(cfg: DictConfig) -> Path:
"""Determines the appropriate output path for the MJCF XML."""
dump_out = cfg.get("dump_out", None)
if dump_out is not None:
return Path(hydra.utils.to_absolute_path(str(dump_out)))
morph_name = "morphology"
for arg in sys.argv[1:]:
if arg.startswith("morphology="):
morph_name = arg.split("=", 1)[1]
break
default_out = (
f"runs/morphologies/{morph_name}.xml" if morph_name != "morphology" else "morphology.xml"
)
return Path(hydra.utils.to_absolute_path(default_out))
@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3")
def main(cfg: DictConfig) -> None:
"""Main entry point to construct the morphology and dump its XML."""
logger.info("Initializing morphology construction...")
# Extract morphology config safely using dict `.get()` to avoid OmegaConf AttributeErrors
simulation_cfg = cfg.get("simulation", cfg)
override_path = simulation_cfg.get("morphology_override", None)
if override_path:
logger.info(f"Using morphology override: {override_path}")
with open(hydra.utils.to_absolute_path(override_path), "r") as f:
data = yaml.safe_load(f) or {}
morph_cfg = MorphologyConfig(**data)
else:
# Fallback to default simulation morphology, or an empty base config
morph_node = simulation_cfg.get("morphology", cfg.get("morphology", None))
if morph_node is not None:
# Convert OmegaConf node to dict and instantiate MorphologyConfig.
# This ensures any missing keys gracefully fall back to the dataclass defaults.
morph_dict = OmegaConf.to_container(morph_node, resolve=True)
if isinstance(morph_dict, dict):
# Filter to avoid unexpected kwargs if the dataclass is strictly defined
if dataclasses.is_dataclass(MorphologyConfig):
valid_keys = {f.name for f in dataclasses.fields(MorphologyConfig)}
morph_dict = {k: v for k, v in morph_dict.items() if k in valid_keys}
morph_cfg = MorphologyConfig(**morph_dict)
else:
morph_cfg = MorphologyConfig()
else:
morph_cfg = MorphologyConfig()
morphology = BrittleStarEnvFactory.create_morphology(morph_cfg)
xml_text = extract_xml_string(morphology)
if not xml_text:
raise RuntimeError("Failed to serialize morphology to MJCF/XML. ")
out_path = resolve_output_path(cfg)
out_path.parent.mkdir(parents=True, exist_ok=True)
with out_path.open("w", encoding="utf-8") as f:
f.write(xml_text)
logger.info(f"Successfully exported MJCF XML to: {out_path}")
if __name__ == "__main__":
register_configs()
main()

138
scripts/tools/extract_observation_bounds.py Normal file
View file

@ -0,0 +1,138 @@
#!/usr/bin/env python3
"""Empirically extract observation bounds (focused on joint velocities).
This script creates a MuJoCo environment using the project's factory and
randomly samples actions to discover observed maxima for selected
observation keys (joint_velocity, joint_position, joint_actuator_force).
Usage:
python scripts/extract_observation_bounds.py \
--morphology configs/morphology/3_arms.yaml --num-steps 5000 --seed 42
If `--morphology` is omitted the default `MorphologyConfig()` is used.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import yaml
import numpy as np
from brittle_star_project import BrittleStarEnvFactory, BrittleStarEnv, Backend
from brittle_star_project.environment.env_config import (
MorphologyConfig,
ArenaConfig,
EnvConfig,
)
def load_morphology(path: str | None) -> MorphologyConfig:
if path is None:
return MorphologyConfig()
p = Path(path)
if not p.exists():
raise FileNotFoundError(f"Morphology file not found: {p}")
with open(p, "r") as f:
data = yaml.safe_load(f) or {}
return MorphologyConfig(**data)
def _extract_observations(state):
# Under different backends the returned state may be a dict or an object
obs = getattr(state, "observations", None)
if obs is None and isinstance(state, dict):
obs = state.get("observations", state)
return obs
def find_empirical_bounds(
morph_cfg: MorphologyConfig,
arena_cfg: ArenaConfig,
env_cfg: EnvConfig,
num_steps: int = 5000,
seed: int = 42,
) -> None:
factory = BrittleStarEnvFactory()
raw_env = factory.create_environment(Backend.MJC, morph_cfg, arena_cfg, env_cfg)
env = BrittleStarEnv(raw_env, backend=Backend.MJC, config=env_cfg, morphology_config=morph_cfg)
# Initial reset
state = env.reset(seed=seed)
# Determine action bounds
action_space = getattr(raw_env, "action_space", None)
if action_space is None:
raise RuntimeError("Environment missing `action_space`; cannot sample actions.")
action_low = np.asarray(action_space.low, dtype=np.float32)
action_high = np.asarray(action_space.high, dtype=np.float32)
action_shape = action_low.shape
# Track maximum absolute observed values
tracked_keys = ["joint_velocity", "joint_position", "joint_actuator_force"]
max_observed = {k: 0.0 for k in tracked_keys}
# Include observation at reset
obs0 = _extract_observations(state)
if isinstance(obs0, dict):
for k in tracked_keys:
if k in obs0:
max_observed[k] = max(max_observed[k], float(np.max(np.abs(np.asarray(obs0[k])))))
rng = np.random.RandomState(seed)
for i in range(num_steps):
u = rng.uniform(size=action_shape)
action = action_low + (action_high - action_low) * u
# Provide a numpy RNG to the env step; wrapper will pass it if accepted.
step_out = env.step(state=state, action=action, rng=env.make_rng(seed + i + 1))
# Unpack next state from common return conventions
if hasattr(step_out, "state"):
next_state = step_out.state
elif isinstance(step_out, (tuple, list)) and len(step_out) >= 1:
next_state = step_out[0]
else:
next_state = step_out
obs = _extract_observations(next_state)
if isinstance(obs, dict):
for k in tracked_keys:
if k in obs:
val = float(np.max(np.abs(np.asarray(obs[k]))))
if val > max_observed[k]:
max_observed[k] = val
state = next_state
# Print recommended bounds with a 20% safety margin
print("\n--- Recommended Observation Bounds (20% margin) ---")
for k, v in max_observed.items():
if v == 0.0:
print(f"{k}: observed max 0.0 (increase sampling or inspect env)")
else:
safe = v * 1.2
print(f"{k}: [-{safe:.6f}, {safe:.6f}] (observed max: {v:.6f})")
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument(
"--morphology", type=str, default=None, help="Path to morphology YAML (optional)"
)
parser.add_argument(
"--num-steps", type=int, default=5000, help="Number of random steps to sample"
)
parser.add_argument("--seed", type=int, default=42, help="RNG seed")
args = parser.parse_args()
morph_cfg = load_morphology(args.morphology)
arena_cfg = ArenaConfig()
env_cfg = EnvConfig()
find_empirical_bounds(morph_cfg, arena_cfg, env_cfg, num_steps=args.num_steps, seed=args.seed)
if __name__ == "__main__":
main()

58
scripts/train.py Normal file
View file

@ -0,0 +1,58 @@
import os
import torch
import hydra
from omegaconf import DictConfig, OmegaConf
from brittle_star_project.configs.main_config import BrittleStarConfig
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.trainers.PPOTrainer import PPOTrainer
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
from experiment_logger import init_logger, get_logger
def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper:
"""Create the environment using the structured configuration."""
return BrittleStarJaxEnvWrapper(
morphology=cfg.morphology,
arena=cfg.arena,
env_config=cfg.environment,
num_envs=cfg.ppo.num_envs,
)
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
def main(dict_cfg: DictConfig):
# 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly.
config: BrittleStarConfig = OmegaConf.to_object(
OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
)
# 2. Setup run metadata
# Hydra changes CWD to the output directory by default.
run_dir = os.getcwd()
run_name = os.path.basename(run_dir)
# 3. Initialize Logger
cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True)
init_logger(
run_name=run_name,
full_config=cfg_dict,
logging_cfg=config.logging,
base_dir=os.path.dirname(run_dir),
)
logger = get_logger()
logger.info(f"Hydra-initialized run: {run_name}")
logger.info(f"Output directory: {run_dir}")
# 4. Setup Environment and Torch
env = make_env(config)
torch.backends.cudnn.deterministic = config.experiment.torch_deterministic
# 5. Train - pass structured config directly
ppo_trainer = PPOTrainer(config, env, run_dir, run_name)
ppo_trainer.train()
if __name__ == "__main__":
register_configs()
main()