[64, 64]) and utilize tanh
+ They are configured as standard dense networks with 3 hidden layers of 300 nodes each ([300, 300, 300]) and utilize tanh
activation functions.mean and log_std) using a single dense
diff --git a/search/search_index.json b/search/search_index.json
index 0f6f92a..be2b6da 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":"uv, including GPU configuration. For High-Performance Computing (HPC) setup details, see the HPC Guide..\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/design)","text":"If you are interested in the \"why did you do it like this?\"
/api)","text":"If you are interested in the \"how do I use it?\"
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":"Code readability is paramount, as code is read far more frequently than it is written.
dev branch serves as the integration branch for pushing and merging code. Only stable releases may be pushed to the main branch.git-lfs installed locally (see DEVELOPMENT.md for setup).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.uv using ruff and pre-commit hooks.This project supports AI-assisted development to enhance productivity, but contributors must take full responsibility for all AI-generated outputs.
.agents/skills/), including linting and testing workflows. Contributors should leverage these skills to maintain consistency and quality.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.
uv.lock manually.uv add <package>.uv lock --upgrade.uv sync --frozen.All developers must have Git LFS installed locally. This repository tracks model weights (.pt, .safetensors, etc.), recordings (.mp4), and datasets using Git LFS.
git lfs install after cloning this repository. If you are using the .devcontainer or flake.nix, LFS is typically available automatically.git lfs pull after installation to fetch the actual data files instead of the small pointer files.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":"post-create.sh script will:nvidia-smi.uv sync --frozen --extra cuda if a GPU is found.uv sync --frozen otherwise..venv to ensure persistence and performance..devcontainer/devcontainer.json file./workspaces/project/.venv.If you prefer not to use Docker:
uv sync --frozen (CPU) or uv sync --frozen --extra cuda (GPU).Verify your setup by running the JAX initialization test:
uv run pytest tests/test_jax_init.py\nIn 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.
This project uses a unified logging system through the experiment_logger package.
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":"$VSC_SCRATCH during the job (fast I/O) and copied to $VSC_DATA at the end for persistence.$VSC_DATA by mirroring configuration files. This avoids the 3GB home quota without requiring symlinks in the project root.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\nOur 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.
The donphan cluster does not support GPUs. Simply run the scripts without extra resource flags:
module swap cluster/donphan\nqsub scripts/hpc/train.pbs\nThese 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\nTo 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\nAfter 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.\"\nVerify 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)\nVerify GPU Access:
python -c \"import torch, jax; print(f'GPU: {torch.cuda.is_available()}'); print(f'JAX: {jax.devices()}')\"\nenv/hpc/requirements.txt is auto-generated from pyproject.toml. To regenerate:
uv run scripts/hpc/export_requirements.py\nModules listed in env/hpc/modules.txt are automatically excluded from the pip requirements to save space and use HPC-optimized binaries.
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:
The scripts/plots/analyze_comparisons.py script generates grouped bar charts comparing the performance of different architectures across various morphologies.
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\ninput_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.The script generates four key plots, each saved as both .png and .svg:
The scripts/plots/analyze_convergence.py script determines the convergence point of training runs.
uv run python scripts/plots/analyze_convergence.py --output_dir runs/convergence/\nFILE_MAPPING dictionary. Update these paths to point to your specific run evaluation files.--show_titles, --font_size, and --fig_width/height flags as the comparison script.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.We recommend using the SVG outputs for poster design in Figma:
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.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.
The data classes in env_config have default values as stated in the tutorials.
The Backend enum specifies either an MJC or MJX backend.
The Task enum specifies which task to use. 2 items are present:
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.
In your experiment config or via CLI:
python scripts/train.py evaluation.evaluate_checkpoints=true evaluation.eval_max_steps=5000\nResults are saved to runs/<run_dir>/metrics/checkpoint_evaluation.csv and synced to Weights & Biases if enabled.
To measure how well different controllers handle damage (amputations), use scripts/compare_models.py. This script performs a grid search over models x morphologies.
configs/evaluation.python scripts/compare_models.py evaluation=poster\nThe 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).
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\nThis 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:
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:
...) menu.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\nRun 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\nReproducing 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:
uv run python scripts/evaluate_checkpoints.py simulation.model_path=runs/your_run_dir/final_model.flax\nThis script runs deterministic rollouts for every checkpoint in runs/your_run_dir/checkpoints/.
runs/your_run_dir/metrics/checkpoint_evaluation.csv\nAnalyze 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).
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\nuv run python scripts/compare_models.py evaluation=poster\nThis 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.
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.
To simulate a model in the MuJoCo viewer:
uv run scripts/simulate.py simulation.model_path=runs/your_run/final_model.flax\nYou 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\nRecording 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\nVideos 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:
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/\nThis 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\nAlternatively, you can log in using the CLI:
uv run wandb login\nTo enable online sync during a training run, set logging.track=true on the command line:
uv run python scripts/train.py logging.track=true\nYou 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\"\nThese can also be set in your configuration YAML file under the logging key.
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/\nAccess 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/\nSee the detailed description in /scripts/analysis/README.md.
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\nEdit configs/experiment/my_experiment.yaml to set your experiment parameters:
# @package _global_\nexperiment:\n exp_name: \"my_custom_run\"\n seed: 42\nTo start a training run with the default settings defined in configs/main_config.yaml:
uv run python scripts/train.py\nTo run with your custom experiment file:
uv run python scripts/train.py experiment=my_experiment\nuv run python scripts/train.py ppo.learning_rate=0.001 ppo.num_envs=32 logging.track=true\nBy 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\nFor 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":"ConfigStore). Misspelled keys throw a ConfigAttributeError immediately.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).Run a quick test with minimal iterations:
python scripts/train.py experiment=dev_test ppo=fast\nTest a decentralized controller on a 3-arm robot:
python scripts/train.py architecture=decentralized morphology=3_arms\nRun stable PPO with WandB enabled (HPC submission scripts handle the hydra.run.dir redirection):
python scripts/train.py ppo=stable logging=wandb_enabled\nCheck if your configuration is valid without starting the simulation:
python scripts/train.py --cfg job\nconfigs/ and register the new dataclass in src/brittle_star_project/configs/register_configs.py.ConfigAttributeError, check for typos in your YAML keys or CLI overrides.experiment.base_run_dir to configure where logs and models are stored (defaults to runs/).python scripts/train.py experiment.base_run_dir=/path/to/custom/dirTo 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).
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 --> OutCritDecentralized Architecture
This pipeline utilizes the \"Centralized Training with Decentralized Execution\" principle, specifically the NerveNet-MLP variant.
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:
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\"|PropInspired 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:
[64, 64]) and utilize tanh activation functions.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
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
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:
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:
Local inputs, routed directly to specific nodes:
Outputs (action space)
The action space defines how the agent interacts with the environment.
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?
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:
Alternative state and action formulations include:
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_rotationThis 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\nTo 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:
References
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 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\nRequirements: - pandas - tensorboard - tensorflow-cpu (or tensorflow)
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.
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)\nUnifiedLogger","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.
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\")\nlogger.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.
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.
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":"uv, including GPU configuration. For High-Performance Computing (HPC) setup details, see the HPC Guide..\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/design)","text":"If you are interested in the \"why did you do it like this?\"
/api)","text":"If you are interested in the \"how do I use it?\"
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":"Code readability is paramount, as code is read far more frequently than it is written.
dev branch serves as the integration branch for pushing and merging code. Only stable releases may be pushed to the main branch.git-lfs installed locally (see DEVELOPMENT.md for setup).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.uv using ruff and pre-commit hooks.This project supports AI-assisted development to enhance productivity, but contributors must take full responsibility for all AI-generated outputs.
.agents/skills/), including linting and testing workflows. Contributors should leverage these skills to maintain consistency and quality.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.
uv.lock manually.uv add <package>.uv lock --upgrade.uv sync --frozen.All developers must have Git LFS installed locally. This repository tracks model weights (.pt, .safetensors, etc.), recordings (.mp4), and datasets using Git LFS.
git lfs install after cloning this repository. If you are using the .devcontainer or flake.nix, LFS is typically available automatically.git lfs pull after installation to fetch the actual data files instead of the small pointer files.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":"post-create.sh script will:nvidia-smi.uv sync --frozen --extra cuda if a GPU is found.uv sync --frozen otherwise..venv to ensure persistence and performance..devcontainer/devcontainer.json file./workspaces/project/.venv.If you prefer not to use Docker:
uv sync --frozen (CPU) or uv sync --frozen --extra cuda (GPU).Verify your setup by running the JAX initialization test:
uv run pytest tests/test_jax_init.py\nIn 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.
This project uses a unified logging system through the experiment_logger package.
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":"$VSC_SCRATCH during the job (fast I/O) and copied to $VSC_DATA at the end for persistence.$VSC_DATA by mirroring configuration files. This avoids the 3GB home quota without requiring symlinks in the project root.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\nOur 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.
The donphan cluster does not support GPUs. Simply run the scripts without extra resource flags:
module swap cluster/donphan\nqsub scripts/hpc/train.pbs\nThese 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\nTo 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\nAfter 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.\"\nVerify 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)\nVerify GPU Access:
python -c \"import torch, jax; print(f'GPU: {torch.cuda.is_available()}'); print(f'JAX: {jax.devices()}')\"\nenv/hpc/requirements.txt is auto-generated from pyproject.toml. To regenerate:
uv run scripts/hpc/export_requirements.py\nModules listed in env/hpc/modules.txt are automatically excluded from the pip requirements to save space and use HPC-optimized binaries.
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:
The scripts/plots/analyze_comparisons.py script generates grouped bar charts comparing the performance of different architectures across various morphologies.
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\ninput_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.The script generates four key plots, each saved as both .png and .svg:
The scripts/plots/analyze_convergence.py script determines the convergence point of training runs.
uv run python scripts/plots/analyze_convergence.py --output_dir runs/convergence/\nFILE_MAPPING dictionary. Update these paths to point to your specific run evaluation files.--show_titles, --font_size, and --fig_width/height flags as the comparison script.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.We recommend using the SVG outputs for poster design in Figma:
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.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.
The data classes in env_config have default values as stated in the tutorials.
The Backend enum specifies either an MJC or MJX backend.
The Task enum specifies which task to use. 2 items are present:
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.
In your experiment config or via CLI:
python scripts/train.py evaluation.evaluate_checkpoints=true evaluation.eval_max_steps=5000\nResults are saved to runs/<run_dir>/metrics/checkpoint_evaluation.csv and synced to Weights & Biases if enabled.
To measure how well different controllers handle damage (amputations), use scripts/compare_models.py. This script performs a grid search over models x morphologies.
configs/evaluation.python scripts/compare_models.py evaluation=poster\nThe 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).
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\nThis 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:
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:
...) menu.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\nRun 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\nReproducing 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:
uv run python scripts/evaluate_checkpoints.py simulation.model_path=runs/your_run_dir/final_model.flax\nThis script runs deterministic rollouts for every checkpoint in runs/your_run_dir/checkpoints/.
runs/your_run_dir/metrics/checkpoint_evaluation.csv\nAnalyze 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).
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\nuv run python scripts/compare_models.py evaluation=poster\nThis 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.
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.
To simulate a model in the MuJoCo viewer:
uv run scripts/simulate.py simulation.model_path=runs/your_run/final_model.flax\nYou 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\nRecording 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\nVideos 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:
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/\nThis 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\nAlternatively, you can log in using the CLI:
uv run wandb login\nTo enable online sync during a training run, set logging.track=true on the command line:
uv run python scripts/train.py logging.track=true\nYou 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\"\nThese can also be set in your configuration YAML file under the logging key.
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/\nAccess 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/\nSee the detailed description in /scripts/analysis/README.md.
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\nEdit configs/experiment/my_experiment.yaml to set your experiment parameters:
# @package _global_\nexperiment:\n exp_name: \"my_custom_run\"\n seed: 42\nTo start a training run with the default settings defined in configs/main_config.yaml:
uv run python scripts/train.py\nTo run with your custom experiment file:
uv run python scripts/train.py experiment=my_experiment\nuv run python scripts/train.py ppo.learning_rate=0.001 ppo.num_envs=32 logging.track=true\nBy 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\nFor 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":"ConfigStore). Misspelled keys throw a ConfigAttributeError immediately.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).Run a quick test with minimal iterations:
python scripts/train.py experiment=dev_test ppo=fast\nTest a decentralized controller on a 3-arm robot:
python scripts/train.py architecture=decentralized morphology=3_arms\nRun stable PPO with WandB enabled (HPC submission scripts handle the hydra.run.dir redirection):
python scripts/train.py ppo=stable logging=wandb_enabled\nCheck if your configuration is valid without starting the simulation:
python scripts/train.py --cfg job\nconfigs/ and register the new dataclass in src/brittle_star_project/configs/register_configs.py.ConfigAttributeError, check for typos in your YAML keys or CLI overrides.experiment.base_run_dir to configure where logs and models are stored (defaults to runs/).python scripts/train.py experiment.base_run_dir=/path/to/custom/dirTo 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).
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 --> OutCritDecentralized Architecture
This pipeline utilizes the \"Centralized Training with Decentralized Execution\" principle, specifically the NerveNet-MLP variant.
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:
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\"|PropInspired 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:
[300, 300, 300]) and utilize tanh activation functions.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
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
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:
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:
Local inputs, routed directly to specific nodes:
Outputs (action space)
The action space defines how the agent interacts with the environment.
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?
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:
Alternative state and action formulations include:
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_rotationThis 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\nTo 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:
References
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 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\nRequirements: - pandas - tensorboard - tensorflow-cpu (or tensorflow)
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.
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)\nUnifiedLogger","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.
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\")\nlogger.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.
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.