diff --git a/src/brittle_star_project/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py index cc8c226..2c5989f 100644 --- a/src/brittle_star_project/MLPs/mlps.py +++ b/src/brittle_star_project/MLPs/mlps.py @@ -40,16 +40,17 @@ class Actor(nn.Module): class MessagePasser(nn.Module): hidden_dim: int num_propagation_steps: int + adj_matrix: jnp.ndarray @nn.compact - def __call__(self, x: jnp.ndarray, adj_matrix: jnp.ndarray): + def __call__(self, x: jnp.ndarray): for _ in range(self.num_propagation_steps): # (n_nodes, feat) messages = nn.Dense(self.hidden_dim)(x) messages = nn.tanh(messages) # note: if mean is wanted: adj_matrix / (adj.sum(axis=-1, keepdims=True) + 1e-8) - agg = adj_matrix + agg = self.adj_matrix aggregated = agg @ messages x_concat = jnp.concatenate([x, aggregated], axis=-1) diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index 87ebe10..5733177 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -72,6 +72,10 @@ def create_obs_processor( padded[key] = arr return padded + # TODO + def _split_to_agents() -> dict: + return {} + def _flatten_features(obs: dict) -> jnp.ndarray: ordered_keys = [ "disk_z_tilt", @@ -93,7 +97,7 @@ def create_obs_processor( if v.size == 0: continue - # reshape scalars + # reshape scalars to 1D array if v.ndim == 0: v = v.reshape(1) @@ -102,19 +106,8 @@ def create_obs_processor( values.append(v.reshape(1, -1)) # (1, feat) continue - # -------- SPLIT TO SEGMENTS -------- + # -------- SCALE WITH SEGMENTS -------- if key in _JOINT_SCALED_KEYS: - if morph_mode == MorphMode.SEGMENT: - joint_count = 3 - axis_per_joint = 2 - - center_size = num_arms * joint_count * axis_per_joint - v_center = v[:center_size].reshape( - num_arms, joint_count * axis_per_joint - ) # (arms, 6) - v_segs = v[center_size:].reshape(-1, 2) # (segs, 2) - values.append(jnp.concatenate([v_center, v_segs], axis=0)) # (arms+segs, ?) - continue v = v.reshape(num_arms, -1) # (n_arms, 2) elif key in _SEGMENT_SCALED_KEYS: @@ -122,25 +115,20 @@ def create_obs_processor( else: # global key, broadcast to all nodes - n_nodes = (num_segments + num_arms) if morph_mode == MorphMode.SEGMENT else num_arms - v = jnp.repeat(v[None, :], n_nodes, axis=0) # (n_nodes, feat) + v = jnp.repeat(v[None, :], num_arms, axis=0) # (num_arms, feat) - # -------- SEGMENT MODE -------- - if morph_mode == MorphMode.SEGMENT: - values.append(v) # (n_nodes, feat) - continue - - # -------- ARM MODE -------- + # RING + FULLY CONNECTED v = v.reshape(num_arms, -1) values.append(v) # (n_arms, feat) - return jnp.concatenate(values, axis=-1) + return jnp.concatenate(values, axis=-1) # (agent, feat) def _process_single(obs_dict: dict) -> jnp.ndarray: - processed = _add_derived_features(obs_dict) - processed = _normalize_features(processed) + processed = _add_derived_features(obs_dict) # key |--> (feat-count, feat-lengths, ) + processed = _normalize_features(processed) # key |--> (feat-count, feat-lengths, ) + # TODO: split to agents # key |--> (agents, feat-count, feat-lengths) if padding_masks is not None: - processed = _pad_features(processed) - return _flatten_features(processed) + processed = _pad_features(processed) # key |--> (agents, feat-count, feat-lengths') + return _flatten_features(processed) # (agents, feat) return jax.jit(jax.vmap(_process_single)) diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index 87e1dfe..635e7b8 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -119,7 +119,7 @@ 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(hidden_sensor) + 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 29c5d19..43de64b 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -62,7 +62,6 @@ def _get_action_and_value_noise( key, action_low, action_high, - adj_matrix: jnp.ndarray, ): # (B, n_nodes, feat) hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs) @@ -70,7 +69,7 @@ def _get_action_and_value_noise( if message_passer is not None: params = agent_state.params["message_passer_params"] # (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension - hidden = jax.vmap(lambda x: message_passer.apply(params, x, adj_matrix))(hidden) + hidden = jax.vmap(lambda x: message_passer.apply(params, x))(hidden) hidden_critic = apply_shared( feature_extractor, agent_state.params["feature_extractor_params"], next_obs @@ -100,7 +99,6 @@ def _step_once( carry, _, env_step_fn, - adj_matrix, sensor: nn.Module, feature_extractor: nn.Module, actor: nn.Module, @@ -121,7 +119,6 @@ def _step_once( key, action_low, action_high, - adj_matrix, ) logger11.debug(f"[_step_once] raw_action: {raw_action.shape}") logger11.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}") @@ -242,7 +239,6 @@ def _rollout_jit( message_passer: Optional[nn.Module], action_low, action_high, - adj_matrix, ): (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan( partial( @@ -255,7 +251,6 @@ def _rollout_jit( env_step_fn=step_env_fn, action_low=action_low, action_high=action_high, - adj_matrix=adj_matrix, ), (agent_state, episode_stats, next_obs, next_done, key, env_state), (), @@ -393,7 +388,6 @@ class PPOTrainer: message_passer=self.message_passer, action_low=action_low, action_high=action_high, - adj_matrix=self.adj, ) ) self._compute_gae_jit = logged_jit( @@ -419,7 +413,13 @@ class PPOTrainer: def apply_feature(p, x): return apply_shared(self.feature_extractor, p, x) - self._ppo = PPO(self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature) + def apply_message_passer(p, x): + assert self.message_passer is not None + 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.agent_state = self._init_agent_state() @@ -453,6 +453,7 @@ class PPOTrainer: MessagePasser( hidden_dim=300, num_propagation_steps=self.cfg.architecture.message_passing_steps or 4, + adj_matrix=self.adj, ) if self.morph_mode != MorphMode.CENTRALIZED else None @@ -517,7 +518,6 @@ class PPOTrainer: message_passer_params = self.message_passer.init( message_passer_key, self.sensor.apply(single_sensor_param, sample_obs), - self.adj, ) self.logger.debug( f"[_init_agent_state] message_passer_params: {