diff --git a/configs/morphology/5_arms_full_fullconnected.yaml b/configs/morphology/5_arms_full_fullconnected.yaml new file mode 100644 index 0000000..ab08e59 --- /dev/null +++ b/configs/morphology/5_arms_full_fullconnected.yaml @@ -0,0 +1,7 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 4, 4, 4] +use_p_control: true +use_torque_control: false +morph_mode: FULLY_CONNECTED diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index 635e7b8..79fea1c 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -119,7 +119,8 @@ def get_action_and_value( ): hidden_sensor = sensor_apply(params["sensor_params"], x) hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x) - hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor) + if message_passer is not None: + hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor) debug.callback(logger.debug, f"[SHAPE] hidden_sensor: {hidden_sensor.shape}") debug.callback(logger.debug, f"[SHAPE] hidden_critic: {hidden_critic.shape}") diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 9cbcf7c..f7c6489 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -514,7 +514,12 @@ class PPOTrainer: return jax.vmap(lambda x_in: self.message_passer.apply(p, x_in))(x) self._ppo = PPO( - self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature, apply_message_passer + self.ppo, + apply_sensor, + apply_actor, + apply_critic, + apply_feature, + apply_message_passer if self.message_passer is not None else None, ) self.agent_state = self._init_agent_state()