From 34e8ba91c21a46f06cea9b91dd3863dfe48b1a8b Mon Sep 17 00:00:00 2001
From: Robin Meersman <echteenrobin@gmail.com>
Date: Sat, 18 Apr 2026 11:07:01 +0200
Subject: [PATCH] fix(PR#41): applied comments for PR #41 review

---
 configs/morphology/2_arms.yaml                |  2 +-
 configs/ppo/fast.yaml                         | 19 -------------------
 configs/ppo/fast2.yaml                        | 16 ----------------
 .../environment/BrittleStarJaxEnvWrapper.py   |  2 +-
 4 files changed, 2 insertions(+), 37 deletions(-)
 delete mode 100644 configs/ppo/fast.yaml
 delete mode 100644 configs/ppo/fast2.yaml

diff --git a/configs/morphology/2_arms.yaml b/configs/morphology/2_arms.yaml
index dc6f159..c301ec4 100644
--- a/configs/morphology/2_arms.yaml
+++ b/configs/morphology/2_arms.yaml
@@ -1,5 +1,5 @@
 # 2 Arms Morphology Configuration
 
-segments_per_arm: [4, 4]
+segments_per_arm: [4, 0, 4, 0, 0]
 use_p_control: true
 use_torque_control: false
\ No newline at end of file
diff --git a/configs/ppo/fast.yaml b/configs/ppo/fast.yaml
deleted file mode 100644
index ecd65fa..0000000
--- a/configs/ppo/fast.yaml
+++ /dev/null
@@ -1,19 +0,0 @@
-# Fast PPO Configuration
-# Lower timestep count for quick iterations/testing.
-
-learning_rate: 0.0005
-total_timesteps: 500000
-num_envs: 32
-num_steps: 128
-anneal_lr: true
-gamma: 0.99
-gae_lambda: 0.95
-num_minibatches: 4
-update_epochs: 4
-norm_adv: true
-clip_coef: 0.2
-clip_vloss: true
-ent_coef: 0.01
-vf_coef: 0.5
-max_grad_norm: 0.5
-target_kl: null
diff --git a/configs/ppo/fast2.yaml b/configs/ppo/fast2.yaml
deleted file mode 100644
index b2f629a..0000000
--- a/configs/ppo/fast2.yaml
+++ /dev/null
@@ -1,16 +0,0 @@
-learning_rate: 0.0003
-total_timesteps: 1228800
-num_envs: 32
-num_steps: 64
-anneal_lr: true
-gamma: 0.99
-gae_lambda: 0.95
-num_minibatches: 32
-update_epochs: 4
-norm_adv: true
-clip_coef: 0.2
-clip_vloss: true
-ent_coef: 0.005
-vf_coef: 1.0
-max_grad_norm: 0.5
-target_kl: null
\ No newline at end of file
diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
index 3122696..a5175a7 100644
--- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
+++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
@@ -27,7 +27,7 @@ class BrittleStarJaxEnvWrapper:
         )
 
         # Pre-compute masks for observation padding
-        self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm, (4, 4))
+        self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm)
 
         self._vectorized_reset = jax.jit(jax.vmap(self._env.reset))
         self._vectorized_step = jax.jit(jax.vmap(self._env.step))