Videos and evaluation metadata are stored in timestamped folders alongside the model:
runs/your_run/final_model_evaluations/eval_<timestamp>/simulation.mp4
-Using the following script, you can render a top-down and follow camera view for multiple models at once:
+_metadata.yaml` alongside the model_path.
metadata_path: null
diff --git a/scripts/plots/analyze_comparisons.py b/scripts/plots/analyze_comparisons.py
index 3ac6e8a..8a66b4c 100644
--- a/scripts/plots/analyze_comparisons.py
+++ b/scripts/plots/analyze_comparisons.py
@@ -6,18 +6,17 @@ Rate, Distance Remaining).
"""
import os
-import pandas as pd
-import numpy as np
-import matplotlib.pyplot as plt
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
from plot_config import (
- COLORS,
- apply_style,
- BEST_PERFORMER_MARKER,
- BEST_PERFORMER_TEXT,
BEST_PERFORMER_COLOR,
- create_common_parser,
+ BEST_PERFORMER_TEXT,
+ COLORS,
LEGEND_KWARGS,
+ apply_style,
+ create_common_parser,
)
@@ -84,7 +83,8 @@ def plot_grouped_bar(
fig, ax = plt.subplots(figsize=figsize)
bar_width = 0.35
- x_indices = np.arange(len(morphologies))
+ group_spacing = 1.3
+ x_indices = np.arange(len(morphologies)) * group_spacing
all_bars = {}
all_means = []
@@ -118,7 +118,7 @@ def plot_grouped_bar(
)
all_bars[arch] = (x_pos, means, stds, bars)
- for m_idx, m in enumerate(morphologies):
+ for m_idx, _ 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)
@@ -144,12 +144,13 @@ def plot_grouped_bar(
x_ticks_pos = (
x_indices
+ + bar_width # center the label in the 3 bars
+ (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
+ ax.tick_params(axis="x") # More padding for the squares
# X-axis at zero
ax.axhline(0, color="black", linewidth=1.5)
@@ -172,20 +173,7 @@ def plot_grouped_bar(
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.legend(**LEGEND_KWARGS, ncol=len(architectures))
ax.set_facecolor("white")
fig.patch.set_facecolor("white")
@@ -305,22 +293,7 @@ def plot_grouped_bar_alt(
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.legend(**LEGEND_KWARGS, ncol=len(morphologies))
ax.set_facecolor("white")
fig.patch.set_facecolor("white")
@@ -358,8 +331,8 @@ if __name__ == "__main__":
plot_grouped_bar(
df=df,
metric_col="approx_max_velocity",
- ylabel="Max Forward Velocity (cm/s)",
- title="Graceful Degradation: Velocity Across Morphologies",
+ ylabel="",
+ title="Maximal forward velocity (in cm/s)",
output_filename="poster_plot_velocity.png",
output_dir=OUTPUT_DIR,
higher_is_better=True,
diff --git a/scripts/plots/analyze_convergence.py b/scripts/plots/analyze_convergence.py
index 2612bb9..bc12c79 100644
--- a/scripts/plots/analyze_convergence.py
+++ b/scripts/plots/analyze_convergence.py
@@ -44,20 +44,24 @@ 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."""
+ CHECKPOINT = "checkpoint"
ARCH = "architecture"
- TIMESTEPS = "total_trained_timesteps"
- REWARD = "accumulated_reward"
+ TIMESTEPS = "trained_timesteps"
+ REWARD = "eval_return"
VELOCITY = "velocity"
+ EVAL_STEPS = "eval_steps"
+ FINAL_XY_DIST = "final_xy_dist"
+ INITIAL_XY_DIST = "initial_xy_dist"
+ REACHED_TARGET = "reached_target"
# 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",
+ # "centralized 2 arms": "runs/dummy/dummy_centralized_2_arms.csv",
+ "centralized 5 arms": "runs/final-v2-centralized/checkpoint_evaluation.csv",
+ "decentralized fully connected": "runs/final-v2-fully-conn/checkpoint_evaluation.csv",
+ "decentralized ring-level": "runs/final-v2-ring/checkpoint_evaluation.csv",
}
# Architecture profiles for dummy data generation: (max_reward, max_velocity, sigmoid_speed)
@@ -108,7 +112,14 @@ 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]
+ required = [
+ Columns.CHECKPOINT,
+ Columns.TIMESTEPS,
+ Columns.REWARD,
+ Columns.INITIAL_XY_DIST,
+ Columns.FINAL_XY_DIST,
+ Columns.EVAL_STEPS,
+ ]
dfs = []
for arch_name, filepath in file_mapping.items():
@@ -124,17 +135,30 @@ def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame:
continue
df = df[required].copy()
+ df[Columns.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[
+ Columns.EVAL_STEPS
+ ]
df[Columns.ARCH] = arch_name
+ df[Columns.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[
+ Columns.EVAL_STEPS
+ ]
+
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:
+def _convergence_timestep(
+ series: pd.Series, timesteps: pd.Series, checkpoints: pd.Series
+) -> tuple[float, int, int]:
"""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]
+
+ mask = smoothed >= threshold
+ first_idx = mask.idxmax()
+
+ return timesteps.loc[first_idx], first_idx, checkpoints.loc[first_idx]
def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame:
@@ -144,21 +168,40 @@ def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame:
"""
results = []
+ centralized_base = 0
+
for arch in df[Columns.ARCH].unique():
arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
+ reward_timestep, reward_checkpoint_idx, reward_checkpoint = _convergence_timestep(
+ arch_data[Columns.REWARD],
+ arch_data[Columns.TIMESTEPS],
+ arch_data[Columns.CHECKPOINT],
+ )
+
+ velocity_timestep, velocity_checkpoint_idx, velocity_checkpoint = _convergence_timestep(
+ arch_data[Columns.VELOCITY],
+ arch_data[Columns.TIMESTEPS],
+ arch_data[Columns.CHECKPOINT],
+ )
+
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]
- ),
+ "Reward_Convergence_Timestep": reward_timestep,
+ "Reward_Convergence_Checkpoint_Idx": reward_checkpoint_idx,
+ "Reward_Convergence_Checkpoint": reward_checkpoint,
+ "Velocity_Convergence_Timestep": velocity_timestep,
+ "Velocity_Convergence_Checkpoint_Idx": velocity_checkpoint_idx,
+ "Velocity_Convergence_Checkpoint": velocity_checkpoint,
}
)
+ if arch == "centralized 5 arms":
+ centralized_base = reward_checkpoint
+ else:
+ print(arch, "speedup:", 1 - reward_checkpoint / centralized_base)
+
return pd.DataFrame(results)
@@ -303,6 +346,7 @@ def plot_results(df: pd.DataFrame, results: pd.DataFrame, output_dir: str, **kwa
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)
@@ -319,6 +363,12 @@ def run_analysis(output_dir: str, **kwargs):
return
results = analyze_convergence(df)
+ print(
+ results[
+ ["Architecture", "Reward_Convergence_Checkpoint_Idx", "Reward_Convergence_Checkpoint"]
+ ]
+ )
+
plot_results(df, results, output_dir, **kwargs)
logger.info("Analysis complete. Plots saved to disk.")
diff --git a/scripts/plots/plot_config.py b/scripts/plots/plot_config.py
index 5fd6ffd..48f9106 100644
--- a/scripts/plots/plot_config.py
+++ b/scripts/plots/plot_config.py
@@ -4,26 +4,24 @@ 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
+ "CENTRALIZED": "#0D567C", # Blue
+ "FULLY_CONNECTED": "#8C0E0F", # Reddish
+ "RING": "#FCB305", # Pale Yellow
}
-def apply_style(font_size=28):
+def apply_style(font_size=36):
"""
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.labelsize": font_size,
+ "axes.titlesize": font_size,
+ "xtick.labelsize": font_size,
+ "ytick.labelsize": font_size,
+ "legend.fontsize": font_size,
"axes.linewidth": 2,
"axes.spines.top": False,
"axes.spines.right": False,
@@ -44,7 +42,7 @@ BEST_PERFORMER_COLOR = "#D4AF37" # Gold
# Centralized Legend Configuration
LEGEND_KWARGS = {
"loc": "upper center",
- "bbox_to_anchor": (0.5, -0.5),
+ "bbox_to_anchor": (0.5, -0.12),
"frameon": False,
}
diff --git a/scripts/poster_visualisations/render_poster_videos.py b/scripts/poster_visualisations/render_poster_videos.py
new file mode 100644
index 0000000..eca60fc
--- /dev/null
+++ b/scripts/poster_visualisations/render_poster_videos.py
@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+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.video import record_episode_multi_camera
+
+_ARCH_DIR_MAP = {
+ "CENTRALIZED": "centralized",
+ "FULLY_CONNECTED": "fully-connected",
+ "RING": "ring",
+ "SEGMENT": "segment",
+}
+
+ROBOT_COLOR_MAP = {
+ "CENTRALIZED": "#0D567C", # Blue
+ "FULLY_CONNECTED": "#8C0E0F", # Reddish
+ "RING": "#FCB304", # Pale Yellow
+}
+
+
+def _arch_dir(name: str) -> str:
+ return _ARCH_DIR_MAP.get(name, name.lower())
+
+
+def _resolve_overrides(overrides: list[str], count: int) -> list[str | None]:
+ if not overrides:
+ return [None] * count
+ if len(overrides) == 1 and count > 1:
+ return overrides * count
+ if len(overrides) != count:
+ raise ValueError("morphology overrides must match the number of models")
+ return overrides
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description="Render top-down and follow videos for poster.")
+ parser.add_argument("models", nargs="+", help="Paths to .flax checkpoints")
+ parser.add_argument(
+ "--morphology-override",
+ action="append",
+ default=[],
+ help="Override morphology YAML path (repeat to match models)",
+ )
+ parser.add_argument("--output-root", default="vids/poster")
+ parser.add_argument("--seed", type=int, default=0)
+ parser.add_argument("--max-steps", type=int, default=5000)
+ parser.add_argument("--topdown-camera", type=int, default=0)
+ parser.add_argument("--follow-camera", type=int, default=1)
+ parser.add_argument("--topdown-camera-x", type=float, default=-3.0)
+ parser.add_argument("--topdown-camera-y", type=float, default=0.0)
+ parser.add_argument("--topdown-camera-z", type=float, default=4.5)
+ parser.add_argument("--topdown-camera-fovy", type=float, default=None)
+ parser.add_argument("--target-x", type=float, default=-6.0)
+ parser.add_argument("--target-y", type=float, default=0.0)
+ parser.add_argument("--width", type=int, default=2160)
+ parser.add_argument("--height", type=int, default=960)
+ parser.add_argument("--fps", type=int, default=60)
+ parser.add_argument(
+ "--robot-color",
+ default="#2B4162",
+ help="Hex color for the brittle star robot",
+ )
+ args = parser.parse_args()
+
+ if (args.target_x is None) != (args.target_y is None):
+ raise ValueError("target-x and target-y must be provided together")
+
+ target_xy = None
+ if args.target_x is not None:
+ target_xy = (float(args.target_x), float(args.target_y))
+
+ camera_fovy = None
+ if args.topdown_camera_fovy is not None:
+ camera_fovy = {args.topdown_camera: float(args.topdown_camera_fovy)}
+
+ camera_x = None
+ if args.topdown_camera_x is not None:
+ camera_x = {args.topdown_camera: float(args.topdown_camera_x)}
+
+ camera_y = None
+ if args.topdown_camera_y is not None:
+ camera_y = {args.topdown_camera: float(args.topdown_camera_y)}
+
+ camera_z = None
+ if args.topdown_camera_z is not None:
+ camera_z = {args.topdown_camera: float(args.topdown_camera_z)}
+
+ camera_xyz = (camera_x, camera_y, camera_z)
+
+ overrides = _resolve_overrides(args.morphology_override, len(args.models))
+ output_root = Path(args.output_root)
+
+ for model_path_str, override in zip(args.models, overrides):
+ model_path = Path(model_path_str)
+ metadata = load_metadata(model_path, None)
+ training = metadata_to_configs(metadata)
+
+ bundle = build_eval_env(
+ model_path=model_path,
+ training=training,
+ metadata=metadata,
+ morphology_override_path=override,
+ )
+
+ arch_dir = _arch_dir(bundle.architecture)
+ arms_dir = f"{bundle.num_active_arms}arms"
+ out_dir = output_root / arms_dir / arch_dir
+ out_dir.mkdir(parents=True, exist_ok=True)
+
+ output_paths = {
+ args.topdown_camera: out_dir / "topdown.mp4",
+ args.follow_camera: out_dir / "follow.mp4",
+ }
+
+ print(bundle.architecture)
+ color = ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color)
+
+ result = record_episode_multi_camera(
+ env=bundle.env,
+ policy=bundle.policy,
+ seed=args.seed,
+ max_steps=args.max_steps,
+ action_low=bundle.action_low,
+ action_high=bundle.action_high,
+ action_mask=bundle.action_mask,
+ output_paths=output_paths,
+ camera_ids=[args.topdown_camera, args.follow_camera],
+ camera_fovy=camera_fovy,
+ camera_xyz=camera_xyz,
+ target_xy=target_xy,
+ robot_color=color,
+ width=args.width,
+ height=args.height,
+ fps=args.fps,
+ )
+
+ final_dist = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
+ print(
+ f"{arms_dir}/{arch_dir}: return={result.return_:.6f}, len={result.length}, "
+ f"target_reached={result.reached_target}, final_xy_dist={final_dist}"
+ )
+
+ bundle.env.close()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/poster_visualisations/render_poster_videos.sh b/scripts/poster_visualisations/render_poster_videos.sh
new file mode 100644
index 0000000..10648f3
--- /dev/null
+++ b/scripts/poster_visualisations/render_poster_videos.sh
@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+
+# Multi-camera renders per model -> vids/poster/{arms}arms/{arch}/topdown.mp4 + follow.mp4
+
+path=$1
+
+uv run scripts/poster_visualisations/render_poster_videos.py \
+ "$path"/centralized.flax \
+ "$path"/fully-connected.flax \
+ "$path"/ring.flax \
\ No newline at end of file
diff --git a/scripts/poster_visualisations/render_static_path_image.py b/scripts/poster_visualisations/render_static_path_image.py
new file mode 100644
index 0000000..8229f52
--- /dev/null
+++ b/scripts/poster_visualisations/render_static_path_image.py
@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+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 (
+ _get_observations,
+ _maybe_clip_action,
+ _target_reached,
+)
+from brittle_star_project.evaluation.video import (
+ _apply_camera_overrides,
+ _ensure_offscreen_size,
+ hex_to_rgba,
+)
+
+ROBOT_COLOR_MAP = {
+ "CENTRALIZED": "#0D567C", # Blue
+ "FULLY_CONNECTED": "#8C0E0F", # Reddish
+ "RING": "#FCB304", # Pale Yellow
+}
+
+
+def _enum_value(enum_obj, *names: str) -> int:
+ for name in names:
+ if hasattr(enum_obj, name):
+ return int(getattr(enum_obj, name))
+ raise AttributeError(f"Could not find any of {names!r} on {enum_obj!r}")
+
+
+def _append_sphere(scene, mujoco, center: np.ndarray, radius: float, rgba: np.ndarray) -> None:
+ geom = scene.geoms[scene.ngeom]
+ mujoco.mjv_initGeom(
+ geom,
+ mujoco.mjtGeom.mjGEOM_SPHERE,
+ np.asarray([radius, 0.0, 0.0], dtype=np.float32),
+ center,
+ np.eye(3, dtype=np.float32).reshape(-1),
+ rgba,
+ )
+ scene.ngeom += 1
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description="Render a static path image from a rollout.")
+ parser.add_argument("model", help="Path to .flax checkpoint")
+ parser.add_argument("--morphology-override", default=None)
+ parser.add_argument("--output-path", required=True)
+ parser.add_argument("--seed", type=int, default=0)
+ parser.add_argument("--max-steps", type=int, default=5000)
+ parser.add_argument("--body-name", default="BrittleStarMorphology/central_disk")
+ parser.add_argument("--camera-id", type=int, default=0)
+ parser.add_argument("--camera-x", type=float, default=-3.0)
+ parser.add_argument("--camera-y", type=float, default=0.0)
+ parser.add_argument("--camera-z", type=float, default=4.5)
+ parser.add_argument("--camera-fovy", type=float, default=None)
+ parser.add_argument("--target-x", type=float, default=-6.0)
+ parser.add_argument("--target-y", type=float, default=0.0)
+ parser.add_argument("--width", type=int, default=2160)
+ parser.add_argument("--height", type=int, default=960)
+ parser.add_argument("--frame-stride", type=int, default=15)
+ parser.add_argument(
+ "--path-color", default="#FA9F42"
+ ) # ring = #888888, centralized = #2B4162, fully connected = FA9F42
+ parser.add_argument(
+ "--robot-color",
+ default="#FA9F42",
+ help="Hex color for brittle star robot (e.g. #ff0000)",
+ )
+ args = parser.parse_args()
+
+ model_path = Path(args.model)
+ metadata = load_metadata(model_path, None)
+ training = metadata_to_configs(metadata)
+
+ bundle = build_eval_env(
+ model_path=model_path,
+ training=training,
+ metadata=metadata,
+ morphology_override_path=args.morphology_override,
+ )
+
+ if (args.target_x is None) != (args.target_y is None):
+ raise ValueError("target-x and target-y must be provided together")
+
+ target_xy = None
+ if args.target_x is not None:
+ target_xy = (float(args.target_x), float(args.target_y))
+
+ try:
+ import imageio
+ import mujoco
+ except ImportError as e:
+ raise ImportError(
+ "Static image rendering requires 'mujoco' and 'imageio'. "
+ "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
+ ) from e
+
+ reset_kwargs = {}
+ if target_xy is not None:
+ reset_kwargs["target_position"] = (target_xy[0], target_xy[1], 0.0)
+
+ state = bundle.env.reset(seed=args.seed, **reset_kwargs)
+ model = state.mj_model
+ data = state.mj_data
+
+ _apply_camera_overrides(
+ model,
+ camera_fovy={args.camera_id: float(args.camera_fovy)}
+ if args.camera_fovy is not None
+ else None,
+ camera_xyz=(
+ {args.camera_id: float(args.camera_x)} if args.camera_x is not None else None,
+ {args.camera_id: float(args.camera_y)} if args.camera_y is not None else None,
+ {args.camera_id: float(args.camera_z)} if args.camera_z is not None else None,
+ ),
+ )
+ _ensure_offscreen_size(model, args.width, args.height)
+
+ body_id = mujoco.mj_name2id(
+ model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
+ )
+
+ robot_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color), 1.0)
+
+ # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
+ if args.robot_color is not None:
+ # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
+ nbody = int(model.nbody)
+ body_parent = model.body_parentid
+ robot_body_ids = set([int(body_id)])
+ for i in range(1, nbody):
+ cur = int(i)
+ # walk up until root (0) or until we hit the robot root
+ while cur not in (-1, 0, int(body_id)):
+ cur = int(body_parent[cur])
+ if cur == int(body_id):
+ robot_body_ids.add(i)
+
+ # Recolor geoms whose body id is in the robot subtree
+ for g in range(int(model.ngeom)):
+ if int(model.geom_bodyid[g]) in robot_body_ids:
+ model.geom_rgba[g][:] = robot_rgba
+
+ # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
+ if args.robot_color is not None:
+ # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
+ nbody = int(model.nbody)
+ body_parent = model.body_parentid
+ robot_body_ids = set([int(body_id)])
+ for i in range(1, nbody):
+ cur = int(i)
+ # walk up until root (0) or until we hit the robot root
+ while cur not in (-1, 0, int(body_id)):
+ cur = int(body_parent[cur])
+ if cur == int(body_id):
+ robot_body_ids.add(i)
+
+ # Recolor geoms whose body id is in the robot subtree
+ for g in range(int(model.ngeom)):
+ if int(model.geom_bodyid[g]) in robot_body_ids:
+ model.geom_rgba[g][:] = robot_rgba
+
+ positions = []
+ observations = _get_observations(state)
+
+ for _ in range(int(args.max_steps)):
+ positions.append(np.asarray(data.xpos[body_id], dtype=np.float32))
+
+ obs_dict = observations or {}
+ action = bundle.policy.act(observations=obs_dict)
+ if bundle.action_mask is not None:
+ action = action[bundle.action_mask]
+ action = _maybe_clip_action(action, bundle.action_low, bundle.action_high)
+
+ state = bundle.env.step(state=state, action=action)
+ data = state.mj_data
+ observations = _get_observations(state)
+
+ if _target_reached(state=state):
+ break
+
+ positions_arr = np.vstack(positions)
+ if len(positions_arr) < 2:
+ raise ValueError("Need at least two rollout positions to render a path")
+
+ path_points = positions_arr.copy()
+ path_points[:, 2] -= 0.02
+
+ path_step = max(1, int(args.frame_stride))
+ path_points_visible = path_points[::path_step]
+ path_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.path_color), 0.92)
+
+ ctx = mujoco.GLContext(args.width, args.height)
+ ctx.make_current()
+ try:
+ catmask = _enum_value(mujoco.mjtCatBit, "mjCAT_ALL")
+ camera_type = _enum_value(mujoco.mjtCamera, "mjCAMERA_FIXED")
+ font_scale = _enum_value(mujoco.mjtFontScale, "mjFONTSCALE_100")
+
+ maxgeom = int(model.ngeom + len(path_points_visible) + 8)
+ scene = mujoco.MjvScene(model, maxgeom=maxgeom)
+ option = mujoco.MjvOption()
+ perturb = mujoco.MjvPerturb()
+ camera = mujoco.MjvCamera()
+ mujoco.mjv_defaultOption(option)
+ mujoco.mjv_defaultPerturb(perturb)
+ mujoco.mjv_defaultCamera(camera)
+ camera.type = camera_type
+ camera.fixedcamid = int(args.camera_id)
+ if hasattr(camera, "trackbodyid"):
+ camera.trackbodyid = -1
+
+ context = mujoco.MjrContext(model, font_scale)
+ viewport = mujoco.MjrRect(0, 0, args.width, args.height)
+
+ mujoco.mjv_updateScene(model, data, option, perturb, camera, catmask, scene)
+
+ for idx, path_point in enumerate(path_points_visible):
+ path_rgba[3] = 0.10 + 0.70 * (idx / max(len(path_points_visible) - 1, 1))
+ _append_sphere(scene, mujoco, path_point, 0.03, path_rgba)
+
+ rgb = np.empty((args.height, args.width, 3), dtype=np.uint8)
+ depth = np.empty((args.height, args.width), dtype=np.float32)
+ mujoco.mjr_render(viewport, scene, context)
+ mujoco.mjr_readPixels(rgb, depth, viewport, context)
+ imageio.imwrite(args.output_path, np.flipud(rgb))
+
+ context.free()
+ finally:
+ ctx.free()
+
+ bundle.env.close()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/poster_visualisations/render_static_path_image.sh b/scripts/poster_visualisations/render_static_path_image.sh
new file mode 100644
index 0000000..ab21ca0
--- /dev/null
+++ b/scripts/poster_visualisations/render_static_path_image.sh
@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+
+# Static path image + optional ghost render
+
+path=$1 # path to .flax model with metadata.yaml alongside it
+
+uv run scripts/poster_visualisations/render_static_path_image.py \
+ "$path" \
+ --output-path vids/poster/5arms/centralized/path.png \
+ --ghost-overlay
diff --git a/scripts/simulate.py b/scripts/simulate.py
index 9f3e550..db8f70d 100644
--- a/scripts/simulate.py
+++ b/scripts/simulate.py
@@ -107,6 +107,9 @@ def main(dict_cfg: DictConfig) -> None:
action_mask=action_mask,
output_path=output_path,
camera_id=sim_cfg.camera_id,
+ width=sim_cfg.video_width,
+ height=sim_cfg.video_height,
+ fps=sim_cfg.video_fps,
)
save_evaluation_metadata(
diff --git a/scripts/tools/download_wandb_project.py b/scripts/tools/download_wandb_project.py
new file mode 100644
index 0000000..bf2fb11
--- /dev/null
+++ b/scripts/tools/download_wandb_project.py
@@ -0,0 +1,142 @@
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import threading
+import wandb
+import argparse
+
+# tune these depending on network / W&B limits
+MAX_RUN_WORKERS = 8
+MAX_FILE_WORKERS = 16
+MAX_ARTIFACT_WORKERS = 8
+
+api = wandb.Api()
+
+print_lock = threading.Lock()
+
+
+def safe_print(*args, **kwargs):
+ with print_lock:
+ print(*args, **kwargs)
+
+
+def download_file(file, run_dir):
+ target = run_dir / file.name
+
+ try:
+ # skip existing files
+ if target.exists():
+ return f"SKIP FILE {target}"
+
+ target.parent.mkdir(parents=True, exist_ok=True)
+
+ file.download(root=run_dir, replace=False)
+
+ return f"DONE FILE {target}"
+
+ except Exception as e:
+ return f"FAIL FILE {target}: {e}"
+
+
+def sanitize_artifact_name(name: str):
+ return name.replace(":", "_")
+
+
+def download_artifact(artifact, artifact_root):
+ try:
+ artifact_name = sanitize_artifact_name(artifact.name)
+ artifact_dir = artifact_root / artifact_name
+
+ if artifact_dir.exists() and any(artifact_dir.iterdir()):
+ return f"SKIP ARTIFACT {artifact.name}"
+
+ artifact_dir.mkdir(parents=True, exist_ok=True)
+
+ artifact.download(root=artifact_dir)
+
+ return f"DONE ARTIFACT {artifact.name}"
+
+ except Exception as e:
+ return f"FAIL ARTIFACT {artifact.name}: {e}"
+
+
+def download_run(run, root):
+ run_dir = root / f"{run.name}"
+ run_dir.mkdir(parents=True, exist_ok=True)
+
+ safe_print(f"\n=== {run.name} ({run.id}) ===")
+
+ # -------------------------
+ # Download regular run files
+ # -------------------------
+ files = list(run.files())
+
+ with ThreadPoolExecutor(max_workers=MAX_FILE_WORKERS) as executor:
+ futures = [executor.submit(download_file, file, run_dir) for file in files]
+
+ for future in as_completed(futures):
+ safe_print(future.result())
+
+ # -------------------------
+ # Download logged artifacts
+ # -------------------------
+ artifact_root = run_dir / "artifacts"
+
+ try:
+ artifacts = list(run.logged_artifacts())
+ safe_print(f"Found {len(artifacts)} artifacts for {run.name}")
+
+ with ThreadPoolExecutor(max_workers=MAX_ARTIFACT_WORKERS) as executor:
+ futures = [
+ executor.submit(download_artifact, artifact, artifact_root)
+ for artifact in artifacts
+ ]
+
+ for future in as_completed(futures):
+ safe_print(future.result())
+
+ except Exception as e:
+ safe_print(f"Artifact download failed for {run.name}: {e}")
+
+ # -------------------------
+ # OPTIONAL: download used/input artifacts
+ # -------------------------
+ # try:
+ # used_artifacts = list(run.used_artifacts())
+ # used_root = run_dir / "used_artifacts"
+ #
+ # for artifact in used_artifacts:
+ # download_artifact(artifact, used_root)
+ # except Exception as e:
+ # safe_print(f"Used artifact download failed: {e}")
+
+ safe_print(f"Finished {run.name}")
+
+
+def main(entity: str, project: str, root: Path):
+ root.mkdir(exist_ok=True)
+
+ runs = list(api.runs(f"{entity}/{project}"))
+
+ safe_print(f"Found {len(runs)} runs")
+
+ with ThreadPoolExecutor(max_workers=MAX_RUN_WORKERS) as executor:
+ futures = [executor.submit(download_run, run, root) for run in runs]
+
+ for future in as_completed(futures):
+ try:
+ future.result()
+ except Exception as e:
+ safe_print("RUN FAILED:", e)
+
+ safe_print("\nAll downloads complete.")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--entity", type=str, default="SEL3-2026-Groep-4")
+ parser.add_argument("--project", type=str, required=True)
+ parser.add_argument("--root", type=str, default="runs")
+ args = parser.parse_args()
+
+ root = Path(args.root)
+ main(entity=args.entity, project=args.project, root=root)
diff --git a/search/search_index.json b/search/search_index.json
index 17731f7..0f6f92a 100644
--- a/search/search_index.json
+++ b/search/search_index.json
@@ -1 +1 @@
-{"config":{"lang":["en"],"separator":"[\\s\\-]+","pipeline":["stopWordFilter"],"fields":{"title":{"boost":1000.0},"text":{"boost":1.0},"tags":{"boost":1000000.0}}},"docs":[{"location":"","title":"Documentation","text":"Welcome to the Brittle Star project documentation. This codebase contains the implementations and research for the scientific evaluation of controller modularity in brittle-star-like robots trained using Reinforcement Learning.
For the core codebase, scripts, and contribution history, visit our GitHub Repository.
"},{"location":"#core-requirements-guides","title":"Core Requirements & Guides","text":" - Installation Instructions: Steps to set up your development environment locally or in a devcontainer using
uv, including GPU configuration. For High-Performance Computing (HPC) setup details, see the HPC Guide. - How to Run Experiments: A complete guide on running training jobs, setting custom hyperparameters, and overriding config options using Hydra.
- Results & Reproduction: Guide on how to access our public WandB training runs table and reproduce our training and evaluation phases (determining the best checkpoint vs. comparing architectures).
- Contribution Guidelines: Standards, rules, and best practices for developing and adding code to the repository.
- Repository Structure: Overview of the directories and files within the codebase.
"},{"location":"#repository-structure","title":"Repository Structure","text":".\n\u251c\u2500\u2500 configs/ # Hydra configuration files (YAML)\n\u251c\u2500\u2500 docs/ # Comprehensive documentation and API guides\n\u251c\u2500\u2500 runs/ # Default output directory for Hydra and training artifacts\n\u251c\u2500\u2500 scripts/ # High-level entrypoints for training, simulation, and evaluation\n\u251c\u2500\u2500 src/\n\u2502 \u251c\u2500\u2500 brittle_star_project/ # Core library and environment logic\n\u2502 \u2502 \u251c\u2500\u2500 evaluation/ # Checkpoint evaluation, rollout logic, and metrics persistence\n\u2502 \u2502 \u2514\u2500\u2500 trainers/ # Training loop implementations (e.g., PPO)\n\u2502 \u2514\u2500\u2500 experiment_logger/ # Standalone logging package\n\u2514\u2500\u2500 tests/ # Unit and integration tests\n
"},{"location":"#design-architecture-design","title":"Design & architecture (/design)","text":"If you are interested in the \"why did you do it like this?\"
- Actor-Critic Architecture: Description of the actor-critic pipeline.
- Communication Scheme: Message propagation, Nerve-Net style.
- Modularity & Topology: Macroscopic brain topology, centralized, arm-level, segment-level.
- Input & Action Spaces: Description of the model's input and output.
- Reinforcement Learning Algorithm: RL techniques, i.e. PPO.
- Reward Function & Observation Space: Goals, fitness tracking, and reward structures.
"},{"location":"#api-reference-api","title":"API reference (/api)","text":"If you are interested in the \"how do I use it?\"
- Brittle Star Environment: MuJoCo environment interaction and configuration.
- Training Models: How to configure and run experiments.
- Tracking & Monitoring: Setting up WandB and TensorBoard to monitor runs.
- Checkpoint & Model Evaluation: Evaluating checkpoints and comparing fault tolerance.
- Interactive Simulation & Visualization: Visualizing models in the MuJoCo viewer or rendering simulation videos.
- Analysis & Plotting Tools: Comparing checkpoints and generating plots.
- Results & Reproduction: Accessing WandB results and running reproduction pipelines.
"},{"location":"CONTRIBUTING/","title":"Contribution Guidelines","text":"This document outlines the contribution protocols for the scientific software engineering project focusing on bio-inspired control architectures for brittle-star-like robots. The primary objective of this project is to produce scientific insight, rather than a commercial product.
"},{"location":"CONTRIBUTING/#1-scientific-context-methodology","title":"1. Scientific Context & Methodology","text":" - Research Focus: The goal is to study how controller modularity affects learning speed, coordination, and fault tolerance in brittle-star locomotion.
- Hypothesis-Driven Design: Clear hypotheses must dictate a structured methodology and rigorous evaluation. All design decisions must be formally documented prior to implementation.
- Scaffolding Approach: Development must start with simple setups before progressively increasing the complexity of environments and morphologies.
- Evaluation of Results: Negative results possess scientific validity when thoroughly analyzed. If a controller fails to learn locomotion, providing a comprehensive analysis of the failure is considered a strong scientific contribution.
- Reproducibility: Contributors must utilize fixed library versions. Configuration systems (such as json, gin, or yaml) must be employed to ensure reproducible runs.
"},{"location":"CONTRIBUTING/#2-clean-code-code-quality","title":"2. Clean Code & Code Quality","text":"Code readability is paramount, as code is read far more frequently than it is written.
- Naming Conventions: Variables and functions must utilize consistent, intention-revealing names. A long, descriptive name is strictly preferred over a short name accompanied by a comment.
- Function Design: Functions must be modular and adhere to the single responsibility principle. Arguments must be minimized, and boolean flag arguments controlling behavior should be avoided.
- Commenting: Code must document the \"how,\" while comments are strictly reserved for documenting the \"why\". Commented-out code is prohibited and must be deleted via version control.
- YAGNI: Contributors must adhere to the \"You Aren't Gonna Need It\" (YAGNI) principle and actively avoid premature optimization.
- Notebooks: Jupyter Notebooks are strictly limited to quick prototyping, tutorials, demonstrations, or post-processing analysis. They are explicitly forbidden for general software development because they discourage modularity.
"},{"location":"CONTRIBUTING/#3-version-control-repository-structure","title":"3. Version Control & Repository Structure","text":" - Git Practices: Commits must be frequent and small. Each commit should relate to exactly one piece of functionality.
- Branching Strategy: The
dev branch serves as the integration branch for pushing and merging code. Only stable releases may be pushed to the main branch. - Artifact Management: Data files, trained models, and large datasets must never be committed directly to Git. Git Large File Storage (LFS) must be used for tracking large files. All developers must have
git-lfs installed locally (see DEVELOPMENT.md for setup). - Repository Layout: The repository must maintain the following core directories:
src/ for algorithms, env/ for MuJoCo wrappers, config/ for experiment configurations, experiments/ for scripts, docs/ for Doxygen or ReadTheDocs documentation, and tests/ for unit tests.
"},{"location":"CONTRIBUTING/#4-architecture-tooling","title":"4. Architecture & Tooling","text":" - Algorithms & Frameworks: Proximal Policy Optimization (PPO) is the recommended baseline algorithm. CleanRL should be used as a starting point and adapted for continuous action spaces. All Artificial Neural Network (ANN) controller architectures must be implemented using Flax.
- Simulation: The simulation environment utilizes a MuJoCo brittle star. XML MuJoCo structures must remain realistic and respect morphological constraints.
- Experiment Tracking: Weights & Biases (wandb) must be utilized for tracking and logging all experiments.
- Code Styling: All code must conform to the chosen style guide (Google standard). This is enforced via
uv using ruff and pre-commit hooks.
"},{"location":"CONTRIBUTING/#5-ai-assisted-development-code-review","title":"5. AI-Assisted Development & Code Review","text":"This project supports AI-assisted development to enhance productivity, but contributors must take full responsibility for all AI-generated outputs.
- Self-Review Requirement: Contributors must thoroughly self-review all AI-assisted code, documentation, and configurations before requesting peer review. This includes verifying correctness, adherence to project standards, scientific validity, and integration with existing code.
- Quality Standards: AI-generated content must meet the same rigorous standards as manually written code, including proper testing, documentation, and alignment with the scientific methodology outlined in Section 1.
- Available Skills: This project provides specific AI skills for common tasks (located in
.agents/skills/), including linting and testing workflows. Contributors should leverage these skills to maintain consistency and quality. - Transparency: When using AI assistance for complex algorithmic decisions or scientific design choices, contributors should document the rationale in commit messages or code comments where appropriate.
"},{"location":"DEVELOPMENT/","title":"Development Guide","text":"This guide outlines how to set up the development environment for this project, prioritizing reproducible builds, environment parity, and cross-hardware compatibility.
"},{"location":"DEVELOPMENT/#reproducibility-uv","title":"Reproducibility & uv","text":"This project uses uv to manage dependencies and virtual environments. The uv.lock file is the absolute source of truth for package versions and must always be committed.
"},{"location":"DEVELOPMENT/#source-of-truth","title":"Source of Truth","text":" - Never modify
uv.lock manually. - To add a dependency, run
uv add <package>. - To update dependencies, run
uv lock --upgrade. - To sync your environment with the lockfile, run
uv sync --frozen.
"},{"location":"DEVELOPMENT/#git-lfs-critical","title":"Git LFS (Critical)","text":"All developers must have Git LFS installed locally. This repository tracks model weights (.pt, .safetensors, etc.), recordings (.mp4), and datasets using Git LFS.
- Setup: Run
git lfs install after cloning this repository. If you are using the .devcontainer or flake.nix, LFS is typically available automatically. - If you clone without LFS installed, run
git lfs pull after installation to fetch the actual data files instead of the small pointer files.
"},{"location":"DEVELOPMENT/#devcontainer-setup-recommended","title":"Devcontainer Setup (Recommended)","text":"The devcontainer provides an identical experience to local development but with all system dependencies pre-configured. It automatically detects your hardware (GPU vs CPU) and syncs the appropriate dependencies.
"},{"location":"DEVELOPMENT/#prerequisites","title":"Prerequisites","text":" - Docker Desktop or Docker Engine.
- NVIDIA Container Toolkit (for GPU support).
"},{"location":"DEVELOPMENT/#setup-for-vs-code","title":"Setup for VS Code","text":" - Install the Dev Containers extension.
- Open the project and click Reopen in Container.
- On first launch, the
post-create.sh script will: - Detect if an NVIDIA GPU is available via
nvidia-smi. - Run
uv sync --frozen --extra cuda if a GPU is found. - Run
uv sync --frozen otherwise. - The environment is stored in a named volume for
.venv to ensure persistence and performance.
"},{"location":"DEVELOPMENT/#setup-for-jetbrains-ides","title":"Setup for JetBrains IDEs","text":" - The IDE will detect the
.devcontainer/devcontainer.json file. - The environment is pre-configured to point to
/workspaces/project/.venv. - The hardware-aware sync will run automatically during container creation.
"},{"location":"DEVELOPMENT/#local-development-alternative","title":"Local Development (Alternative)","text":"If you prefer not to use Docker:
- Install uv.
- Run
uv sync --frozen (CPU) or uv sync --frozen --extra cuda (GPU).
"},{"location":"DEVELOPMENT/#hardware-acceleration-jax","title":"Hardware Acceleration (JAX)","text":"Verify your setup by running the JAX initialization test:
uv run pytest tests/test_jax_init.py\n
In the devcontainer, this will succeed on both CPU and GPU. A GpuDevice is expected if a GPU is detected and the cuda extra was installed.
"},{"location":"DEVELOPMENT/#logging-monitoring","title":"Logging & Monitoring","text":"This project uses a unified logging system through the experiment_logger package.
- Usage in Code: To use the logger in your scripts, refer to the package README for the API reference.
- WandB/TensorBoard Setup: For information on how to configure tracking for experiments, see the Tracking & Monitoring API Guide.
The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job), adjusting progress bars and fallback modes accordingly.
"},{"location":"HPC/","title":"HPC Guide","text":"Full documentation: https://docs.hpc.ugent.be/
"},{"location":"HPC/#storage-overview","title":"Storage Overview","text":" - Run Outputs: Written to
$VSC_SCRATCH during the job (fast I/O) and copied to $VSC_DATA at the end for persistence. - Virtual Environments: Managed on
$VSC_DATA by mirroring configuration files. This avoids the 3GB home quota without requiring symlinks in the project root.
"},{"location":"HPC/#initial-environment-setup","title":"Initial Environment Setup","text":"Run once after cloning the repository. This script handles all modules, mirroring, and environment synchronization.
# Option A: Interactive (on a compute node)\nmodule swap cluster/donphan # Debug cluster (CPU only)\n# OR for GPU clusters:\n# module swap cluster/joltik\n# module swap cluster/accelgor\n# module swap cluster/litleo\n\nqsub -I -l nodes=1:gpus=1 # Only for GPU clusters\ncd \"${PBS_O_WORKDIR}\"\nbash scripts/hpc/install.sh\n\n# Option B: Batch (Run in background)\n# NOTE: GPU clusters (joltik/accelgor/litleo) require -l gpus=1 at runtime\nqsub -l gpus=1 scripts/hpc/install.sh\n
"},{"location":"HPC/#production-vs-debug-clusters","title":"Production vs. Debug Clusters","text":"Our scripts are cluster-agnostic and do not have hardcoded GPU requirements. Instead, you must request GPUs at runtime using the -l gpus=1 flag when submitting to a production GPU cluster.
"},{"location":"HPC/#debugging-donphan","title":"Debugging (Donphan)","text":"The donphan cluster does not support GPUs. Simply run the scripts without extra resource flags:
module swap cluster/donphan\nqsub scripts/hpc/train.pbs\n
"},{"location":"HPC/#production-joltik-accelgor-litleo","title":"Production (Joltik, Accelgor, Litleo)","text":"These clusters provide GPU acceleration and require a GPU request at runtime:
module swap cluster/joltik # or accelgor/litleo\nqsub -l gpus=1 scripts/hpc/train.pbs\n
"},{"location":"HPC/#interactive-debugging","title":"Interactive Debugging","text":"To activate your environment for interactive work, simply run the same install.sh script.
qsub -I -l nodes=1:ppn=4 -l walltime=1:00:00\ncd \"$PBS_O_WORKDIR\"\nbash scripts/hpc/install.sh\n
"},{"location":"HPC/#verification-commands","title":"Verification Commands","text":"After installation, run these commands to ensure your environment is set up correctly:
-
Verify Quota Safety:
ls -d venvs 2>/dev/null && echo \"FAIL\" || echo \">>> PASS: Project root is clean.\"\n
-
Verify Library Versions (NumPy Fix):
python -c \"import numpy; print(f'NumPy: {numpy.__version__}')\"\n# Expected: 2.x.x (Venv version), not 1.2x (System version)\n
-
Verify GPU Access:
python -c \"import torch, jax; print(f'GPU: {torch.cuda.is_available()}'); print(f'JAX: {jax.devices()}')\"\n
"},{"location":"HPC/#managing-dependencies","title":"Managing Dependencies","text":"env/hpc/requirements.txt is auto-generated from pyproject.toml. To regenerate:
uv run scripts/hpc/export_requirements.py\n
Modules listed in env/hpc/modules.txt are automatically excluded from the pip requirements to save space and use HPC-optimized binaries.
"},{"location":"api/analysis/","title":"Analysis & Plotting Tools","text":"This guide outlines the tools available for analyzing experimental data and generating poster-quality visualizations for the Brittle Star project.
"},{"location":"api/analysis/#shared-configuration","title":"Shared Configuration","text":"All plotting scripts share a central configuration in scripts/plots/plot_config.py. This file defines:
- Color Palette: A color-blind friendly, high-contrast palette for different architectures.
- Typography: Consistent font sizes and styles tailored for A0 posters.
- Markers: Shared visual indicators, such as the \u2605 used for best performers.
"},{"location":"api/analysis/#comparison-visualization","title":"Comparison Visualization","text":"The scripts/plots/analyze_comparisons.py script generates grouped bar charts comparing the performance of different architectures across various morphologies.
"},{"location":"api/analysis/#usage","title":"Usage","text":"Run the script from the root of the project, providing the path to your evaluation CSV:
# Basic usage (saves PNG and SVG to runs/evaluation/plots/)\nuv run python scripts/plots/analyze_comparisons.py path/to/results.csv\n\n# Advanced usage for Figma/Poster integration\nuv run python scripts/plots/analyze_comparisons.py path/to/results.csv \\\n --output_dir docs/assets/plots/ \\\n --font_size 30 \\\n --fig_width 14 \\\n --fig_height 10\n
"},{"location":"api/analysis/#cli-arguments","title":"CLI Arguments","text":" input_csv: (Required) Path to the CSV file containing evaluation results.--output_dir, -o: Directory where plots will be saved (default: runs/evaluation/plots).--show_titles: Include titles in the plots. Default is False, as titles are typically added natively in design tools like Figma.--font_size: Base font size in points (default: 28).--fig_width / --fig_height: Physical dimensions of the plot in inches. Match these to your Figma layout to maintain exact font sizes.
"},{"location":"api/analysis/#outputs","title":"Outputs","text":"The script generates four key plots, each saved as both .png and .svg:
- Forward Velocity: Grouped bar chart (cm/s).
- Accumulated Reward: Mean cumulative reward.
- Success Rate: Target acquisition percentage.
- Distance Remaining: Navigational accuracy.
"},{"location":"api/analysis/#convergence-analysis","title":"Convergence Analysis","text":"The scripts/plots/analyze_convergence.py script determines the convergence point of training runs.
"},{"location":"api/analysis/#usage_1","title":"Usage","text":"uv run python scripts/plots/analyze_convergence.py --output_dir runs/convergence/\n
"},{"location":"api/analysis/#configuration","title":"Configuration","text":" - File Mapping: The script uses hardcoded paths in the
FILE_MAPPING dictionary. Update these paths to point to your specific run evaluation files. - CLI Arguments: Supports the same
--show_titles, --font_size, and --fig_width/height flags as the comparison script.
"},{"location":"api/analysis/#outputs_1","title":"Outputs","text":"Generates three plots (PNG & SVG):
convergence_comparison: Grouped horizontal bar chart.progress_reward_curves: Line plots of reward over time.progress_velocity_curves: Line plots of velocity over time.
"},{"location":"api/analysis/#poster-integration-figma","title":"Poster Integration (Figma)","text":""},{"location":"api/analysis/#svg-scaling","title":"SVG & Scaling","text":"We recommend using the SVG outputs for poster design in Figma:
- No Resolution Loss: SVGs are vector-based and will remain sharp at any size.
- Native Text: Text in the SVG imports as native text layers in Figma.
- Exact Font Matching: To ensure a
28pt font in the plot matches a 28pt font in your poster, set the --fig_width and --fig_height to match the physical dimensions of the plot box in your Figma layout. - Editable: You can \"Ungroup\" the SVG in Figma to manually move labels, adjust colors, or tweak individual bars.
"},{"location":"api/analysis/#image-placeholders","title":"Image Placeholders","text":"The comparison charts include light-gray square placeholders below the X-axis. These are designed as guides; in Figma, you can drop your morphology renders or illustrations directly on top of these squares.
"},{"location":"api/environment/","title":"Brittle star environment","text":""},{"location":"api/environment/#creation","title":"Creation","text":"The environment package contains a factory class BrittleStarEnvFactory that creates instances of the environment/morphologies/... It uses the configuration classes defined in env_config.py to create the instances.
"},{"location":"api/environment/#configuration","title":"Configuration","text":"The data classes in env_config have default values as stated in the tutorials.
- MorphologyConfig: configuration for the morphology of the brittle star. Contains number of arms, number of segments per arm, and control mode.
- ArenaConfig: configuration for the arena. Sets the size of the arena, whether to set the ground floor to sand, attach a target and sizes of the walls.
- EnvConfig: configuration for the environment. These set shared settings such as camera locations, simulation time and the task.
"},{"location":"api/environment/#backend-and-task-enums","title":"Backend and Task enums","text":"The Backend enum specifies either an MJC or MJX backend.
- MJC: runs on CPU
- MJX: uses jax on the gpu
The Task enum specifies which task to use. 2 items are present:
- DIRECTED_LOCOMOTION: move to a target location
- LIGHT_ESCAPE: situation where the robot must move to a darker location
"},{"location":"api/evaluation/","title":"Checkpoint & Model Evaluation","text":"This guide covers how to evaluate trained brittle star models, with a focus on measuring defect tolerance (amputations) across different controller architectures.
"},{"location":"api/evaluation/#checkpoint-evaluation-during-training","title":"Checkpoint Evaluation (During Training)","text":"The PPOTrainer can automatically evaluate every saved checkpoint using the fast MJX backend. This is enabled via configuration.
"},{"location":"api/evaluation/#configuration","title":"Configuration","text":"In your experiment config or via CLI:
python scripts/train.py evaluation.evaluate_checkpoints=true evaluation.eval_max_steps=5000\n
Results are saved to runs/<run_dir>/metrics/checkpoint_evaluation.csv and synced to Weights & Biases if enabled.
"},{"location":"api/evaluation/#cross-model-fault-tolerance-analysis","title":"Cross-Model & Fault Tolerance Analysis","text":"To measure how well different controllers handle damage (amputations), use scripts/compare_models.py. This script performs a grid search over models x morphologies.
- Create or update a YAML file in
configs/evaluation. - Run the benchmark:
python scripts/compare_models.py evaluation=poster\n
The script will evaluate every combination of model and morphology for the specified number of episodes.
The results are saved to a CSV (default: metrics/model_comparison.csv).
"},{"location":"api/evaluation/#csv-schema","title":"CSV Schema","text":"Column Description model_path Path to the trained weights. architecture The morph_mode of the model (e.g., CENTRALIZED, RING). arm_0 ... arm_4 Number of segments in each arm slot (0 = amputated). num_active_arms Total number of arms with segments > 0. seed The episode seed. eval_return Accumulated shaped reward. approx_max_velocity Average velocity: (initial_dist - final_dist) / steps. reached_target Whether the robot finished within the success radius."},{"location":"api/evaluation/#post-hoc-checkpoint-scanning","title":"Post-hoc Checkpoint Scanning","text":"If you need to re-evaluate every saved checkpoint in a run (e.g., to generate a learning curve with different metrics):
python scripts/evaluate_checkpoints.py \\\n simulation.model_path=runs/<run_id>/final_model.flax \\\n evaluation.eval_max_steps=2000\n
This script scans the checkpoints/ directory of the specified run and evaluates every .flax file it finds using the model's training morphology.
For a step-by-step walkthrough on using these evaluation phases to reproduce our project results, see the Results & Reproduction Guide.
"},{"location":"api/reproduction/","title":"Results & Reproduction","text":"This guide explains how to access our official training logs and reproduce our results.
Our official training runs, model configurations, and metrics are publicly hosted on Weights & Biases (WandB).
"},{"location":"api/reproduction/#weights-biases-wandb-project","title":"Weights & Biases (WandB) Project","text":"All experiments, final models, and training logs are tracked in our public WandB project:
- Official Runs Table: WandB final-models-v2 Table
This page lists the verified runs with their architecture types, morphology definitions, evaluation metrics, and final model performance.
"},{"location":"api/reproduction/#how-to-reproduce-a-run-from-wandb","title":"How to Reproduce a Run from WandB","text":"Weights & Biases provides a built-in feature to extract the exact parameters and commands used for any given run:
- Open the WandB final-models-v2 Table.
- Click on the name of the run you wish to reproduce to open its detail page.
- In the top-right corner of the run header (next to the run name, not the main workspace header), click the three dots (
...) menu. - Select \"Reproduce run\". This will display the exact command-line arguments and configuration settings used to execute that run.
"},{"location":"api/reproduction/#local-hpc-reproduction-workflow","title":"Local & HPC Reproduction Workflow","text":"To reproduce our training and evaluation phases locally or on an HPC cluster, follow the procedures below.
"},{"location":"api/reproduction/#1-environment-setup","title":"1. Environment Setup","text":"To ensure identical package versions (including JAX, Flax, and MuJoCo), sync your environment using the lockfile:
uv sync --frozen\n
"},{"location":"api/reproduction/#2-training-phase","title":"2. Training Phase","text":"Run the training script using the exact parameters retrieved from WandB's \"Reproduce run\" page or from a downloaded _metadata.yaml file:
uv run python scripts/train.py experiment=my_experiment ppo.learning_rate=0.001 experiment.seed=42\n
"},{"location":"api/reproduction/#evaluation-phases","title":"Evaluation Phases","text":"Reproducing our evaluation results is divided into two distinct phases:
"},{"location":"api/reproduction/#phase-1-determining-the-best-checkpoint","title":"Phase 1: Determining the Best Checkpoint","text":"During training, checkpoints are saved at regular intervals. To determine which of these checkpoints performed the best:
- Evaluate Checkpoints Post-Training: If checkpoint evaluation was not run during training, scan the completed run's checkpoints folder by pointing to the final model path:
uv run python scripts/evaluate_checkpoints.py simulation.model_path=runs/your_run_dir/final_model.flax\n
This script runs deterministic rollouts for every checkpoint in runs/your_run_dir/checkpoints/.
- Locate the Results: The evaluations are saved to:
runs/your_run_dir/metrics/checkpoint_evaluation.csv\n
Analyze this CSV to find the checkpoint iteration with the highest average return or target success rate. This checkpoint will be used for cross-architecture comparisons.
"},{"location":"api/reproduction/#phase-2-comparing-checkpoints-between-architectures","title":"Phase 2: Comparing Checkpoints Between Architectures","text":"Once the best checkpoints for each architecture are identified, they are compared under shared, standardized environments (including fault tolerance checks such as leg amputations).
- Configure the Comparison Models: Open or create an evaluation config file (e.g.,
configs/evaluation/poster.yaml) and add the paths to the best checkpoints:
# configs/evaluation/poster.yaml\nevaluation:\n comparison_models:\n - runs/run_arch_centralized/checkpoints/checkpoint_best.flax\n - runs/run_arch_decentralized/checkpoints/checkpoint_best.flax\n
- Execute the Comparison Script: Run the comparison script using your config:
uv run python scripts/compare_models.py evaluation=poster\n
This script runs multiple sequential evaluation episodes (defined by comparison_num_episodes starting at comparison_base_seed) for every model across the selected morphologies.
-
Analyze Comparison Metrics: The script writes a consolidated CSV file to metrics/model_comparison.csv containing:
-
eval_return: The cumulative return.
approx_max_velocity: The distance covered per step.reached_target: Navigational success rates.arm_0 to arm_4: Active segments per arm (indicating damage/amputations).
This CSV can then be passed to the plotting scripts (e.g., scripts/plots/analyze_comparisons.py) to generate visualization plots. For details on configuration and outputs, see the Analysis & Plotting Guide.
"},{"location":"api/simulation/","title":"Interactive Simulation & Visualization","text":"The simulation pipeline allows you to visualize trained models and observe their behavior under various conditions.
"},{"location":"api/simulation/#overview","title":"Overview","text":"The simulation pipeline is metadata-driven. Training-specific configurations (morphology, arena, environment, etc.) are automatically loaded from the _metadata.yaml file associated with the model checkpoint.
"},{"location":"api/simulation/#basic-simulation","title":"Basic Simulation","text":"To simulate a model in the MuJoCo viewer:
uv run scripts/simulate.py simulation.model_path=runs/your_run/final_model.flax\n
"},{"location":"api/simulation/#amputation-morphology-overrides","title":"Amputation & Morphology Overrides","text":"You can test trained models on different morphologies (e.g., amputating legs) by providing a morphology override. The observations will be automatically padded up to the training morphology's dimensions:
uv run scripts/simulate.py \\\n simulation.model_path=runs/your_run/final_model.flax \\\n simulation.morphology_override=configs/morphology/3_arms.yaml\n
"},{"location":"api/simulation/#video-recording","title":"Video Recording","text":"Recording videos requires the [evaluation] extra:
uv run scripts/simulate.py \\\n simulation.model_path=runs/your_run/final_model.flax \\\n simulation.record_video=true \\\n simulation.max_steps=1000\n
Videos and evaluation metadata are stored in timestamped folders alongside the model: runs/your_run/final_model_evaluations/eval_<timestamp>/simulation.mp4
For batch evaluation, checkpoint analysis, and cross-model architecture comparisons, see the Checkpoint & Model Evaluation Guide.
"},{"location":"api/tracking/","title":"Tracking & Monitoring","text":"This guide explains how to monitor your experiments using Weights & Biases (WandB) and TensorBoard.
"},{"location":"api/tracking/#weights-biases-wandb","title":"Weights & Biases (WandB)","text":"WandB is used for online synchronization and visualization of training metrics.
"},{"location":"api/tracking/#authorization","title":"Authorization","text":"Export your API key in your terminal to enable WandB synchronization:
export WANDB_API_KEY=your_copied_api_key_here\n
Alternatively, you can log in using the CLI:
uv run wandb login\n
"},{"location":"api/tracking/#enabling-tracking","title":"Enabling Tracking","text":"To enable online sync during a training run, set logging.track=true on the command line:
uv run python scripts/train.py logging.track=true\n
You can also configure your project and entity:
uv run python scripts/train.py \\\n logging.track=true \\\n logging.wandb_project_name=\"MyProject\" \\\n logging.wandb_entity=\"my-team\"\n
These can also be set in your configuration YAML file under the logging key.
"},{"location":"api/tracking/#local-monitoring-with-tensorboard","title":"Local Monitoring with TensorBoard","text":"All runs are recorded locally in the runs/ directory (or the directory specified in experiment.base_run_dir). You can view scalars and other metrics with TensorBoard:
tensorboard --logdir runs/\n
Access the interface at http://localhost:6006.
"},{"location":"api/tracking/#cli-exploration-tool","title":"CLI Exploration Tool","text":"For quick diagnostics or to export data to CSV without launching the full TensorBoard UI, you can use the explore_tensorboard.py script:
uv run python scripts/analysis/explore_tensorboard.py runs/your_run_name/\n
See the detailed description in /scripts/analysis/README.md.
"},{"location":"api/tracking/#developer-logging-api","title":"Developer Logging API","text":"For details on the developer API of our internal logging library (how backend routing, checkpoint synchronization, and singleton initialization works), see the Experiment Logger API Guide.
"},{"location":"api/training/","title":"Training Models","text":"This guide covers how to configure and run training experiments for the Brittle Star project using Hydra-based configurations.
"},{"location":"api/training/#configuration","title":"Configuration","text":"The project uses a modular configuration system powered by Hydra. Instead of passing many command-line flags, you select and override configuration groups.
For a detailed guide on the structure, validation, and usage of our Hydra configuration files, see the Brittle Star Configuration System Guide.
"},{"location":"api/training/#creating-a-custom-experiment","title":"Creating a Custom Experiment","text":" -
Create a new experiment file: Create a file at configs/experiment/my_experiment.yaml. You can copy an existing one as a template:
cp configs/experiment/base.yaml configs/experiment/my_experiment.yaml\n
-
Edit configs/experiment/my_experiment.yaml to set your experiment parameters:
# @package _global_\nexperiment:\n exp_name: \"my_custom_run\"\n seed: 42\n
"},{"location":"api/training/#training-execution","title":"Training Execution","text":"To start a training run with the default settings defined in configs/main_config.yaml:
uv run python scripts/train.py\n
"},{"location":"api/training/#using-a-custom-experiment-configuration","title":"Using a Custom Experiment Configuration","text":"To run with your custom experiment file:
uv run python scripts/train.py experiment=my_experiment\n
uv run python scripts/train.py ppo.learning_rate=0.001 ppo.num_envs=32 logging.track=true\n
"},{"location":"api/training/#evaluation-during-training","title":"Evaluation During Training","text":"By default, the trainer saves checkpoints but does not evaluate them. To enable automatic headless evaluation of every saved checkpoint, set evaluation.evaluate_checkpoints=true:
uv run python scripts/train.py evaluation.evaluate_checkpoints=true\n
"},{"location":"api/training/#reproducing-experiments","title":"Reproducing Experiments","text":"For detailed steps on how to reproduce training runs, locate run configuration metadata, or reproduce our experiments using Weights & Biases (WandB), see the Results & Reproduction Guide.
For more details on evaluation metrics and comparison tools, see Checkpoint & Model Evaluation.
For more details on tracking your experiments, see Tracking & Monitoring.
"},{"location":"configs/","title":"Brittle Star Configuration System","text":"This project uses Hydra for a modular, hierarchical, and strictly-typed configuration system.
"},{"location":"configs/#core-concepts","title":"Core Concepts","text":" - Composition over Inheritance: Instead of one giant config file, the configuration is composed of small, domain-specific modules (PPO settings, architecture, morphology, etc.).
- Strict Typing: Every configuration is validated against a Python dataclass schema (
ConfigStore). Misspelled keys throw a ConfigAttributeError immediately. - CLI Swapping: You can swap entire modules or override individual values from the command line without touching code.
"},{"location":"configs/#directory-structure","title":"Directory Structure","text":" main_config.yaml: The root entry point defining the default composition.experiment/: High-level experiment settings (seed, device).logging/: WandB and checkpointing configuration.ppo/: PPO training hyperparameters.architecture/: Polymorphic network architectures (centralized vs. decentralized).morphology/: Physical robot definitions (number of segments, amputations).arena/: Environment physics and visual settings.environment/: Task-specific settings (Directed Locomotion, Light Escape).
"},{"location":"configs/#common-commands","title":"Common Commands","text":""},{"location":"configs/#local-debugging","title":"Local Debugging","text":"Run a quick test with minimal iterations:
python scripts/train.py experiment=dev_test ppo=fast\n
"},{"location":"configs/#swapping-architectures-or-morphologies","title":"Swapping Architectures or Morphologies","text":"Test a decentralized controller on a 3-arm robot:
python scripts/train.py architecture=decentralized morphology=3_arms\n
"},{"location":"configs/#hpc-production","title":"HPC Production","text":"Run stable PPO with WandB enabled (HPC submission scripts handle the hydra.run.dir redirection):
python scripts/train.py ppo=stable logging=wandb_enabled\n
"},{"location":"configs/#dry-run-validation","title":"Dry-Run Validation","text":"Check if your configuration is valid without starting the simulation:
python scripts/train.py --cfg job\n
"},{"location":"configs/#developer-notes","title":"Developer Notes","text":" - Adding a new group: Create a subdirectory in
configs/ and register the new dataclass in src/brittle_star_project/configs/register_configs.py. - Typo Catching: If you see a
ConfigAttributeError, check for typos in your YAML keys or CLI overrides. - Output Redirection: We use
experiment.base_run_dir to configure where logs and models are stored (defaults to runs/). - To change it locally:
python scripts/train.py experiment.base_run_dir=/path/to/custom/dir - On HPC, ensure this points to a fast scratch storage.
"},{"location":"design/actor-critic/","title":"Actor-Critic Architecture","text":"To process observations into actions, our controllers utilize an Actor-Critic architecture. Because we use Proximal Policy Optimization (PPO), the pipeline fundamentally requires separate networks for the policy (Actor) and the value estimation (Critic).
Centralized Architecture (Baseline)
This pipeline treats the agent as a single entity and uses standard Proximal Policy Optimization (PPO).
- Centralized Actor: Composed of two chained MLPs (Sensor \\(\\rightarrow\\) Motor) passing a hidden state between them. The centralized sensor receives the concatenated global state vector of all limbs at once and processes it into a hidden state. The centralized motor receives this hidden state and outputs the joint offsets for all actuators simultaneously. This is mathematically equivalent to using one large MLP with hidden layers, but splitting makes the implementation easier by allowing us to reuse the same components for the decentralized modules.
- Centralized Critic: Composed of two sequential MLPs (Feature Extractor \\(\\rightarrow\\) Critic). Because PPO evaluates the state-value function, this network only receives the concatenated global state vector (no actions). It outputs a single scalar estimating the expected future reward for the entire agent.
Our policy and value networks use separate input networks/feature extractors as advised by the SEL3 course assistants and the blog. For continuous actions this should allow better learning at a small cost.
graph TD\n Obs([Global Observation])\n\n Sens[Sensor]\n Act[Motor]\n OutAct([Action Distribution<br/>mean, log_std])\n\n Feat[Feature extractor]\n Crit[Critic]\n OutCrit([Value Estimate<br/>scalar])\n\n Obs --> Sens\n Obs --> Feat\n\n Sens -->|\"Hidden state\"| Act\n Feat -->|\"Hidden state\"| Crit\n\n Act --> OutAct\n Crit --> OutCrit
Decentralized Architecture
This pipeline utilizes the \"Centralized Training with Decentralized Execution\" principle, specifically the NerveNet-MLP variant.
- Decentralized Actor, split into three distinct models:
- Sensor: A local model at each node. It receives its local state plus the goal vector directly, processing them into an initial hidden state.
- Propagator: Nodes synchronously compute and exchange messages with connected neighbors for \\(N\\) steps to update their hidden states. See communication.md for details.
- Motor: A local model uses its final updated hidden state to output the joint offset strictly for its own actuator.
- Centralized Critic: Composed of two sequential MLPs (Feature Extractor \\(\\rightarrow\\) Critic). During training, it acts globally by taking the concatenated state vectors from all sensors to output a single, global state-value scalar evaluating the entire agent's pose.
To keep the implementation simple, we should use one critic per node in our architecture, but only a single, global critic for all nodes at once, for the following reasons:
- Credit Assignment Problem (Ha, 2017): The MuJoCo simulator provides an overall reward based on the brittle star movement progression, e.g. total distance travelled. Using an isolated critic for each node in the network would not allow to determine which local action contributed to the global success. A global critic solves this by evaluating the combined state of the agent at once.
- Implementation simplicity: Building a second decentralized message-passing graph for the critic (NerveNet-2) would require more coding. Using a standard MLP that concatenates all raw input vectors is much easier to program while mathematically equivalent.
graph TD\n Obs([Local Observation])\n\n Sens[Sensor]\n Prop[Propagator]\n Feat[Feature extractor]\n\n Mot[Motor]\n Crit[Critic]\n\n OutMot([Action Distribution<br/>mean, log_std])\n OutCrit([Value Estimate<br/>scalar])\n\n Obs --> Sens\n Sens -->|\"Hidden state\"| Prop\n Obs --> Feat\n\n Prop -->|\"Hidden state\"| Mot\n\n\n Feat -->|\"Hidden state\"| Crit\n\n Mot --> OutMot\n Crit --> OutCrit\n\n Prop -.->|\"message passing\"|Prop
"},{"location":"design/actor-critic/#implementation-details-network-depth","title":"Implementation Details (Network Depth)","text":"Inspired by: PPO Implementation Details
The MLPs used in both pipelines are defined with specific hidden layer configurations to balance learning capability and computational cost. As of right now, though this might change as we make progress in our experiments, we use:
- Input Networks (Sensors & Feature Extractors): These networks map the raw state inputs to internal hidden states. They are configured as standard dense networks with 2 hidden layers of 64 nodes each (
[64, 64]) and utilize tanh activation functions. - Output Networks (Motors, Actors & Critics): The final output models are intentionally kept shallow. The Actor directly projects the hidden state to a continuous action distribution (
mean and log_std) using a single dense output layer (zero hidden layers) initialized orthogonally. The Critic functions similarly, mapping the hidden representation to a single scalar value.
Note: For the continuous action distributions outputted by the Motor, we explicitly use mean and log_std as advised by previous research to maintain learning stability.
References
- Ha, D. (2017, October 29). A Visual Guide to Evolution Strategies. \u5927\u30c8\u30ed \u30fb Machine Learning. https://blog.otoro.net/2017/10/29/visual-evolution-strategies/
- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. \u2018Proximal Policy Optimization Algorithms\u2019. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. https://doi.org/10.48550/arXiv.1707.06347.
- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. \u2018NerveNet: Learning Structured Policy with Graph Neural Networks\u2019. Conference paper presented at International Conference on Learning Representations. 15 February 2018. https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613.
"},{"location":"design/communication/","title":"Communication scheme (Message Passing)","text":"Remember our research question:
\"What is the impact of different levels of controller modularity on learning speed, coordination, and fault tolerance (e.g. amputations) in brittle-star-like robots trained with Reinforcement Learning?\"
To test decentralized modularity (such as arm-level or segment-level controllers), the various modules must be able to communicate with each other to achieve coordinated locomotion. This is accomplished through message passing in a Graph Neural Network (GNN)-like architecture. Two prominent communication styles from the literature are N-step NerveNet (Wang et al., 2018) and bottom-up top-down Shared Modular Policies (Huang et al., 2020).
We have chosen to apply one uniform communication style across all modular architectures, specifically opting for N-step NerveNet.
"},{"location":"design/communication/#rationale","title":"Rationale","text":"Initially, our idea was to equip arm-level controllers with NerveNet message passing and segment-level controllers with SMP. However, we evaluated that this introduces a threat to the validity of our research question. If we observe differences in performance, it would be impossible to determine whether the variance is caused by the level of modularity, or by the difference in the message passing scheme. To purely compare modularity, the communication scheme style must remain constant.
Second, we decided that NerveNet is a better fit for our research. The morphology of our brittle star contains cycles at the decentralized level (e.g., a ring of segments or arms around the body). NerveNet has proven to be robust for arbitrary structures, including graphs with cycles. SMP inherently expects a tree structure for its bottom-up and top-down pass. Applying SMP to a ring structure requires a workaround to break that cycle.
"},{"location":"design/communication/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Choosing NerveNet introduces a scalability issue as the morphology grows. In NerveNet, a message advances only one segment or node per propagation step. When dealing with long arms (e.g., > 5 segments), this requires a large number of propagation steps to transmit information from one tip of an arm to another.
If we were to use SMP instead - which is possible - the inner states of nodes are shared across the entire graph in just two passes. For very large or long morphologies, this would be much more scalable.
By rejecting SMP, we accept that our model might learn slower or require more computational power for highly segmented, extended morphologies.
References
- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. \u2018NerveNet: Learning Structured Policy with Graph Neural Networks\u2019. Conference paper presented at International Conference on Learning Representations. 15 February 2018. https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613.
- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. \u2018One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control\u2019. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. https://doi.org/10.48550/arXiv.2007.04976.
"},{"location":"design/controllers/","title":"Levels of modularity and topology","text":"The brittle star can be controlled at different levels. A monolithic controller processes all inputs and outputs at once, whereas modular controllers divide the brains across the body, inspired by the biology of brittle stars.
We define four architectures to compare:
- Centralized, monolithic: A single Multi Layer Perceptron per robot that receives all observations and outputs all actions.
- Fully connected arm-level: Each arm contains an MLP that processes the inputs for that arm, an MLP that processes the communicated inner-states, and an MLP that outputs the actions for that arm. One policy for these MLPs is shared across the arms. The controllers in each arm are connected to each other and form a fully connected graph. There is no central disk, but the controllers are fully connected.
- Ring arm-level: Identical setup to the fully connected arm-level, but the controllers are connected in a ring structure. This setup is considered less centralized than the fully connected graph.
- Segment-level: Each segment contains the three MLPs discussed above. The base segments, attached to the body, form a ring structure, with the remaining segments attached as extended \"strings\". Segments can only communicate with segments that are physically connected to it.
"},{"location":"design/controllers/#rationale","title":"Rationale","text":"To fairly compare decentralized modularity against centralized control, the decentralized models should not be allowed to contain a central organ acting as a bottleneck or coordinator. By removing the central disk in the decentralized models and replacing it with a ring topology, we closely approximate the biological reality of the brittle star and test a decentralized morphology.
The fully connected graph functions as an intermediate step in between a fully centralized and a decentralized ring. We use it to test whether our models scale to more complex structures.
"},{"location":"design/input_action_spaces/","title":"Input (state) and output (action) spaces","text":"To effectively learn locomotion and navigation, the agent requires a well-defined observation space (inputs) and action space (outputs). The control models map these observations directly to physical movements.
Inputs (state space)
The observation space provides the agent with its current physical state and its navigational objective. With a decentralized control architecture in mind, we divide these inputs into global and local states.
Global inputs, always broadcasted to all nodes:
- Vertical orientation/tilt: A single, simplified metric representing the tilt/vertical alignment of the agent's central body/disk, a.k.a. the deviation from the global Z-axis. Its value is derived from the environment's raw disk rotation 3D vector \\([roll, pitch, yaw]\\): $$ tilt = sqrt(roll^2 + pitch^2) $$
- Goal vector: A 2D unit vector representing the egocentric direction to the target. A value of \\([1.0, 0.0]\\) indicates that the target is directly in front of the agent (angle 0).
Local inputs, routed directly to specific nodes:
- Joint positions: The current angles of all joints within the morphology.
- Joint velocities: The current angular velocities of the joints.
- Joint actuator forces: The physical forces currently exerted at each specific joint.
- Segment contact: These values indicate whether each physical segment of the agent is currently touching the ground.
Outputs (action space)
The action space defines how the agent interacts with the environment.
- Joint offsets: absolute target positions (offsets) for the joints, i.e. the exact angle the joint should move to.
"},{"location":"design/input_action_spaces/#normalization-and-scaling","title":"Normalization and Scaling","text":"Both the input (observation) and output (action) spaces are rescaled to the range \\([-1, 1]\\).
For the input space, all raw physical values (angles, velocities, forces, distances) are normalized based on their defined physical bounds. If a value exceeds these bounds during simulation, it is clipped to the \\([-1, 1]\\) range.
For the output space, the neural network's tanh-activated outputs (which naturally fall in \\([-1, 1]\\)) are linearly mapped to the physical joint limits defined in the robot's morphology.
"},{"location":"design/input_action_spaces/#rationale","title":"Rationale","text":"When designing the state space, we must ask: Could a human operator perform this task given only these inputs?
- Inclusion of Joint Velocities: Because our control models do not inherently possess memory of previous timesteps, providing only the joint position is insufficient to determine the direction a limb is currently moving. By explicitly including joint velocities, the agent can immediately infer momentum and movement direction without needing to memorize past states.
- Absolute Joint Offsets: The physical Brittle Star robot relies on servo motors (if we were to build this simulated robot), which are inherently position-controlled devices. (Continuous rotation servos exist, but they are less commonly used for joints.) If our network outputted continuous torques (forces), a significant portion of the reinforcement learning process would be wasted on learning low-level PID control dynamics (i.e., how much force to apply to hold a position). Abstracting this away forces the learning algorithm to focus entirely on higher-level gait generation and locomotion.
- Simplified vertical orientation: We drop the full 3D spatial rotation and angular velocity arrays in favor of a single vertical orientation metric (tilt). For a brittle star moving accross a flat plane, this metric is sufficient for the agent to sense if it is losing balance or flipping over.
- Force representation: We strictly retain the joint actuator forces and drop the more generic actuator force. Forces that are explicitly tied to individual joints are significantly easier to route into decentralized, local limb nodes, which is necessary for our message-passing architecture.
- Goal Vector (Distance + Angle): Providing only the scalar \"distance to the goal\" as an input is akin to blindfolding the robot and asking it to find a target by playing \"hot or cold.\" By providing a full vector, the agent knows exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion.
The goal representation is explicitly divided into a directional vector and a scalar distance. Keeping the direction as a normalized unit vector bounds the values to the \\([-1, 1]\\) range, which stabilizes neural network training. Providing only a scalar \"distance to the goal\" would force the agent to learning localized searching behaviors (e.g. random walks or spiraling) to deduce the direction, drastically increasing the difficulty of the learning task.
NOTE: We later dropped the \"distance to vector\", switching to only a direction as the input. Our reasoning is the agent should always move towards the goal (it should not learn to stop at the goal), which allows for this simplification that decreases the model input size.
The environment provides a raw unit_xy_direction_to_target (global), which we transform into a calculated robot_direction_to_target (egocentric) before passing it to the MLPs. This vector consists of the X and Y direction, where a value of \\([1.0, 0.0]\\) (mapping to an angle of \\(0\\)) means the robot is facing directly towards the target. - Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait transitions. - Zero-Centered Rescaling (\\([-1, 1]\\)): Using a zero-centered range is standard best practice for continuous control tasks. It provides several mathematical and physical advantages: - Improved Gradient Flow: Neural networks optimize faster when inputs are zero-centered. If all inputs were positive (e.g., \\([0, 1]\\)), the gradients during backpropagation would be forced to the same sign, causing inefficient \"zig-zag\" weight updates. - Meaningful Neutral State: In robotics, \\(0.0\\) naturally represents a resting state (zero velocity, centered position, no force). In a \\([-1, 1]\\) system, this physical rest maps to a neutral \\(0.0\\) signal in the network. This also correctly communicaties a \"neutral/dead\" signal for amputated limbs that are padded with \\(0.0\\) values.
Specifically, we do not include some available inputs:
- Global position: Absolute spatial coordinates can cause the agent to overfit to a specific coordinate frame or map, rather than learning general, adaptable locomotion strategies.
"},{"location":"design/input_action_spaces/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Alternative state and action formulations include:
- Torque-based continuous control: In many continuous control tasks (like standard MuJoCo benchmarks), actions represent continuous torques applied to joints. While this provides more granular, low-level physical control, it heavily complicates training and does not align well with the physical reality of servo-driven hardware.
- Recurrent Neural Networks (RNNs) / Frame Stacking: Instead of explicitly passing velocities in the state space, the network could infer momentum by observing a history of past states. Using RNNs or frame stacking allows the agent to build an internal memory of movement. However, this significantly increases architectural complexity and training time compared to explicitly providing the velocity data.
- Scalar Goal Distance: Giving the agent only the scalar distance to the target would force it to learn a localized searching behavior (e.g., spiraling or random walks) to determine the correct direction. While biologically plausible for simpler organisms following chemical gradients, it drastically increases the difficulty of the learning task.
- \\([0, 1]\\) Rescaling: While some domains (like computer vision) use \\([0, 1]\\) scaling, it is generally avoided in robotics. Scaling to \\([0, 1]\\) would mean that a resting joint (velocity = 0) maps to an input of \\(0.5\\). This constant positive bias forces the network to waste capacity learning to ignore or subtract this baseline signal just to stand still. Furthermore, it breaks the \"dead signal\" interpretation of zero-padding used for amputations.
"},{"location":"design/input_action_spaces/#mujoco","title":"MuJoCo","text":"This is what the filtered input vectors look like in MuJoCo, with \\(J\\) joints and \\(S\\) segments:
joint_position: shape=(J,), dtype=float64joint_velocity: shape=(J,), dtype=float64joint_actuator_force: shape=(J,), dtype=float64segment_contact: shape=(S,), dtype=float64robot_direction_to_target: shape=(2,), dtype=float64, egocentricdisk_z_tilt: shape=(1,), dtype=float64, derived from disk_rotation
This brings the entire input space down to \\(3J + S + 4\\) float64's, compared to \\(4J + S + 15\\) float64's for the unfiltered inputs.
For reference, these are all the inputs that are available in the MuJoCo environment:
obs keys: ['joint_position', 'joint_velocity', 'joint_actuator_force', 'actuator_force', 'disk_position', 'disk_rotation', 'disk_linear_velocity', 'disk_angular_velocity', 'tendon_position', 'tendon_velocity', 'segment_contact', 'unit_xy_direction_to_target', 'xy_distance_to_target']\n\nraw observations dict:\n{'joint_position': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'joint_velocity': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'joint_actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'disk_position': array([0. , 0. , 0.11]),\n 'disk_rotation': (0.0, -0.0, 0.0),\n 'disk_linear_velocity': array([0., 0., 0.]),\n 'disk_angular_velocity': array([0., 0., 0.]),\n 'tendon_position': array([], dtype=float64),\n 'tendon_velocity': array([], dtype=float64),\n 'segment_contact': array([0., 0., 0., 0., 0., 0.]),\n 'unit_xy_direction_to_target': array([-0.95333378, -0.30191837]),\n 'xy_distance_to_target': array([3.])}\n\n(shapes)\njoint_position: shape=(12,), dtype=float64, size=12\njoint_velocity: shape=(12,), dtype=float64, size=12\njoint_actuator_force: shape=(12,), dtype=float64, size=12\nactuator_force: shape=(12,), dtype=float64, size=12\ndisk_position: shape=(3,), dtype=float64, size=3\ndisk_rotation: shape=(3,), dtype=float64, size=3\ndisk_linear_velocity: shape=(3,), dtype=float64, size=3\ndisk_angular_velocity: shape=(3,), dtype=float64, size=3\ntendon_position: shape=(0,), dtype=float64, size=0\ntendon_velocity: shape=(0,), dtype=float64, size=0\nsegment_contact: shape=(6,), dtype=float64, size=6\nxy_distance_to_target: shape=(1,), dtype=float64, size=1\n
"},{"location":"design/learning_algorithm/","title":"Reinforcement Learning Algorithm","text":"To control the continuous action space (the joints of the robot) based on sensor data, we require a reliable Reinforcement Learning (RL) algorithm or optimization strategy.
We have chosen Proximal Policy Optimization (PPO) (Schulman et al., 2017).
"},{"location":"design/learning_algorithm/#rationale","title":"Rationale","text":"PPO is an on-policy algorithm known for its stability and robustness (safe training without excessive variance). More importantly, it requires relatively little hyperparameter tuning compared to other algorithms. Since NerveNet was successfully trained using PPO (Wang et al., 2018), selecting PPO significantly reduces the risk of convergence failures.
"},{"location":"design/learning_algorithm/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Alternative learning algorithms include:
- Twin Delayed DDPG (Fujimoto et al., 2018): TD3 is a strong off-policy alternative used in the SMP paper (Huang et al., 2020). It is highly sample-efficient and reportedly excels at zero-shot adaptations. However, this approach would be more complex and error-prone than with PPO.
- Evolution strategies (ES): Evolution strategies are useful for optimizing Central Pattern Generators (CPGs), e.g. CMA-ES, OpenAI-ES. While this method is easier to distribute and parallelize, ES typically scales worse with exceptionally large observation spaces compared to gradient-based RL methods like PPO.
References
- Fujimoto, Scott, Herke Hoof, and David Meger. \u2018Addressing Function Approximation Error in Actor-Critic Methods\u2019. Proceedings of the 35th International Conference on Machine Learning, 3 July 2018, 1587-96. https://proceedings.mlr.press/v80/fujimoto18a.html.
- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. \u2018One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control\u2019. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. https://doi.org/10.48550/arXiv.2007.04976.
- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. \u2018Proximal Policy Optimization Algorithms\u2019. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. https://doi.org/10.48550/arXiv.1707.06347.
"},{"location":"design/reward_function/","title":"Reward function and observation space","text":"The robot needs to know whether its movements contribute to the ultimate goal of locomotion towards a target. Sensor inputs must be distributed fairly to guarantee an objective comparison between different architectures.
- The distance from the robot to the target and/or the light intensity are treated as global inputs.
- Positions and joints, which are normalized to floating-point values between 0 and 1, are considered local inputs.
- The reward function is centered around minimizing the distance to the goal or maximizing the movement towards the goal within a finite number of timesteps \\(T\\).
- To motivate efficient movement, the amount of timesteps taken to reach the goal will be used as penalty.
- An extra penalty based on movement relative to the current step and the previous is used to penalize a movement away from the target.
"},{"location":"design/reward_function/#from-reward-to-ppo","title":"From reward to PPO","text":"The resulting reward is passed to our PPO library. Our critic network (value function) predicts how good our eventual reward will be for the current state, this value is combined with the reward from the reward function to get advantages. These advantages are then used to calculate the losses to update both our critic and actor pipeline.
"},{"location":"design/reward_function/#rationale","title":"Rationale","text":"Using a light source (or a gradient) is biologically plausible for many simple organisms. By normalizing all signals between 0 and 1, PPO training is highly stabilized. The timesteps must be finite to reset the environment in a timely manner if the policy gets stuck in a local minimum.
"},{"location":"design/reward_function/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Providing global information to all individual decentralized segments can be considered biologically cheating or practically infeasible once the robot would be physically built. Some sensory input cannot be put in each joint, for example.
The alternative is to provide the global input to the outermost segments of the arms, or a specific set of segments assigned with this functionality. The network would then have to learn to propagate this signal throughout the body via message passing. While biologically more accurate, this drastically complicates the learning process. We have written this down as potential future research.
"},{"location":"scripts/analysis/","title":"Experiment Analysis Tools","text":"This directory contains scripts for post-processing and analyzing experiment results, including TensorBoard logs and saved model weights.
"},{"location":"scripts/analysis/#scripts","title":"Scripts","text":""},{"location":"scripts/analysis/#1-explore_tensorboardpy","title":"1. explore_tensorboard.py","text":"A CLI tool to summarize TensorBoard tfevents files without a GUI.
Key Features: - Displays last values, min, max, and step counts for all scalar metrics. - Calculates total run duration and estimated completion percentage. - Exports granular scalar data to CSV for analysis in Excel/Pandas.
Usage:
# General usage\npython explore_tensorboard.py <run_directory>\n\n# Exporting data\npython explore_tensorboard.py <run_directory> --csv data.csv\n
Requirements: - pandas - tensorboard - tensorflow-cpu (or tensorflow)
"},{"location":"src/experiment_logger/","title":"Experiment Logger","text":"A standardized, unified interface for logging experiments across multiple backends (WandB, TensorBoard, and Local Disk).
This library is designed to be a standalone package that decouples the logging logic from the core training routines in the brittle_star_project.
"},{"location":"src/experiment_logger/#quick-start","title":"Quick Start","text":"The recommended way to use the logger is through the get_logger() singleton:
from experiment_logger import UnifiedLogger, get_logger\n\n# Initialize at the start of your script (e.g., in train.py)\nlogger = UnifiedLogger(\n run_name=\"my_experiment_run\",\n config={\"learning_rate\": 3e-4},\n project_name=\"MyProject\",\n base_dir=\"runs\",\n use_wandb=True\n)\n\n# In other files, retrieve the initialized singleton:\n# logger = get_logger()\n\n# Log metrics (Scalar values, numpy scalars, or JAX types)\nlogger.log({\"loss\": 0.5, \"accuracy\": 0.98}, step=100)\n\n# Standard logging (Mirrored to disk and stdout)\nlogger.info(\"Training started\")\nlogger.warning(\"Learning rate is very high\")\n\n# Save checkpoints (Automatically synced to WandB as artifacts)\nlogger.save_checkpoint(params, step=5000)\n
"},{"location":"src/experiment_logger/#logger-classes","title":"Logger Classes","text":""},{"location":"src/experiment_logger/#unifiedlogger","title":"UnifiedLogger","text":"The full suite for production training. It manages: - WandB: Syncs metrics and uploads model checkpoints as artifacts. - TensorBoard: Writes events for local visualization. - Local Disk: Stores metrics in metrics.yaml and textual logs in run.log.
"},{"location":"src/experiment_logger/#simplelogger","title":"SimpleLogger","text":"A zero-dependency fallback that uses standard Python print() statements. Use this for standalone testing or minimal environments where you don't need persistent monitoring.
from experiment_logger import SimpleLogger\nlogger = SimpleLogger(run_name=\"test_run\")\n
"},{"location":"src/experiment_logger/#api-features","title":"API Features","text":""},{"location":"src/experiment_logger/#loggerprogress_bariterable-kwargs","title":"logger.progress_bar(iterable, **kwargs)","text":"A smart wrapper around tqdm that automatically detects its environment. - Interactive Terminal: Displays a normal progress bar. - Non-Interactive (HPC): Automatically disables the bar to prevent log file bloat in slurm.out.
"},{"location":"src/experiment_logger/#loggerlog_non_interactivemsg-str","title":"logger.log_non_interactive(msg: str)","text":"Prints a message only when running in non-interactive environments. Useful for high-level progress tracking (e.g., \"Epoch 5 Complete\") without interactive noise.
"},{"location":"src/experiment_logger/#loggersave_checkpointparams-step-prefixcheckpoint","title":"logger.save_checkpoint(params, step, prefix=\"checkpoint\")","text":"Saves model parameters using Flax serialization. - Local Location: runs/<run_name>/checkpoints/ - WandB Logic: Automatically uploads the .flax file as a model artifact for lineage tracking.
"}]}
\ No newline at end of file
+{"config":{"lang":["en"],"separator":"[\\s\\-]+","pipeline":["stopWordFilter"],"fields":{"title":{"boost":1000.0},"text":{"boost":1.0},"tags":{"boost":1000000.0}}},"docs":[{"location":"","title":"Documentation","text":"Welcome to the Brittle Star project documentation. This codebase contains the implementations and research for the scientific evaluation of controller modularity in brittle-star-like robots trained using Reinforcement Learning.
For the core codebase, scripts, and contribution history, visit our GitHub Repository.
"},{"location":"#core-requirements-guides","title":"Core Requirements & Guides","text":" - Installation Instructions: Steps to set up your development environment locally or in a devcontainer using
uv, including GPU configuration. For High-Performance Computing (HPC) setup details, see the HPC Guide. - How to Run Experiments: A complete guide on running training jobs, setting custom hyperparameters, and overriding config options using Hydra.
- Results & Reproduction: Guide on how to access our public WandB training runs table and reproduce our training and evaluation phases (determining the best checkpoint vs. comparing architectures).
- Contribution Guidelines: Standards, rules, and best practices for developing and adding code to the repository.
- Repository Structure: Overview of the directories and files within the codebase.
"},{"location":"#repository-structure","title":"Repository Structure","text":".\n\u251c\u2500\u2500 configs/ # Hydra configuration files (YAML)\n\u251c\u2500\u2500 docs/ # Comprehensive documentation and API guides\n\u251c\u2500\u2500 runs/ # Default output directory for Hydra and training artifacts\n\u251c\u2500\u2500 scripts/ # High-level entrypoints for training, simulation, and evaluation\n\u251c\u2500\u2500 src/\n\u2502 \u251c\u2500\u2500 brittle_star_project/ # Core library and environment logic\n\u2502 \u2502 \u251c\u2500\u2500 evaluation/ # Checkpoint evaluation, rollout logic, and metrics persistence\n\u2502 \u2502 \u2514\u2500\u2500 trainers/ # Training loop implementations (e.g., PPO)\n\u2502 \u2514\u2500\u2500 experiment_logger/ # Standalone logging package\n\u2514\u2500\u2500 tests/ # Unit and integration tests\n
"},{"location":"#design-architecture-design","title":"Design & architecture (/design)","text":"If you are interested in the \"why did you do it like this?\"
- Actor-Critic Architecture: Description of the actor-critic pipeline.
- Communication Scheme: Message propagation, Nerve-Net style.
- Modularity & Topology: Macroscopic brain topology, centralized, arm-level, segment-level.
- Input & Action Spaces: Description of the model's input and output.
- Reinforcement Learning Algorithm: RL techniques, i.e. PPO.
- Reward Function & Observation Space: Goals, fitness tracking, and reward structures.
"},{"location":"#api-reference-api","title":"API reference (/api)","text":"If you are interested in the \"how do I use it?\"
- Brittle Star Environment: MuJoCo environment interaction and configuration.
- Training Models: How to configure and run experiments.
- Tracking & Monitoring: Setting up WandB and TensorBoard to monitor runs.
- Checkpoint & Model Evaluation: Evaluating checkpoints and comparing fault tolerance.
- Interactive Simulation & Visualization: Visualizing models in the MuJoCo viewer or rendering simulation videos.
- Analysis & Plotting Tools: Comparing checkpoints and generating plots.
- Results & Reproduction: Accessing WandB results and running reproduction pipelines.
"},{"location":"CONTRIBUTING/","title":"Contribution Guidelines","text":"This document outlines the contribution protocols for the scientific software engineering project focusing on bio-inspired control architectures for brittle-star-like robots. The primary objective of this project is to produce scientific insight, rather than a commercial product.
"},{"location":"CONTRIBUTING/#1-scientific-context-methodology","title":"1. Scientific Context & Methodology","text":" - Research Focus: The goal is to study how controller modularity affects learning speed, coordination, and fault tolerance in brittle-star locomotion.
- Hypothesis-Driven Design: Clear hypotheses must dictate a structured methodology and rigorous evaluation. All design decisions must be formally documented prior to implementation.
- Scaffolding Approach: Development must start with simple setups before progressively increasing the complexity of environments and morphologies.
- Evaluation of Results: Negative results possess scientific validity when thoroughly analyzed. If a controller fails to learn locomotion, providing a comprehensive analysis of the failure is considered a strong scientific contribution.
- Reproducibility: Contributors must utilize fixed library versions. Configuration systems (such as json, gin, or yaml) must be employed to ensure reproducible runs.
"},{"location":"CONTRIBUTING/#2-clean-code-code-quality","title":"2. Clean Code & Code Quality","text":"Code readability is paramount, as code is read far more frequently than it is written.
- Naming Conventions: Variables and functions must utilize consistent, intention-revealing names. A long, descriptive name is strictly preferred over a short name accompanied by a comment.
- Function Design: Functions must be modular and adhere to the single responsibility principle. Arguments must be minimized, and boolean flag arguments controlling behavior should be avoided.
- Commenting: Code must document the \"how,\" while comments are strictly reserved for documenting the \"why\". Commented-out code is prohibited and must be deleted via version control.
- YAGNI: Contributors must adhere to the \"You Aren't Gonna Need It\" (YAGNI) principle and actively avoid premature optimization.
- Notebooks: Jupyter Notebooks are strictly limited to quick prototyping, tutorials, demonstrations, or post-processing analysis. They are explicitly forbidden for general software development because they discourage modularity.
"},{"location":"CONTRIBUTING/#3-version-control-repository-structure","title":"3. Version Control & Repository Structure","text":" - Git Practices: Commits must be frequent and small. Each commit should relate to exactly one piece of functionality.
- Branching Strategy: The
dev branch serves as the integration branch for pushing and merging code. Only stable releases may be pushed to the main branch. - Artifact Management: Data files, trained models, and large datasets must never be committed directly to Git. Git Large File Storage (LFS) must be used for tracking large files. All developers must have
git-lfs installed locally (see DEVELOPMENT.md for setup). - Repository Layout: The repository must maintain the following core directories:
src/ for algorithms, env/ for MuJoCo wrappers, config/ for experiment configurations, experiments/ for scripts, docs/ for Doxygen or ReadTheDocs documentation, and tests/ for unit tests.
"},{"location":"CONTRIBUTING/#4-architecture-tooling","title":"4. Architecture & Tooling","text":" - Algorithms & Frameworks: Proximal Policy Optimization (PPO) is the recommended baseline algorithm. CleanRL should be used as a starting point and adapted for continuous action spaces. All Artificial Neural Network (ANN) controller architectures must be implemented using Flax.
- Simulation: The simulation environment utilizes a MuJoCo brittle star. XML MuJoCo structures must remain realistic and respect morphological constraints.
- Experiment Tracking: Weights & Biases (wandb) must be utilized for tracking and logging all experiments.
- Code Styling: All code must conform to the chosen style guide (Google standard). This is enforced via
uv using ruff and pre-commit hooks.
"},{"location":"CONTRIBUTING/#5-ai-assisted-development-code-review","title":"5. AI-Assisted Development & Code Review","text":"This project supports AI-assisted development to enhance productivity, but contributors must take full responsibility for all AI-generated outputs.
- Self-Review Requirement: Contributors must thoroughly self-review all AI-assisted code, documentation, and configurations before requesting peer review. This includes verifying correctness, adherence to project standards, scientific validity, and integration with existing code.
- Quality Standards: AI-generated content must meet the same rigorous standards as manually written code, including proper testing, documentation, and alignment with the scientific methodology outlined in Section 1.
- Available Skills: This project provides specific AI skills for common tasks (located in
.agents/skills/), including linting and testing workflows. Contributors should leverage these skills to maintain consistency and quality. - Transparency: When using AI assistance for complex algorithmic decisions or scientific design choices, contributors should document the rationale in commit messages or code comments where appropriate.
"},{"location":"DEVELOPMENT/","title":"Development Guide","text":"This guide outlines how to set up the development environment for this project, prioritizing reproducible builds, environment parity, and cross-hardware compatibility.
"},{"location":"DEVELOPMENT/#reproducibility-uv","title":"Reproducibility & uv","text":"This project uses uv to manage dependencies and virtual environments. The uv.lock file is the absolute source of truth for package versions and must always be committed.
"},{"location":"DEVELOPMENT/#source-of-truth","title":"Source of Truth","text":" - Never modify
uv.lock manually. - To add a dependency, run
uv add <package>. - To update dependencies, run
uv lock --upgrade. - To sync your environment with the lockfile, run
uv sync --frozen.
"},{"location":"DEVELOPMENT/#git-lfs-critical","title":"Git LFS (Critical)","text":"All developers must have Git LFS installed locally. This repository tracks model weights (.pt, .safetensors, etc.), recordings (.mp4), and datasets using Git LFS.
- Setup: Run
git lfs install after cloning this repository. If you are using the .devcontainer or flake.nix, LFS is typically available automatically. - If you clone without LFS installed, run
git lfs pull after installation to fetch the actual data files instead of the small pointer files.
"},{"location":"DEVELOPMENT/#devcontainer-setup-recommended","title":"Devcontainer Setup (Recommended)","text":"The devcontainer provides an identical experience to local development but with all system dependencies pre-configured. It automatically detects your hardware (GPU vs CPU) and syncs the appropriate dependencies.
"},{"location":"DEVELOPMENT/#prerequisites","title":"Prerequisites","text":" - Docker Desktop or Docker Engine.
- NVIDIA Container Toolkit (for GPU support).
"},{"location":"DEVELOPMENT/#setup-for-vs-code","title":"Setup for VS Code","text":" - Install the Dev Containers extension.
- Open the project and click Reopen in Container.
- On first launch, the
post-create.sh script will: - Detect if an NVIDIA GPU is available via
nvidia-smi. - Run
uv sync --frozen --extra cuda if a GPU is found. - Run
uv sync --frozen otherwise. - The environment is stored in a named volume for
.venv to ensure persistence and performance.
"},{"location":"DEVELOPMENT/#setup-for-jetbrains-ides","title":"Setup for JetBrains IDEs","text":" - The IDE will detect the
.devcontainer/devcontainer.json file. - The environment is pre-configured to point to
/workspaces/project/.venv. - The hardware-aware sync will run automatically during container creation.
"},{"location":"DEVELOPMENT/#local-development-alternative","title":"Local Development (Alternative)","text":"If you prefer not to use Docker:
- Install uv.
- Run
uv sync --frozen (CPU) or uv sync --frozen --extra cuda (GPU).
"},{"location":"DEVELOPMENT/#hardware-acceleration-jax","title":"Hardware Acceleration (JAX)","text":"Verify your setup by running the JAX initialization test:
uv run pytest tests/test_jax_init.py\n
In the devcontainer, this will succeed on both CPU and GPU. A GpuDevice is expected if a GPU is detected and the cuda extra was installed.
"},{"location":"DEVELOPMENT/#logging-monitoring","title":"Logging & Monitoring","text":"This project uses a unified logging system through the experiment_logger package.
- Usage in Code: To use the logger in your scripts, refer to the package README for the API reference.
- WandB/TensorBoard Setup: For information on how to configure tracking for experiments, see the Tracking & Monitoring API Guide.
The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job), adjusting progress bars and fallback modes accordingly.
"},{"location":"HPC/","title":"HPC Guide","text":"Full documentation: https://docs.hpc.ugent.be/
"},{"location":"HPC/#storage-overview","title":"Storage Overview","text":" - Run Outputs: Written to
$VSC_SCRATCH during the job (fast I/O) and copied to $VSC_DATA at the end for persistence. - Virtual Environments: Managed on
$VSC_DATA by mirroring configuration files. This avoids the 3GB home quota without requiring symlinks in the project root.
"},{"location":"HPC/#initial-environment-setup","title":"Initial Environment Setup","text":"Run once after cloning the repository. This script handles all modules, mirroring, and environment synchronization.
# Option A: Interactive (on a compute node)\nmodule swap cluster/donphan # Debug cluster (CPU only)\n# OR for GPU clusters:\n# module swap cluster/joltik\n# module swap cluster/accelgor\n# module swap cluster/litleo\n\nqsub -I -l nodes=1:gpus=1 # Only for GPU clusters\ncd \"${PBS_O_WORKDIR}\"\nbash scripts/hpc/install.sh\n\n# Option B: Batch (Run in background)\n# NOTE: GPU clusters (joltik/accelgor/litleo) require -l gpus=1 at runtime\nqsub -l gpus=1 scripts/hpc/install.sh\n
"},{"location":"HPC/#production-vs-debug-clusters","title":"Production vs. Debug Clusters","text":"Our scripts are cluster-agnostic and do not have hardcoded GPU requirements. Instead, you must request GPUs at runtime using the -l gpus=1 flag when submitting to a production GPU cluster.
"},{"location":"HPC/#debugging-donphan","title":"Debugging (Donphan)","text":"The donphan cluster does not support GPUs. Simply run the scripts without extra resource flags:
module swap cluster/donphan\nqsub scripts/hpc/train.pbs\n
"},{"location":"HPC/#production-joltik-accelgor-litleo","title":"Production (Joltik, Accelgor, Litleo)","text":"These clusters provide GPU acceleration and require a GPU request at runtime:
module swap cluster/joltik # or accelgor/litleo\nqsub -l gpus=1 scripts/hpc/train.pbs\n
"},{"location":"HPC/#interactive-debugging","title":"Interactive Debugging","text":"To activate your environment for interactive work, simply run the same install.sh script.
qsub -I -l nodes=1:ppn=4 -l walltime=1:00:00\ncd \"$PBS_O_WORKDIR\"\nbash scripts/hpc/install.sh\n
"},{"location":"HPC/#verification-commands","title":"Verification Commands","text":"After installation, run these commands to ensure your environment is set up correctly:
-
Verify Quota Safety:
ls -d venvs 2>/dev/null && echo \"FAIL\" || echo \">>> PASS: Project root is clean.\"\n
-
Verify Library Versions (NumPy Fix):
python -c \"import numpy; print(f'NumPy: {numpy.__version__}')\"\n# Expected: 2.x.x (Venv version), not 1.2x (System version)\n
-
Verify GPU Access:
python -c \"import torch, jax; print(f'GPU: {torch.cuda.is_available()}'); print(f'JAX: {jax.devices()}')\"\n
"},{"location":"HPC/#managing-dependencies","title":"Managing Dependencies","text":"env/hpc/requirements.txt is auto-generated from pyproject.toml. To regenerate:
uv run scripts/hpc/export_requirements.py\n
Modules listed in env/hpc/modules.txt are automatically excluded from the pip requirements to save space and use HPC-optimized binaries.
"},{"location":"api/analysis/","title":"Analysis & Plotting Tools","text":"This guide outlines the tools available for analyzing experimental data and generating poster-quality visualizations for the Brittle Star project.
"},{"location":"api/analysis/#shared-configuration","title":"Shared Configuration","text":"All plotting scripts share a central configuration in scripts/plots/plot_config.py. This file defines:
- Color Palette: A color-blind friendly, high-contrast palette for different architectures.
- Typography: Consistent font sizes and styles tailored for A0 posters.
- Markers: Shared visual indicators, such as the \u2605 used for best performers.
"},{"location":"api/analysis/#comparison-visualization","title":"Comparison Visualization","text":"The scripts/plots/analyze_comparisons.py script generates grouped bar charts comparing the performance of different architectures across various morphologies.
"},{"location":"api/analysis/#usage","title":"Usage","text":"Run the script from the root of the project, providing the path to your evaluation CSV:
# Basic usage (saves PNG and SVG to runs/evaluation/plots/)\nuv run python scripts/plots/analyze_comparisons.py path/to/results.csv\n\n# Advanced usage for Figma/Poster integration\nuv run python scripts/plots/analyze_comparisons.py path/to/results.csv \\\n --output_dir docs/assets/plots/ \\\n --font_size 30 \\\n --fig_width 14 \\\n --fig_height 10\n
"},{"location":"api/analysis/#cli-arguments","title":"CLI Arguments","text":" input_csv: (Required) Path to the CSV file containing evaluation results.--output_dir, -o: Directory where plots will be saved (default: runs/evaluation/plots).--show_titles: Include titles in the plots. Default is False, as titles are typically added natively in design tools like Figma.--font_size: Base font size in points (default: 28).--fig_width / --fig_height: Physical dimensions of the plot in inches. Match these to your Figma layout to maintain exact font sizes.
"},{"location":"api/analysis/#outputs","title":"Outputs","text":"The script generates four key plots, each saved as both .png and .svg:
- Forward Velocity: Grouped bar chart (cm/s).
- Accumulated Reward: Mean cumulative reward.
- Success Rate: Target acquisition percentage.
- Distance Remaining: Navigational accuracy.
"},{"location":"api/analysis/#convergence-analysis","title":"Convergence Analysis","text":"The scripts/plots/analyze_convergence.py script determines the convergence point of training runs.
"},{"location":"api/analysis/#usage_1","title":"Usage","text":"uv run python scripts/plots/analyze_convergence.py --output_dir runs/convergence/\n
"},{"location":"api/analysis/#configuration","title":"Configuration","text":" - File Mapping: The script uses hardcoded paths in the
FILE_MAPPING dictionary. Update these paths to point to your specific run evaluation files. - CLI Arguments: Supports the same
--show_titles, --font_size, and --fig_width/height flags as the comparison script.
"},{"location":"api/analysis/#outputs_1","title":"Outputs","text":"Generates three plots (PNG & SVG):
convergence_comparison: Grouped horizontal bar chart.progress_reward_curves: Line plots of reward over time.progress_velocity_curves: Line plots of velocity over time.
"},{"location":"api/analysis/#poster-integration-figma","title":"Poster Integration (Figma)","text":""},{"location":"api/analysis/#svg-scaling","title":"SVG & Scaling","text":"We recommend using the SVG outputs for poster design in Figma:
- No Resolution Loss: SVGs are vector-based and will remain sharp at any size.
- Native Text: Text in the SVG imports as native text layers in Figma.
- Exact Font Matching: To ensure a
28pt font in the plot matches a 28pt font in your poster, set the --fig_width and --fig_height to match the physical dimensions of the plot box in your Figma layout. - Editable: You can \"Ungroup\" the SVG in Figma to manually move labels, adjust colors, or tweak individual bars.
"},{"location":"api/analysis/#image-placeholders","title":"Image Placeholders","text":"The comparison charts include light-gray square placeholders below the X-axis. These are designed as guides; in Figma, you can drop your morphology renders or illustrations directly on top of these squares.
"},{"location":"api/environment/","title":"Brittle star environment","text":""},{"location":"api/environment/#creation","title":"Creation","text":"The environment package contains a factory class BrittleStarEnvFactory that creates instances of the environment/morphologies/... It uses the configuration classes defined in env_config.py to create the instances.
"},{"location":"api/environment/#configuration","title":"Configuration","text":"The data classes in env_config have default values as stated in the tutorials.
- MorphologyConfig: configuration for the morphology of the brittle star. Contains number of arms, number of segments per arm, and control mode.
- ArenaConfig: configuration for the arena. Sets the size of the arena, whether to set the ground floor to sand, attach a target and sizes of the walls.
- EnvConfig: configuration for the environment. These set shared settings such as camera locations, simulation time and the task.
"},{"location":"api/environment/#backend-and-task-enums","title":"Backend and Task enums","text":"The Backend enum specifies either an MJC or MJX backend.
- MJC: runs on CPU
- MJX: uses jax on the gpu
The Task enum specifies which task to use. 2 items are present:
- DIRECTED_LOCOMOTION: move to a target location
- LIGHT_ESCAPE: situation where the robot must move to a darker location
"},{"location":"api/evaluation/","title":"Checkpoint & Model Evaluation","text":"This guide covers how to evaluate trained brittle star models, with a focus on measuring defect tolerance (amputations) across different controller architectures.
"},{"location":"api/evaluation/#checkpoint-evaluation-during-training","title":"Checkpoint Evaluation (During Training)","text":"The PPOTrainer can automatically evaluate every saved checkpoint using the fast MJX backend. This is enabled via configuration.
"},{"location":"api/evaluation/#configuration","title":"Configuration","text":"In your experiment config or via CLI:
python scripts/train.py evaluation.evaluate_checkpoints=true evaluation.eval_max_steps=5000\n
Results are saved to runs/<run_dir>/metrics/checkpoint_evaluation.csv and synced to Weights & Biases if enabled.
"},{"location":"api/evaluation/#cross-model-fault-tolerance-analysis","title":"Cross-Model & Fault Tolerance Analysis","text":"To measure how well different controllers handle damage (amputations), use scripts/compare_models.py. This script performs a grid search over models x morphologies.
- Create or update a YAML file in
configs/evaluation. - Run the benchmark:
python scripts/compare_models.py evaluation=poster\n
The script will evaluate every combination of model and morphology for the specified number of episodes.
The results are saved to a CSV (default: metrics/model_comparison.csv).
"},{"location":"api/evaluation/#csv-schema","title":"CSV Schema","text":"Column Description model_path Path to the trained weights. architecture The morph_mode of the model (e.g., CENTRALIZED, RING). arm_0 ... arm_4 Number of segments in each arm slot (0 = amputated). num_active_arms Total number of arms with segments > 0. seed The episode seed. eval_return Accumulated shaped reward. approx_max_velocity Average velocity: (initial_dist - final_dist) / steps. reached_target Whether the robot finished within the success radius."},{"location":"api/evaluation/#post-hoc-checkpoint-scanning","title":"Post-hoc Checkpoint Scanning","text":"If you need to re-evaluate every saved checkpoint in a run (e.g., to generate a learning curve with different metrics):
python scripts/evaluate_checkpoints.py \\\n simulation.model_path=runs/<run_id>/final_model.flax \\\n evaluation.eval_max_steps=2000\n
This script scans the checkpoints/ directory of the specified run and evaluates every .flax file it finds using the model's training morphology.
For a step-by-step walkthrough on using these evaluation phases to reproduce our project results, see the Results & Reproduction Guide.
"},{"location":"api/reproduction/","title":"Results & Reproduction","text":"This guide explains how to access our official training logs and reproduce our results.
Our official training runs, model configurations, and metrics are publicly hosted on Weights & Biases (WandB).
"},{"location":"api/reproduction/#weights-biases-wandb-project","title":"Weights & Biases (WandB) Project","text":"All experiments, final models, and training logs are tracked in our public WandB project:
- Official Runs Table: WandB final-models-v2 Table
This page lists the verified runs with their architecture types, morphology definitions, evaluation metrics, and final model performance.
"},{"location":"api/reproduction/#how-to-reproduce-a-run-from-wandb","title":"How to Reproduce a Run from WandB","text":"Weights & Biases provides a built-in feature to extract the exact parameters and commands used for any given run:
- Open the WandB final-models-v2 Table.
- Click on the name of the run you wish to reproduce to open its detail page.
- In the top-right corner of the run header (next to the run name, not the main workspace header), click the three dots (
...) menu. - Select \"Reproduce run\". This will display the exact command-line arguments and configuration settings used to execute that run.
"},{"location":"api/reproduction/#local-hpc-reproduction-workflow","title":"Local & HPC Reproduction Workflow","text":"To reproduce our training and evaluation phases locally or on an HPC cluster, follow the procedures below.
"},{"location":"api/reproduction/#1-environment-setup","title":"1. Environment Setup","text":"To ensure identical package versions (including JAX, Flax, and MuJoCo), sync your environment using the lockfile:
uv sync --frozen\n
"},{"location":"api/reproduction/#2-training-phase","title":"2. Training Phase","text":"Run the training script using the exact parameters retrieved from WandB's \"Reproduce run\" page or from a downloaded _metadata.yaml file:
uv run python scripts/train.py experiment=my_experiment ppo.learning_rate=0.001 experiment.seed=42\n
"},{"location":"api/reproduction/#evaluation-phases","title":"Evaluation Phases","text":"Reproducing our evaluation results is divided into two distinct phases:
"},{"location":"api/reproduction/#phase-1-determining-the-best-checkpoint","title":"Phase 1: Determining the Best Checkpoint","text":"During training, checkpoints are saved at regular intervals. To determine which of these checkpoints performed the best:
- Evaluate Checkpoints Post-Training: If checkpoint evaluation was not run during training, scan the completed run's checkpoints folder by pointing to the final model path:
uv run python scripts/evaluate_checkpoints.py simulation.model_path=runs/your_run_dir/final_model.flax\n
This script runs deterministic rollouts for every checkpoint in runs/your_run_dir/checkpoints/.
- Locate the Results: The evaluations are saved to:
runs/your_run_dir/metrics/checkpoint_evaluation.csv\n
Analyze this CSV to find the checkpoint iteration with the highest average return or target success rate. This checkpoint will be used for cross-architecture comparisons.
"},{"location":"api/reproduction/#phase-2-comparing-checkpoints-between-architectures","title":"Phase 2: Comparing Checkpoints Between Architectures","text":"Once the best checkpoints for each architecture are identified, they are compared under shared, standardized environments (including fault tolerance checks such as leg amputations).
- Configure the Comparison Models: Open or create an evaluation config file (e.g.,
configs/evaluation/poster.yaml) and add the paths to the best checkpoints:
# configs/evaluation/poster.yaml\nevaluation:\n comparison_models:\n - runs/run_arch_centralized/checkpoints/checkpoint_best.flax\n - runs/run_arch_decentralized/checkpoints/checkpoint_best.flax\n
- Execute the Comparison Script: Run the comparison script using your config:
uv run python scripts/compare_models.py evaluation=poster\n
This script runs multiple sequential evaluation episodes (defined by comparison_num_episodes starting at comparison_base_seed) for every model across the selected morphologies.
-
Analyze Comparison Metrics: The script writes a consolidated CSV file to metrics/model_comparison.csv containing:
-
eval_return: The cumulative return.
approx_max_velocity: The distance covered per step.reached_target: Navigational success rates.arm_0 to arm_4: Active segments per arm (indicating damage/amputations).
This CSV can then be passed to the plotting scripts (e.g., scripts/plots/analyze_comparisons.py) to generate visualization plots. For details on configuration and outputs, see the Analysis & Plotting Guide.
"},{"location":"api/simulation/","title":"Interactive Simulation & Visualization","text":"The simulation pipeline allows you to visualize trained models and observe their behavior under various conditions.
"},{"location":"api/simulation/#overview","title":"Overview","text":"The simulation pipeline is metadata-driven. Training-specific configurations (morphology, arena, environment, etc.) are automatically loaded from the _metadata.yaml file associated with the model checkpoint.
"},{"location":"api/simulation/#basic-simulation","title":"Basic Simulation","text":"To simulate a model in the MuJoCo viewer:
uv run scripts/simulate.py simulation.model_path=runs/your_run/final_model.flax\n
"},{"location":"api/simulation/#amputation-morphology-overrides","title":"Amputation & Morphology Overrides","text":"You can test trained models on different morphologies (e.g., amputating legs) by providing a morphology override. The observations will be automatically padded up to the training morphology's dimensions:
uv run scripts/simulate.py \\\n simulation.model_path=runs/your_run/final_model.flax \\\n simulation.morphology_override=configs/morphology/3_arms.yaml\n
"},{"location":"api/simulation/#video-recording","title":"Video Recording","text":"Recording videos requires the [evaluation] extra:
uv run scripts/simulate.py \\\n simulation.model_path=runs/your_run/final_model.flax \\\n simulation.record_video=true \\\n simulation.max_steps=1000\n
Videos and evaluation metadata are stored in timestamped folders alongside the model: runs/your_run/final_model_evaluations/eval_<timestamp>/simulation.mp4
"},{"location":"api/simulation/#top-down-and-follow-cameras","title":"Top-Down and Follow Cameras","text":"Using the following script, you can render a top-down and follow camera view for multiple models at once:
uv run scripts/poster_visualisations/render_poster_videos.py \\\n runs/final-models/centralized/.../final_model.flax \\\n runs/final-models/fully-connected/.../final_model.flax \\\n runs/final-models/ring/.../final_model.flax \\\n --max-steps 10000 --width 640 --height 480 --fps 60 \\\n --output-root vids/poster/\n
For batch evaluation, checkpoint analysis, and cross-model architecture comparisons, see the Checkpoint & Model Evaluation Guide."},{"location":"api/tracking/","title":"Tracking & Monitoring","text":"This guide explains how to monitor your experiments using Weights & Biases (WandB) and TensorBoard.
"},{"location":"api/tracking/#weights-biases-wandb","title":"Weights & Biases (WandB)","text":"WandB is used for online synchronization and visualization of training metrics.
"},{"location":"api/tracking/#authorization","title":"Authorization","text":"Export your API key in your terminal to enable WandB synchronization:
export WANDB_API_KEY=your_copied_api_key_here\n
Alternatively, you can log in using the CLI:
uv run wandb login\n
"},{"location":"api/tracking/#enabling-tracking","title":"Enabling Tracking","text":"To enable online sync during a training run, set logging.track=true on the command line:
uv run python scripts/train.py logging.track=true\n
You can also configure your project and entity:
uv run python scripts/train.py \\\n logging.track=true \\\n logging.wandb_project_name=\"MyProject\" \\\n logging.wandb_entity=\"my-team\"\n
These can also be set in your configuration YAML file under the logging key.
"},{"location":"api/tracking/#local-monitoring-with-tensorboard","title":"Local Monitoring with TensorBoard","text":"All runs are recorded locally in the runs/ directory (or the directory specified in experiment.base_run_dir). You can view scalars and other metrics with TensorBoard:
tensorboard --logdir runs/\n
Access the interface at http://localhost:6006.
"},{"location":"api/tracking/#cli-exploration-tool","title":"CLI Exploration Tool","text":"For quick diagnostics or to export data to CSV without launching the full TensorBoard UI, you can use the explore_tensorboard.py script:
uv run python scripts/analysis/explore_tensorboard.py runs/your_run_name/\n
See the detailed description in /scripts/analysis/README.md.
"},{"location":"api/tracking/#developer-logging-api","title":"Developer Logging API","text":"For details on the developer API of our internal logging library (how backend routing, checkpoint synchronization, and singleton initialization works), see the Experiment Logger API Guide.
"},{"location":"api/training/","title":"Training Models","text":"This guide covers how to configure and run training experiments for the Brittle Star project using Hydra-based configurations.
"},{"location":"api/training/#configuration","title":"Configuration","text":"The project uses a modular configuration system powered by Hydra. Instead of passing many command-line flags, you select and override configuration groups.
For a detailed guide on the structure, validation, and usage of our Hydra configuration files, see the Brittle Star Configuration System Guide.
"},{"location":"api/training/#creating-a-custom-experiment","title":"Creating a Custom Experiment","text":" -
Create a new experiment file: Create a file at configs/experiment/my_experiment.yaml. You can copy an existing one as a template:
cp configs/experiment/base.yaml configs/experiment/my_experiment.yaml\n
-
Edit configs/experiment/my_experiment.yaml to set your experiment parameters:
# @package _global_\nexperiment:\n exp_name: \"my_custom_run\"\n seed: 42\n
"},{"location":"api/training/#training-execution","title":"Training Execution","text":"To start a training run with the default settings defined in configs/main_config.yaml:
uv run python scripts/train.py\n
"},{"location":"api/training/#using-a-custom-experiment-configuration","title":"Using a Custom Experiment Configuration","text":"To run with your custom experiment file:
uv run python scripts/train.py experiment=my_experiment\n
uv run python scripts/train.py ppo.learning_rate=0.001 ppo.num_envs=32 logging.track=true\n
"},{"location":"api/training/#evaluation-during-training","title":"Evaluation During Training","text":"By default, the trainer saves checkpoints but does not evaluate them. To enable automatic headless evaluation of every saved checkpoint, set evaluation.evaluate_checkpoints=true:
uv run python scripts/train.py evaluation.evaluate_checkpoints=true\n
"},{"location":"api/training/#reproducing-experiments","title":"Reproducing Experiments","text":"For detailed steps on how to reproduce training runs, locate run configuration metadata, or reproduce our experiments using Weights & Biases (WandB), see the Results & Reproduction Guide.
For more details on evaluation metrics and comparison tools, see Checkpoint & Model Evaluation.
For more details on tracking your experiments, see Tracking & Monitoring.
"},{"location":"configs/","title":"Brittle Star Configuration System","text":"This project uses Hydra for a modular, hierarchical, and strictly-typed configuration system.
"},{"location":"configs/#core-concepts","title":"Core Concepts","text":" - Composition over Inheritance: Instead of one giant config file, the configuration is composed of small, domain-specific modules (PPO settings, architecture, morphology, etc.).
- Strict Typing: Every configuration is validated against a Python dataclass schema (
ConfigStore). Misspelled keys throw a ConfigAttributeError immediately. - CLI Swapping: You can swap entire modules or override individual values from the command line without touching code.
"},{"location":"configs/#directory-structure","title":"Directory Structure","text":" main_config.yaml: The root entry point defining the default composition.experiment/: High-level experiment settings (seed, device).logging/: WandB and checkpointing configuration.ppo/: PPO training hyperparameters.architecture/: Polymorphic network architectures (centralized vs. decentralized).morphology/: Physical robot definitions (number of segments, amputations).arena/: Environment physics and visual settings.environment/: Task-specific settings (Directed Locomotion, Light Escape).
"},{"location":"configs/#common-commands","title":"Common Commands","text":""},{"location":"configs/#local-debugging","title":"Local Debugging","text":"Run a quick test with minimal iterations:
python scripts/train.py experiment=dev_test ppo=fast\n
"},{"location":"configs/#swapping-architectures-or-morphologies","title":"Swapping Architectures or Morphologies","text":"Test a decentralized controller on a 3-arm robot:
python scripts/train.py architecture=decentralized morphology=3_arms\n
"},{"location":"configs/#hpc-production","title":"HPC Production","text":"Run stable PPO with WandB enabled (HPC submission scripts handle the hydra.run.dir redirection):
python scripts/train.py ppo=stable logging=wandb_enabled\n
"},{"location":"configs/#dry-run-validation","title":"Dry-Run Validation","text":"Check if your configuration is valid without starting the simulation:
python scripts/train.py --cfg job\n
"},{"location":"configs/#developer-notes","title":"Developer Notes","text":" - Adding a new group: Create a subdirectory in
configs/ and register the new dataclass in src/brittle_star_project/configs/register_configs.py. - Typo Catching: If you see a
ConfigAttributeError, check for typos in your YAML keys or CLI overrides. - Output Redirection: We use
experiment.base_run_dir to configure where logs and models are stored (defaults to runs/). - To change it locally:
python scripts/train.py experiment.base_run_dir=/path/to/custom/dir - On HPC, ensure this points to a fast scratch storage.
"},{"location":"design/actor-critic/","title":"Actor-Critic Architecture","text":"To process observations into actions, our controllers utilize an Actor-Critic architecture. Because we use Proximal Policy Optimization (PPO), the pipeline fundamentally requires separate networks for the policy (Actor) and the value estimation (Critic).
Centralized Architecture (Baseline)
This pipeline treats the agent as a single entity and uses standard Proximal Policy Optimization (PPO).
- Centralized Actor: Composed of two chained MLPs (Sensor \\(\\rightarrow\\) Motor) passing a hidden state between them. The centralized sensor receives the concatenated global state vector of all limbs at once and processes it into a hidden state. The centralized motor receives this hidden state and outputs the joint offsets for all actuators simultaneously. This is mathematically equivalent to using one large MLP with hidden layers, but splitting makes the implementation easier by allowing us to reuse the same components for the decentralized modules.
- Centralized Critic: Composed of two sequential MLPs (Feature Extractor \\(\\rightarrow\\) Critic). Because PPO evaluates the state-value function, this network only receives the concatenated global state vector (no actions). It outputs a single scalar estimating the expected future reward for the entire agent.
Our policy and value networks use separate input networks/feature extractors as advised by the SEL3 course assistants and the blog. For continuous actions this should allow better learning at a small cost.
graph TD\n Obs([Global Observation])\n\n Sens[Sensor]\n Act[Motor]\n OutAct([Action Distribution<br/>mean, log_std])\n\n Feat[Feature extractor]\n Crit[Critic]\n OutCrit([Value Estimate<br/>scalar])\n\n Obs --> Sens\n Obs --> Feat\n\n Sens -->|\"Hidden state\"| Act\n Feat -->|\"Hidden state\"| Crit\n\n Act --> OutAct\n Crit --> OutCrit
Decentralized Architecture
This pipeline utilizes the \"Centralized Training with Decentralized Execution\" principle, specifically the NerveNet-MLP variant.
- Decentralized Actor, split into three distinct models:
- Sensor: A local model at each node. It receives its local state plus the goal vector directly, processing them into an initial hidden state.
- Propagator: Nodes synchronously compute and exchange messages with connected neighbors for \\(N\\) steps to update their hidden states. See communication.md for details.
- Motor: A local model uses its final updated hidden state to output the joint offset strictly for its own actuator.
- Centralized Critic: Composed of two sequential MLPs (Feature Extractor \\(\\rightarrow\\) Critic). During training, it acts globally by taking the concatenated state vectors from all sensors to output a single, global state-value scalar evaluating the entire agent's pose.
To keep the implementation simple, we should use one critic per node in our architecture, but only a single, global critic for all nodes at once, for the following reasons:
- Credit Assignment Problem (Ha, 2017): The MuJoCo simulator provides an overall reward based on the brittle star movement progression, e.g. total distance travelled. Using an isolated critic for each node in the network would not allow to determine which local action contributed to the global success. A global critic solves this by evaluating the combined state of the agent at once.
- Implementation simplicity: Building a second decentralized message-passing graph for the critic (NerveNet-2) would require more coding. Using a standard MLP that concatenates all raw input vectors is much easier to program while mathematically equivalent.
graph TD\n Obs([Local Observation])\n\n Sens[Sensor]\n Prop[Propagator]\n Feat[Feature extractor]\n\n Mot[Motor]\n Crit[Critic]\n\n OutMot([Action Distribution<br/>mean, log_std])\n OutCrit([Value Estimate<br/>scalar])\n\n Obs --> Sens\n Sens -->|\"Hidden state\"| Prop\n Obs --> Feat\n\n Prop -->|\"Hidden state\"| Mot\n\n\n Feat -->|\"Hidden state\"| Crit\n\n Mot --> OutMot\n Crit --> OutCrit\n\n Prop -.->|\"message passing\"|Prop
"},{"location":"design/actor-critic/#implementation-details-network-depth","title":"Implementation Details (Network Depth)","text":"Inspired by: PPO Implementation Details
The MLPs used in both pipelines are defined with specific hidden layer configurations to balance learning capability and computational cost. As of right now, though this might change as we make progress in our experiments, we use:
- Input Networks (Sensors & Feature Extractors): These networks map the raw state inputs to internal hidden states. They are configured as standard dense networks with 2 hidden layers of 64 nodes each (
[64, 64]) and utilize tanh activation functions. - Output Networks (Motors, Actors & Critics): The final output models are intentionally kept shallow. The Actor directly projects the hidden state to a continuous action distribution (
mean and log_std) using a single dense output layer (zero hidden layers) initialized orthogonally. The Critic functions similarly, mapping the hidden representation to a single scalar value.
Note: For the continuous action distributions outputted by the Motor, we explicitly use mean and log_std as advised by previous research to maintain learning stability.
References
- Ha, D. (2017, October 29). A Visual Guide to Evolution Strategies. \u5927\u30c8\u30ed \u30fb Machine Learning. https://blog.otoro.net/2017/10/29/visual-evolution-strategies/
- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. \u2018Proximal Policy Optimization Algorithms\u2019. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. https://doi.org/10.48550/arXiv.1707.06347.
- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. \u2018NerveNet: Learning Structured Policy with Graph Neural Networks\u2019. Conference paper presented at International Conference on Learning Representations. 15 February 2018. https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613.
"},{"location":"design/communication/","title":"Communication scheme (Message Passing)","text":"Remember our research question:
\"What is the impact of different levels of controller modularity on learning speed, coordination, and fault tolerance (e.g. amputations) in brittle-star-like robots trained with Reinforcement Learning?\"
To test decentralized modularity (such as arm-level or segment-level controllers), the various modules must be able to communicate with each other to achieve coordinated locomotion. This is accomplished through message passing in a Graph Neural Network (GNN)-like architecture. Two prominent communication styles from the literature are N-step NerveNet (Wang et al., 2018) and bottom-up top-down Shared Modular Policies (Huang et al., 2020).
We have chosen to apply one uniform communication style across all modular architectures, specifically opting for N-step NerveNet.
"},{"location":"design/communication/#rationale","title":"Rationale","text":"Initially, our idea was to equip arm-level controllers with NerveNet message passing and segment-level controllers with SMP. However, we evaluated that this introduces a threat to the validity of our research question. If we observe differences in performance, it would be impossible to determine whether the variance is caused by the level of modularity, or by the difference in the message passing scheme. To purely compare modularity, the communication scheme style must remain constant.
Second, we decided that NerveNet is a better fit for our research. The morphology of our brittle star contains cycles at the decentralized level (e.g., a ring of segments or arms around the body). NerveNet has proven to be robust for arbitrary structures, including graphs with cycles. SMP inherently expects a tree structure for its bottom-up and top-down pass. Applying SMP to a ring structure requires a workaround to break that cycle.
"},{"location":"design/communication/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Choosing NerveNet introduces a scalability issue as the morphology grows. In NerveNet, a message advances only one segment or node per propagation step. When dealing with long arms (e.g., > 5 segments), this requires a large number of propagation steps to transmit information from one tip of an arm to another.
If we were to use SMP instead - which is possible - the inner states of nodes are shared across the entire graph in just two passes. For very large or long morphologies, this would be much more scalable.
By rejecting SMP, we accept that our model might learn slower or require more computational power for highly segmented, extended morphologies.
References
- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. \u2018NerveNet: Learning Structured Policy with Graph Neural Networks\u2019. Conference paper presented at International Conference on Learning Representations. 15 February 2018. https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613.
- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. \u2018One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control\u2019. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. https://doi.org/10.48550/arXiv.2007.04976.
"},{"location":"design/controllers/","title":"Levels of modularity and topology","text":"The brittle star can be controlled at different levels. A monolithic controller processes all inputs and outputs at once, whereas modular controllers divide the brains across the body, inspired by the biology of brittle stars.
We define four architectures to compare:
- Centralized, monolithic: A single Multi Layer Perceptron per robot that receives all observations and outputs all actions.
- Fully connected arm-level: Each arm contains an MLP that processes the inputs for that arm, an MLP that processes the communicated inner-states, and an MLP that outputs the actions for that arm. One policy for these MLPs is shared across the arms. The controllers in each arm are connected to each other and form a fully connected graph. There is no central disk, but the controllers are fully connected.
- Ring arm-level: Identical setup to the fully connected arm-level, but the controllers are connected in a ring structure. This setup is considered less centralized than the fully connected graph.
- Segment-level: Each segment contains the three MLPs discussed above. The base segments, attached to the body, form a ring structure, with the remaining segments attached as extended \"strings\". Segments can only communicate with segments that are physically connected to it.
"},{"location":"design/controllers/#rationale","title":"Rationale","text":"To fairly compare decentralized modularity against centralized control, the decentralized models should not be allowed to contain a central organ acting as a bottleneck or coordinator. By removing the central disk in the decentralized models and replacing it with a ring topology, we closely approximate the biological reality of the brittle star and test a decentralized morphology.
The fully connected graph functions as an intermediate step in between a fully centralized and a decentralized ring. We use it to test whether our models scale to more complex structures.
"},{"location":"design/input_action_spaces/","title":"Input (state) and output (action) spaces","text":"To effectively learn locomotion and navigation, the agent requires a well-defined observation space (inputs) and action space (outputs). The control models map these observations directly to physical movements.
Inputs (state space)
The observation space provides the agent with its current physical state and its navigational objective. With a decentralized control architecture in mind, we divide these inputs into global and local states.
Global inputs, always broadcasted to all nodes:
- Vertical orientation/tilt: A single, simplified metric representing the tilt/vertical alignment of the agent's central body/disk, a.k.a. the deviation from the global Z-axis. Its value is derived from the environment's raw disk rotation 3D vector \\([roll, pitch, yaw]\\): $$ tilt = sqrt(roll^2 + pitch^2) $$
- Goal vector: A 2D unit vector representing the egocentric direction to the target. A value of \\([1.0, 0.0]\\) indicates that the target is directly in front of the agent (angle 0).
Local inputs, routed directly to specific nodes:
- Joint positions: The current angles of all joints within the morphology.
- Joint velocities: The current angular velocities of the joints.
- Joint actuator forces: The physical forces currently exerted at each specific joint.
- Segment contact: These values indicate whether each physical segment of the agent is currently touching the ground.
Outputs (action space)
The action space defines how the agent interacts with the environment.
- Joint offsets: absolute target positions (offsets) for the joints, i.e. the exact angle the joint should move to.
"},{"location":"design/input_action_spaces/#normalization-and-scaling","title":"Normalization and Scaling","text":"Both the input (observation) and output (action) spaces are rescaled to the range \\([-1, 1]\\).
For the input space, all raw physical values (angles, velocities, forces, distances) are normalized based on their defined physical bounds. If a value exceeds these bounds during simulation, it is clipped to the \\([-1, 1]\\) range.
For the output space, the neural network's tanh-activated outputs (which naturally fall in \\([-1, 1]\\)) are linearly mapped to the physical joint limits defined in the robot's morphology.
"},{"location":"design/input_action_spaces/#rationale","title":"Rationale","text":"When designing the state space, we must ask: Could a human operator perform this task given only these inputs?
- Inclusion of Joint Velocities: Because our control models do not inherently possess memory of previous timesteps, providing only the joint position is insufficient to determine the direction a limb is currently moving. By explicitly including joint velocities, the agent can immediately infer momentum and movement direction without needing to memorize past states.
- Absolute Joint Offsets: The physical Brittle Star robot relies on servo motors (if we were to build this simulated robot), which are inherently position-controlled devices. (Continuous rotation servos exist, but they are less commonly used for joints.) If our network outputted continuous torques (forces), a significant portion of the reinforcement learning process would be wasted on learning low-level PID control dynamics (i.e., how much force to apply to hold a position). Abstracting this away forces the learning algorithm to focus entirely on higher-level gait generation and locomotion.
- Simplified vertical orientation: We drop the full 3D spatial rotation and angular velocity arrays in favor of a single vertical orientation metric (tilt). For a brittle star moving accross a flat plane, this metric is sufficient for the agent to sense if it is losing balance or flipping over.
- Force representation: We strictly retain the joint actuator forces and drop the more generic actuator force. Forces that are explicitly tied to individual joints are significantly easier to route into decentralized, local limb nodes, which is necessary for our message-passing architecture.
- Goal Vector (Distance + Angle): Providing only the scalar \"distance to the goal\" as an input is akin to blindfolding the robot and asking it to find a target by playing \"hot or cold.\" By providing a full vector, the agent knows exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion.
The goal representation is explicitly divided into a directional vector and a scalar distance. Keeping the direction as a normalized unit vector bounds the values to the \\([-1, 1]\\) range, which stabilizes neural network training. Providing only a scalar \"distance to the goal\" would force the agent to learning localized searching behaviors (e.g. random walks or spiraling) to deduce the direction, drastically increasing the difficulty of the learning task.
NOTE: We later dropped the \"distance to vector\", switching to only a direction as the input. Our reasoning is the agent should always move towards the goal (it should not learn to stop at the goal), which allows for this simplification that decreases the model input size.
The environment provides a raw unit_xy_direction_to_target (global), which we transform into a calculated robot_direction_to_target (egocentric) before passing it to the MLPs. This vector consists of the X and Y direction, where a value of \\([1.0, 0.0]\\) (mapping to an angle of \\(0\\)) means the robot is facing directly towards the target. - Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait transitions. - Zero-Centered Rescaling (\\([-1, 1]\\)): Using a zero-centered range is standard best practice for continuous control tasks. It provides several mathematical and physical advantages: - Improved Gradient Flow: Neural networks optimize faster when inputs are zero-centered. If all inputs were positive (e.g., \\([0, 1]\\)), the gradients during backpropagation would be forced to the same sign, causing inefficient \"zig-zag\" weight updates. - Meaningful Neutral State: In robotics, \\(0.0\\) naturally represents a resting state (zero velocity, centered position, no force). In a \\([-1, 1]\\) system, this physical rest maps to a neutral \\(0.0\\) signal in the network. This also correctly communicaties a \"neutral/dead\" signal for amputated limbs that are padded with \\(0.0\\) values.
Specifically, we do not include some available inputs:
- Global position: Absolute spatial coordinates can cause the agent to overfit to a specific coordinate frame or map, rather than learning general, adaptable locomotion strategies.
"},{"location":"design/input_action_spaces/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Alternative state and action formulations include:
- Torque-based continuous control: In many continuous control tasks (like standard MuJoCo benchmarks), actions represent continuous torques applied to joints. While this provides more granular, low-level physical control, it heavily complicates training and does not align well with the physical reality of servo-driven hardware.
- Recurrent Neural Networks (RNNs) / Frame Stacking: Instead of explicitly passing velocities in the state space, the network could infer momentum by observing a history of past states. Using RNNs or frame stacking allows the agent to build an internal memory of movement. However, this significantly increases architectural complexity and training time compared to explicitly providing the velocity data.
- Scalar Goal Distance: Giving the agent only the scalar distance to the target would force it to learn a localized searching behavior (e.g., spiraling or random walks) to determine the correct direction. While biologically plausible for simpler organisms following chemical gradients, it drastically increases the difficulty of the learning task.
- \\([0, 1]\\) Rescaling: While some domains (like computer vision) use \\([0, 1]\\) scaling, it is generally avoided in robotics. Scaling to \\([0, 1]\\) would mean that a resting joint (velocity = 0) maps to an input of \\(0.5\\). This constant positive bias forces the network to waste capacity learning to ignore or subtract this baseline signal just to stand still. Furthermore, it breaks the \"dead signal\" interpretation of zero-padding used for amputations.
"},{"location":"design/input_action_spaces/#mujoco","title":"MuJoCo","text":"This is what the filtered input vectors look like in MuJoCo, with \\(J\\) joints and \\(S\\) segments:
joint_position: shape=(J,), dtype=float64joint_velocity: shape=(J,), dtype=float64joint_actuator_force: shape=(J,), dtype=float64segment_contact: shape=(S,), dtype=float64robot_direction_to_target: shape=(2,), dtype=float64, egocentricdisk_z_tilt: shape=(1,), dtype=float64, derived from disk_rotation
This brings the entire input space down to \\(3J + S + 4\\) float64's, compared to \\(4J + S + 15\\) float64's for the unfiltered inputs.
For reference, these are all the inputs that are available in the MuJoCo environment:
obs keys: ['joint_position', 'joint_velocity', 'joint_actuator_force', 'actuator_force', 'disk_position', 'disk_rotation', 'disk_linear_velocity', 'disk_angular_velocity', 'tendon_position', 'tendon_velocity', 'segment_contact', 'unit_xy_direction_to_target', 'xy_distance_to_target']\n\nraw observations dict:\n{'joint_position': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'joint_velocity': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'joint_actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),\n 'disk_position': array([0. , 0. , 0.11]),\n 'disk_rotation': (0.0, -0.0, 0.0),\n 'disk_linear_velocity': array([0., 0., 0.]),\n 'disk_angular_velocity': array([0., 0., 0.]),\n 'tendon_position': array([], dtype=float64),\n 'tendon_velocity': array([], dtype=float64),\n 'segment_contact': array([0., 0., 0., 0., 0., 0.]),\n 'unit_xy_direction_to_target': array([-0.95333378, -0.30191837]),\n 'xy_distance_to_target': array([3.])}\n\n(shapes)\njoint_position: shape=(12,), dtype=float64, size=12\njoint_velocity: shape=(12,), dtype=float64, size=12\njoint_actuator_force: shape=(12,), dtype=float64, size=12\nactuator_force: shape=(12,), dtype=float64, size=12\ndisk_position: shape=(3,), dtype=float64, size=3\ndisk_rotation: shape=(3,), dtype=float64, size=3\ndisk_linear_velocity: shape=(3,), dtype=float64, size=3\ndisk_angular_velocity: shape=(3,), dtype=float64, size=3\ntendon_position: shape=(0,), dtype=float64, size=0\ntendon_velocity: shape=(0,), dtype=float64, size=0\nsegment_contact: shape=(6,), dtype=float64, size=6\nxy_distance_to_target: shape=(1,), dtype=float64, size=1\n
"},{"location":"design/learning_algorithm/","title":"Reinforcement Learning Algorithm","text":"To control the continuous action space (the joints of the robot) based on sensor data, we require a reliable Reinforcement Learning (RL) algorithm or optimization strategy.
We have chosen Proximal Policy Optimization (PPO) (Schulman et al., 2017).
"},{"location":"design/learning_algorithm/#rationale","title":"Rationale","text":"PPO is an on-policy algorithm known for its stability and robustness (safe training without excessive variance). More importantly, it requires relatively little hyperparameter tuning compared to other algorithms. Since NerveNet was successfully trained using PPO (Wang et al., 2018), selecting PPO significantly reduces the risk of convergence failures.
"},{"location":"design/learning_algorithm/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Alternative learning algorithms include:
- Twin Delayed DDPG (Fujimoto et al., 2018): TD3 is a strong off-policy alternative used in the SMP paper (Huang et al., 2020). It is highly sample-efficient and reportedly excels at zero-shot adaptations. However, this approach would be more complex and error-prone than with PPO.
- Evolution strategies (ES): Evolution strategies are useful for optimizing Central Pattern Generators (CPGs), e.g. CMA-ES, OpenAI-ES. While this method is easier to distribute and parallelize, ES typically scales worse with exceptionally large observation spaces compared to gradient-based RL methods like PPO.
References
- Fujimoto, Scott, Herke Hoof, and David Meger. \u2018Addressing Function Approximation Error in Actor-Critic Methods\u2019. Proceedings of the 35th International Conference on Machine Learning, 3 July 2018, 1587-96. https://proceedings.mlr.press/v80/fujimoto18a.html.
- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. \u2018One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control\u2019. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. https://doi.org/10.48550/arXiv.2007.04976.
- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. \u2018Proximal Policy Optimization Algorithms\u2019. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. https://doi.org/10.48550/arXiv.1707.06347.
"},{"location":"design/reward_function/","title":"Reward function and observation space","text":"The robot needs to know whether its movements contribute to the ultimate goal of locomotion towards a target. Sensor inputs must be distributed fairly to guarantee an objective comparison between different architectures.
- The distance from the robot to the target and/or the light intensity are treated as global inputs.
- Positions and joints, which are normalized to floating-point values between 0 and 1, are considered local inputs.
- The reward function is centered around minimizing the distance to the goal or maximizing the movement towards the goal within a finite number of timesteps \\(T\\).
- To motivate efficient movement, the amount of timesteps taken to reach the goal will be used as penalty.
- An extra penalty based on movement relative to the current step and the previous is used to penalize a movement away from the target.
"},{"location":"design/reward_function/#from-reward-to-ppo","title":"From reward to PPO","text":"The resulting reward is passed to our PPO library. Our critic network (value function) predicts how good our eventual reward will be for the current state, this value is combined with the reward from the reward function to get advantages. These advantages are then used to calculate the losses to update both our critic and actor pipeline.
"},{"location":"design/reward_function/#rationale","title":"Rationale","text":"Using a light source (or a gradient) is biologically plausible for many simple organisms. By normalizing all signals between 0 and 1, PPO training is highly stabilized. The timesteps must be finite to reset the environment in a timely manner if the policy gets stuck in a local minimum.
"},{"location":"design/reward_function/#limitations-and-alternatives","title":"Limitations and alternatives","text":"Providing global information to all individual decentralized segments can be considered biologically cheating or practically infeasible once the robot would be physically built. Some sensory input cannot be put in each joint, for example.
The alternative is to provide the global input to the outermost segments of the arms, or a specific set of segments assigned with this functionality. The network would then have to learn to propagate this signal throughout the body via message passing. While biologically more accurate, this drastically complicates the learning process. We have written this down as potential future research.
"},{"location":"scripts/analysis/","title":"Experiment Analysis Tools","text":"This directory contains scripts for post-processing and analyzing experiment results, including TensorBoard logs and saved model weights.
"},{"location":"scripts/analysis/#scripts","title":"Scripts","text":""},{"location":"scripts/analysis/#1-explore_tensorboardpy","title":"1. explore_tensorboard.py","text":"A CLI tool to summarize TensorBoard tfevents files without a GUI.
Key Features: - Displays last values, min, max, and step counts for all scalar metrics. - Calculates total run duration and estimated completion percentage. - Exports granular scalar data to CSV for analysis in Excel/Pandas.
Usage:
# General usage\npython explore_tensorboard.py <run_directory>\n\n# Exporting data\npython explore_tensorboard.py <run_directory> --csv data.csv\n
Requirements: - pandas - tensorboard - tensorflow-cpu (or tensorflow)
"},{"location":"src/experiment_logger/","title":"Experiment Logger","text":"A standardized, unified interface for logging experiments across multiple backends (WandB, TensorBoard, and Local Disk).
This library is designed to be a standalone package that decouples the logging logic from the core training routines in the brittle_star_project.
"},{"location":"src/experiment_logger/#quick-start","title":"Quick Start","text":"The recommended way to use the logger is through the get_logger() singleton:
from experiment_logger import UnifiedLogger, get_logger\n\n# Initialize at the start of your script (e.g., in train.py)\nlogger = UnifiedLogger(\n run_name=\"my_experiment_run\",\n config={\"learning_rate\": 3e-4},\n project_name=\"MyProject\",\n base_dir=\"runs\",\n use_wandb=True\n)\n\n# In other files, retrieve the initialized singleton:\n# logger = get_logger()\n\n# Log metrics (Scalar values, numpy scalars, or JAX types)\nlogger.log({\"loss\": 0.5, \"accuracy\": 0.98}, step=100)\n\n# Standard logging (Mirrored to disk and stdout)\nlogger.info(\"Training started\")\nlogger.warning(\"Learning rate is very high\")\n\n# Save checkpoints (Automatically synced to WandB as artifacts)\nlogger.save_checkpoint(params, step=5000)\n
"},{"location":"src/experiment_logger/#logger-classes","title":"Logger Classes","text":""},{"location":"src/experiment_logger/#unifiedlogger","title":"UnifiedLogger","text":"The full suite for production training. It manages: - WandB: Syncs metrics and uploads model checkpoints as artifacts. - TensorBoard: Writes events for local visualization. - Local Disk: Stores metrics in metrics.yaml and textual logs in run.log.
"},{"location":"src/experiment_logger/#simplelogger","title":"SimpleLogger","text":"A zero-dependency fallback that uses standard Python print() statements. Use this for standalone testing or minimal environments where you don't need persistent monitoring.
from experiment_logger import SimpleLogger\nlogger = SimpleLogger(run_name=\"test_run\")\n
"},{"location":"src/experiment_logger/#api-features","title":"API Features","text":""},{"location":"src/experiment_logger/#loggerprogress_bariterable-kwargs","title":"logger.progress_bar(iterable, **kwargs)","text":"A smart wrapper around tqdm that automatically detects its environment. - Interactive Terminal: Displays a normal progress bar. - Non-Interactive (HPC): Automatically disables the bar to prevent log file bloat in slurm.out.
"},{"location":"src/experiment_logger/#loggerlog_non_interactivemsg-str","title":"logger.log_non_interactive(msg: str)","text":"Prints a message only when running in non-interactive environments. Useful for high-level progress tracking (e.g., \"Epoch 5 Complete\") without interactive noise.
"},{"location":"src/experiment_logger/#loggersave_checkpointparams-step-prefixcheckpoint","title":"logger.save_checkpoint(params, step, prefix=\"checkpoint\")","text":"Saves model parameters using Flax serialization. - Local Location: runs/<run_name>/checkpoints/ - WandB Logic: Automatically uploads the .flax file as a model artifact for lineage tracking.
"}]}
\ No newline at end of file
diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py
index 9fd1507..0a75b8d 100644
--- a/src/brittle_star_project/configs/config_simulation.py
+++ b/src/brittle_star_project/configs/config_simulation.py
@@ -26,6 +26,9 @@ class SimulationSettings:
video_output_path: Optional[str] = None
# Camera ID to use for video recording (1 is usually the close-up camera)
camera_id: int = 1
+ video_width: int = 640
+ video_height: int = 480
+ video_fps: int = 60
# Optional override for the sidecar metadata YAML file.
# If None, it defaults to the model_path with a `_metadata.yaml` suffix.
diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
index b514c11..cfb60cd 100644
--- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
+++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
@@ -65,11 +65,21 @@ class BrittleStarJaxEnvWrapper:
def single_observation_space(self):
return self._env.observation_space
- def reset(self, seed: int = 0):
+ def reset(self, seed: int = 0, target_position: tuple[float, float] | None = None):
self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
- state = self._vectorized_reset(rng=env_rngs)
+
+ # If a target_position is provided, pass it through to the underlying env.reset
+ if target_position is None:
+ state = jax.jit(jax.vmap(lambda rng: self._env.reset(rng=rng)))(env_rngs)
+ else:
+ tp = jnp.asarray(target_position, dtype=jnp.float32)
+ tp_batched = jnp.tile(tp[None, :], (self._num_envs, 1))
+ state = jax.jit(jax.vmap(lambda rng, t: self._env.reset(rng=rng, target_position=t)))(
+ env_rngs, tp_batched
+ )
+
return state
def sample_actions(self):
diff --git a/src/brittle_star_project/environment/env_wrapper.py b/src/brittle_star_project/environment/env_wrapper.py
index 8b3e51f..0d2463c 100644
--- a/src/brittle_star_project/environment/env_wrapper.py
+++ b/src/brittle_star_project/environment/env_wrapper.py
@@ -62,9 +62,12 @@ class BrittleStarEnv:
return jax.random.PRNGKey(seed)
- def reset(self, *, seed: int = 0):
+ def reset(self, *, seed: int = 0, target_position: tuple[float, float, float] | None = None):
rng = self.make_rng(seed)
- state = self._env.reset(rng=rng)
+ if target_position is not None:
+ state = self._env.reset(rng=rng, target_position=target_position)
+ else:
+ state = self._env.reset(rng=rng)
return state
def render(self, *, state: Any):
diff --git a/src/brittle_star_project/evaluation/video.py b/src/brittle_star_project/evaluation/video.py
index b4b44db..cf2e5c2 100644
--- a/src/brittle_star_project/evaluation/video.py
+++ b/src/brittle_star_project/evaluation/video.py
@@ -25,6 +25,59 @@ def create_evaluation_dir(model_path: Path) -> Path:
return eval_dir
+def _ensure_offscreen_size(model, width: int, height: int) -> None:
+ vis_global = getattr(getattr(model, "vis", None), "global_", None)
+ if vis_global is None:
+ return
+ vis_global.offwidth = int(max(width, vis_global.offwidth))
+ vis_global.offheight = int(max(height, vis_global.offheight))
+
+
+def _apply_camera_overrides(
+ model,
+ *,
+ camera_fovy: dict[int, float] | None = None,
+ camera_xyz: tuple[
+ dict[int, float] | None,
+ dict[int, float] | None,
+ dict[int, float] | None,
+ ] = (None, None, None),
+) -> None:
+ if not camera_fovy and not (camera_xyz[0] or camera_xyz[1] or camera_xyz[2]):
+ return
+
+ ncam = int(getattr(model, "ncam", 0))
+ for cam_id, fovy in (camera_fovy or {}).items():
+ if cam_id < 0 or cam_id >= ncam:
+ raise ValueError(f"Camera id {cam_id} is out of range")
+ model.cam_fovy[cam_id] = float(fovy)
+
+ for cam_id, x in (camera_xyz[0] or {}).items():
+ if cam_id < 0 or cam_id >= ncam:
+ raise ValueError(f"Camera id {cam_id} is out of range")
+ model.cam_pos[cam_id][0] = float(x)
+
+ for cam_id, y in (camera_xyz[1] or {}).items():
+ if cam_id < 0 or cam_id >= ncam:
+ raise ValueError(f"Camera id {cam_id} is out of range")
+ model.cam_pos[cam_id][1] = float(y)
+
+ for cam_id, z in (camera_xyz[2] or {}).items():
+ if cam_id < 0 or cam_id >= ncam:
+ raise ValueError(f"Camera id {cam_id} is out of range")
+ model.cam_pos[cam_id][2] = float(z)
+
+
+def hex_to_rgba(hex_color: str, alpha: float) -> np.ndarray:
+ color = hex_color.lstrip("#")
+ if len(color) != 6:
+ raise ValueError(f"Expected a 6-digit hex color, got {hex_color!r}")
+ red = int(color[0:2], 16) / 255.0
+ green = int(color[2:4], 16) / 255.0
+ blue = int(color[4:6], 16) / 255.0
+ return np.asarray([red, green, blue, float(alpha)], dtype=np.float32)
+
+
def save_evaluation_metadata(
eval_dir: Path,
*,
@@ -66,6 +119,7 @@ def record_episode(
fps: int = 60,
width: int = 640,
height: int = 480,
+ target_xy: tuple[float, float] | None = None,
) -> EpisodeResult:
"""Run an episode headlessly and record a video using MuJoCo's Renderer and imageio.
@@ -92,10 +146,12 @@ def record_episode(
"Please install the evaluation dependencies: `uv pip install .[evaluation]`"
) from e
- state = env.reset(seed=seed)
+ state = env.reset(seed=seed, target_position=target_xy)
model = state.mj_model
data = state.mj_data
+ _ensure_offscreen_size(model, width, height)
+
renderer = mujoco.Renderer(model, width=width, height=height)
ep_return = 0.0
observations = _get_observations(state)
@@ -147,3 +203,133 @@ def record_episode(
final_xy_dist=final_dist,
initial_target_distance=initial_dist,
)
+
+
+def record_episode_multi_camera(
+ *,
+ env: BrittleStarEnv,
+ policy: ControlPolicy,
+ seed: int,
+ max_steps: int,
+ action_low: np.ndarray | None,
+ action_high: np.ndarray | None,
+ output_paths: dict[int, Path],
+ action_mask: np.ndarray | None = None,
+ camera_ids: list[int] | None = None,
+ camera_fovy: dict[int, float] | None = None,
+ camera_xyz: tuple[
+ dict[int, float] | None,
+ dict[int, float] | None,
+ dict[int, float] | None,
+ ] = (None, None, None),
+ target_xy: tuple[float, float] | None = None,
+ robot_color: str | None = None,
+ fps: int = 60,
+ width: int = 640,
+ height: int = 480,
+) -> EpisodeResult:
+ """Run one episode and render multiple camera views to separate files."""
+ try:
+ import imageio
+ import mujoco
+ except ImportError as e:
+ raise ImportError(
+ "Video recording requires 'imageio' and 'mujoco'. "
+ "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
+ ) from e
+
+ if camera_ids is None:
+ camera_ids = list(output_paths.keys())
+
+ for cam_id in camera_ids:
+ if cam_id not in output_paths:
+ raise ValueError(f"Missing output path for camera {cam_id}")
+
+ output_paths = {cam_id: output_paths[cam_id] for cam_id in camera_ids}
+
+ for path in output_paths.values():
+ path.parent.mkdir(parents=True, exist_ok=True)
+
+ state = env.reset(seed=seed, target_position=(target_xy[0], target_xy[1], 0.0))
+ model = state.mj_model
+ data = state.mj_data
+
+ _apply_camera_overrides(model, camera_fovy=camera_fovy, camera_xyz=camera_xyz)
+
+ if robot_color is not None:
+ robot_body_id = mujoco.mj_name2id(
+ model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
+ )
+ if robot_body_id < 0:
+ raise ValueError("Body 'BrittleStarMorphology/central_disk' not found in the model")
+
+ robot_rgba = hex_to_rgba(robot_color, 1.0)
+ body_parent = model.body_parentid
+ robot_body_ids = {int(robot_body_id)}
+
+ for body_id in range(1, int(model.nbody)):
+ current_body_id = int(body_id)
+ while current_body_id not in (-1, 0, int(robot_body_id)):
+ current_body_id = int(body_parent[current_body_id])
+ if current_body_id == int(robot_body_id):
+ robot_body_ids.add(body_id)
+
+ for geom_id in range(int(model.ngeom)):
+ if int(model.geom_bodyid[geom_id]) in robot_body_ids:
+ model.geom_rgba[geom_id][:] = robot_rgba
+
+ _ensure_offscreen_size(model, width, height)
+
+ renderer = mujoco.Renderer(model, width=width, height=height)
+ writers = {
+ cam_id: imageio.get_writer(str(path), fps=fps) for cam_id, path in output_paths.items()
+ }
+
+ ep_return = 0.0
+ observations = _get_observations(state)
+ prev_dist = _get_xy_distance_to_target(observations) if observations else None
+ initial_dist = prev_dist
+ reached_target = _target_reached(state=state)
+
+ steps = 0
+ try:
+ for _ in range(int(max_steps)):
+ for cam_id in camera_ids:
+ renderer.update_scene(data, camera=cam_id)
+ writers[cam_id].append_data(renderer.render())
+
+ 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 observations else None
+ 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
+
+ for cam_id in camera_ids:
+ renderer.update_scene(data, camera=cam_id)
+ writers[cam_id].append_data(renderer.render())
+ finally:
+ renderer.close()
+ for writer in writers.values():
+ writer.close()
+
+ final_dist = _get_xy_distance_to_target(observations) if observations else None
+ return EpisodeResult(
+ return_=ep_return,
+ length=steps,
+ reached_target=reached_target,
+ final_xy_dist=final_dist,
+ initial_target_distance=initial_dist,
+ )