From 2ac21e2014eebc0361e1067e493ace9d0e317366 Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Mon, 4 May 2026 22:38:45 +0200 Subject: [PATCH] fix: linting + format --- .../environment/env_config.py | 2 +- .../environment/obs_processing.py | 1 - src/brittle_star_project/messager.py | 26 ------- .../trainers/PPOTrainer.py | 71 +------------------ tests/test_adjacency.py | 4 +- 5 files changed, 5 insertions(+), 99 deletions(-) delete mode 100644 src/brittle_star_project/messager.py diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index adf0736..33e8bdf 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -70,7 +70,7 @@ class EnvConfig: # Per docs in upstream env config: integer factors of 200. light_perlin_noise_scale: int = 0 - + @dataclass class ObservationBoundsConfig: """Physical observation bounds for deterministic min-max normalization.""" diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index 6b7cbe8..14b48f2 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/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( { diff --git a/src/brittle_star_project/messager.py b/src/brittle_star_project/messager.py deleted file mode 100644 index 7178527..0000000 --- a/src/brittle_star_project/messager.py +++ /dev/null @@ -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) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index a0af8b7..29c5d19 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/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 diff --git a/tests/test_adjacency.py b/tests/test_adjacency.py index 5e8a129..e6a1333 100644 --- a/tests/test_adjacency.py +++ b/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):