fix: linting + format

src/brittle_star_project/environment/obs_processing.py

@@ -3,7 +3,6 @@ import jax.numpy as jnp
 from typing import Dict, Tuple, Optional
 
 from brittle_star_project.environment.env_config import MorphMode
-from experiment_logger import get_logger
 
 _JOINT_SCALED_KEYS = frozenset(
     {

src/brittle_star_project/messager.py

@@ -1,26 +0,0 @@
-import jax
-import jax.numpy as jnp
-
-
-def message_passer(params, hidden, adjacency):
-    # take current hidden state per env, per agent, needs to communicate according to adjacency
-    # adjacency can be assumed to be the same per env? or mix it too idk
-    # rest is simple, use adjacency to make which hidden states we can combine
-    # repeat X times with the new combined vectors
-    # return result..
-    # use jax lax scan stuff or vmap for speed
-    def per_env(h):
-        # compute messages per agent
-        def compute_messages(h_i, h_all):
-            # broadcast h_i with all neighbors
-            h_i_rep = jnp.repeat(h_i[None, :], h_all.shape[0], axis=0)
-            msg_input = jnp.concatenate([h_i_rep, h_all], axis=-1)
-            messages = messager_apply(params, msg_input)
-            return messages
-
-        msgs = jax.vmap(compute_messages, in_axes=(0, None))(h, h)
-        msgs = msgs * adjacency[..., None]  # mask neighbors
-        agg = msgs.sum(axis=1)
-        return agg
-
-    return jax.vmap(per_env)(hidden)

src/brittle_star_project/trainers/PPOTrainer.py

@@ -18,87 +18,20 @@ from brittle_star_project.configs.main_config import BrittleStarConfig
 from brittle_star_project.dataclasses import EpisodeStatistics
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 from brittle_star_project.environment.obs_processing import create_obs_processor
-from brittle_star_project.MLPs.mlps import (
+from brittle_star_project.MLPs import (
     Actor,
     AgentParams,
     GenericDenseLayersWithActivation,
     MessagePasser,
     OneDenseLayerMLP,
     Storage,
+    build_adjacency,
 )
 from brittle_star_project.ppo import PPO
 from brittle_star_project.environment import MorphMode
 from brittle_star_project.utils import logged_jit
 
 logger11 = get_logger()
-# TODO: move to config
-
-# TODO: clip scaled reward?
-
-
-def build_adjacency(segments_per_arm, mode: MorphMode):
-    num_arms = sum(1 for s in segments_per_arm if s > 0)
-    num_segments = sum(segments_per_arm)
-
-    # FOR NOW SEMI HARDCODE:
-    # CENTRALIZED: 1 agent, no stress, adja = 1,1 = [[1]]
-    # FULLY CONNECTED: 5 agents: adj = alle 1
-    # CENTRAL DISK:#arms= 5 agents, only neighbor as adjacent so diagonal kinda..
-    # ARM = #segments agents: diago kinda, but extra, center ring too, put center mlps first or..
-
-    if mode == MorphMode.CENTRALIZED:
-        return jnp.ones((1, 1))
-
-    if mode == MorphMode.FULLY_CONNECTED:
-        adj = jnp.ones((num_arms, num_arms))  # everybody adjacent everybody
-        return adj
-
-    if mode == MorphMode.RING:  # ring
-        adj = jnp.zeros((num_arms, num_arms))
-        for i in range(num_arms):
-            adj = adj.at[i, i].set(1)  # self
-            adj = adj.at[i, (i - 1) % num_arms].set(1)
-            adj = adj.at[i, (i + 1) % num_arms].set(1)  # left and right..
-        return adj
-
-    if mode == MorphMode.SEGMENT:
-        num_nodes = num_arms + num_segments
-        adj = jnp.zeros((num_nodes, num_nodes))
-
-        # first ring
-        for i in range(num_arms):
-            # self
-            adj = adj.at[i, i].set(1)
-
-            # ring neighbors
-            adj = adj.at[i, (i - 1) % num_arms].set(1)
-            adj = adj.at[i, (i + 1) % num_arms].set(1)
-
-        # then segment chains
-        idx = 0
-        for arm_idx, seg_count in enumerate(segments_per_arm):
-            for i in range(seg_count):
-                seg_node = num_arms + idx + i
-
-                adj = adj.at[seg_node, seg_node].set(1)
-                if i > 0:
-                    adj = adj.at[seg_node, seg_node - 1].set(1)
-                if i < seg_count - 1:
-                    adj = adj.at[seg_node, seg_node + 1].set(1)
-
-            idx += seg_count
-
-        idx = 0
-        for arm_idx, seg_count in enumerate(segments_per_arm):
-            first_seg = num_arms + idx  # first segment of this arm
-
-            # connect ring node first segment
-            adj = adj.at[arm_idx, first_seg].set(1)
-            adj = adj.at[first_seg, arm_idx].set(1)
-
-            idx += seg_count
-
-        return adj
 
 
 @logged_jit

tests/test_adjacency.py

@@ -1,8 +1,8 @@
 import jax.numpy as jnp
-import pytest
 import numpy as np
 
-from brittle_star_project.trainers.PPOTrainer import build_adjacency, MorphMode
+from brittle_star_project.MLPs import build_adjacency
+from brittle_star_project.environment.env_config import MorphMode
 
 
 def assert_symmetric(adj):