Merge pull request #19 from SELab-3-2026/docs/reward_and_mlp-design

Extended docs on mlp and reward design plus cleanup, cheers.

README.md

@@ -1,8 +1,11 @@
 # Brittle Star
 
-## Quick Start
+> What is the impact of different levels of controller-modularity on the learning-speed, coordination and tolerance for
+defects (e.g. amputations) in brittle-star-like robots trained with Reinforcement Learning?
 
-### Installation
+## Quick start
+
+### Local setup
 
 To set up the UV module, you can run the following command:
 
@@ -49,15 +52,10 @@ Override specific parameters:
 uv run python scripts/train.py --learning-rate 0.001 --num-envs 32 --track
 ```
 
-### Logging
-
-The training script uses a unified logging framework that:
-- Logs to **WandB** (when enabled)
-- Saves metrics to **local disk** (JSON files in `runs/`)
-- Displays progress in **stdout**
-
-All experiment data is preserved locally, even if WandB is unavailable.
-
 ## HPC
 
 See **[docs/HPC.md](docs/HPC.md)** for the full guide, including environment setup, cluster selection, interactive debugging, and job submission.
+
+## Documentation
+
+Please find all documentation and a starting point for more information in [corresponding README](./docs/README.md).

docs/README.md

@@ -0,0 +1,15 @@
+# Documentation
+
+## Design & architecture ([`/design`](./design/))
+
+- [Actor/critic architecture](./design/actor-critic.md): Description of the actor-critic pipeline.
+- [Communication](./design/communication.md): Message propagation, Nerve-Net style.
+- [Controllers](./design/controllers.md): Macroscopig brain toplogy, centralized, arm-level, segment-level.
+- [Input/output](./design/input_action_spaces.md): Description of the model's input and output.
+- [Learning algorithm](./design/learning_algorithm.md): RL techniques, i.e. PPO.
+- [Reward function](./design/learning_algorithm.md): Goals, fitness tracking, and reward structures.
+
+## API reference ([`/api`](./api/))
+
+- [Environment](./api/environment.md): MuJoCo environment interaction, state retrieval, and configuration.
+- [Simulate](./api/simulate.md): Simulation rendering.

docs/design/actor-critic.md

@@ -0,0 +1,121 @@
+# Actor-Critic Architecture
+
+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.
+
+```mermaid
+graph TD
+    Obs([Global Observation])
+
+    Sens[Sensor]
+    Act[Motor]
+    OutAct([Action Distribution<br/>mean, log_std])
+
+    Feat[Feature extractor]
+    Crit[Critic]
+    OutCrit([Value Estimate<br/>scalar])
+
+    Obs --> Sens
+    Obs --> Feat
+
+    Sens -->|"Hidden state"| Act
+    Feat -->|"Hidden state"| Crit
+
+    Act --> OutAct
+    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](./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:
+
+1. 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.
+2. 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.
+
+```mermaid
+graph TD
+    Obs([Local Observation])
+
+    Sens[Sensor]
+    Prop[Propagator]
+    Feat[Feature extractor]
+
+    Mot[Motor]
+    Crit[Critic]
+
+    OutMot([Action Distribution<br/>mean, log_std])
+    OutCrit([Value Estimate<br/>scalar])
+
+    Obs --> Sens
+    Sens -->|"Hidden state"| Prop
+    Obs --> Feat
+
+    Prop -->|"Hidden state"| Mot
+
+
+    Feat -->|"Hidden state"| Crit
+
+    Mot --> OutMot
+    Crit --> OutCrit
+
+    Prop -.->|"message passing"|Prop
+```
+
+## Implementation Details (Network Depth)
+
+Inspired by: https://iclr-blog-track.github.io/2022/03/25/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. 大トロ ・ Machine Learning. https://blog.otoro.net/2017/10/29/visual-evolution-strategies/
+- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. ‘Proximal Policy Optimization Algorithms’. 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. ‘NerveNet: Learning Structured Policy with Graph Neural Networks’. 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.

docs/design/input_action_spaces.md

@@ -0,0 +1,52 @@
+# Input (state) and output (action) spaces
+
+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 objective.
+
+- Joint positions: the current angles of all joints in the morphology.
+- Joint velocities: the current moving speed of the joints.
+- Goal vector: instad of just a scalar distance, the goal is represented asa a vector (distance and ange/direction) to
+  the target.
+
+**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.
+
+## Rationale
+
+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.
+- 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.
+- 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.
+
+## Limitations and alternatives
+
+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.

docs/design/reward_function.md

@@ -7,6 +7,11 @@ inputs must be distributed fairly to guarantee an objective comparison between d
 - 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.
+
+## From reward to PPO
+
+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.
 
 ## Rationale