fix(PR#41): applied comments for PR #41 review

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

configs/ppo/fast.yaml

@@ -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

configs/ppo/fast2.yaml

@@ -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

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))