diff --git a/scripts/evaluate_checkpoints.py b/scripts/evaluate_checkpoints.py index e27f579..a63796b 100644 --- a/scripts/evaluate_checkpoints.py +++ b/scripts/evaluate_checkpoints.py @@ -26,7 +26,7 @@ from brittle_star_project.MLPs.mlps import ( ) from brittle_star_project.MLPs.adjancency_builder import build_adjacency from brittle_star_project.environment import MorphMode -from brittle_star_project.trainers.PPOTrainer import apply_per_node +from brittle_star_project.MLPs.routing import apply_per_node import logging import re from pathlib import Path @@ -219,8 +219,8 @@ def main(dict_cfg: DictConfig) -> None: eval_fn = build_eval_rollout_fn( env=env, obs_processor=obs_processor, - sensor_apply=lambda p, x: apply_per_node(sensor, p, x), - actor_apply=lambda p, x: apply_per_node(actor, p, x), + sensor_apply=lambda p, x: apply_per_node(sensor.apply, p, x), + actor_apply=lambda p, x: apply_per_node(actor.apply, p, x), message_passer_apply=(None if message_passer is None else message_passer.apply), action_low=action_low, action_high=action_high, diff --git a/src/brittle_star_project/MLPs/routing.py b/src/brittle_star_project/MLPs/routing.py new file mode 100644 index 0000000..4033cb9 --- /dev/null +++ b/src/brittle_star_project/MLPs/routing.py @@ -0,0 +1,22 @@ +"""Shared JAX routing utilities for decentralized multi-agent models.""" + +import jax + + +def apply_per_node(apply_fn, params, x): + """Apply a Flax module independently to each node. + + Args: + apply_fn: The module's ``apply`` method (e.g. ``sensor.apply``). + params: Per-node parameters with shape ``(num_nodes, ...)``. + x: Input tensor with shape ``(batch, num_nodes, features)``. + + Returns: + Output tensor with shape ``(batch, num_nodes, out_features)``. + """ + + def apply_single_node(p, x_node): + # x_node: (batch, feat) — one node's input across the batch + return jax.vmap(lambda xi: apply_fn(p, xi))(x_node) + + return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x) diff --git a/src/brittle_star_project/evaluation/policy.py b/src/brittle_star_project/evaluation/policy.py index 5b12d07..00c1e06 100644 --- a/src/brittle_star_project/evaluation/policy.py +++ b/src/brittle_star_project/evaluation/policy.py @@ -7,6 +7,7 @@ import jax import jax.numpy as jnp import numpy as np +from brittle_star_project.MLPs.routing import apply_per_node from brittle_star_project.evaluation.checkpoint import load_params @@ -147,22 +148,12 @@ class PolicyAgent: obs_processor=obs_processor, ) - def _apply_per_node(self, net, params, x): - # params: (nodes, ...) - # x: (batch, nodes, feat) - - def apply_single_node(p, x_node): - # x_node: (batch, feat) - return jax.vmap(lambda xi: net.apply(p, xi))(x_node) - - return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x) - def act(self, *, observations: dict[str, Any]) -> np.ndarray: """Return deterministic action (actor mean, no exploration noise).""" batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations) obs = self._obs_processor(batched_obs) - hidden = self._apply_per_node(self._sensor, self._params["sensor_params"], obs) + hidden = apply_per_node(self._sensor.apply, self._params["sensor_params"], obs) if self._message_passer is not None: mp_params = self._params.get("message_passer_params") @@ -172,6 +163,6 @@ class PolicyAgent: ) hidden = jax.vmap(lambda x: self._message_passer.apply(mp_params, x))(hidden) - mean, _log_std = self._apply_per_node(self._actor, self._params["actor_params"], hidden) + mean, _log_std = apply_per_node(self._actor.apply, self._params["actor_params"], hidden) return np.asarray(mean, dtype=np.float32).ravel() diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 76bb128..98d4045 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -23,6 +23,7 @@ from brittle_star_project.evaluation.evaluate_mjx import ( build_eval_rollout_fn, evaluate_checkpoint_mjx, ) +from brittle_star_project.MLPs.routing import apply_per_node from brittle_star_project.MLPs.mlps import ( Actor, AgentParams, @@ -71,7 +72,7 @@ def _get_action_and_value_noise( action_high, ): # (B, n_nodes, feat) - hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs) + hidden = apply_per_node(sensor.apply, agent_state.params["sensor_params"], next_obs) if message_passer is not None: params = agent_state.params["message_passer_params"] @@ -82,7 +83,7 @@ def _get_action_and_value_noise( feature_extractor, agent_state.params["feature_extractor_params"], next_obs ) - mean, log_std = apply_per_node(actor, agent_state.params["actor_params"], hidden) + mean, log_std = apply_per_node(actor.apply, agent_state.params["actor_params"], hidden) log_std = jnp.clip(log_std, -5, 2) key, subkey = jax.random.split(key) noise = jax.random.normal(subkey, shape=mean.shape) @@ -272,17 +273,6 @@ def _step_env_wrapped( ) -def apply_per_node(net, params, x): - # params: (nodes, ...) - # x: (batch, nodes, feat) - - def apply_single_node(p, x_node): - # x_node: (batch, feat) - return jax.vmap(lambda xi: net.apply(p, xi))(x_node) - - return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x) - - def apply_shared(net, params, x): # x: (batch, nodes, feat) # If the critic expects a single vector per environment: @@ -516,10 +506,10 @@ class PPOTrainer: ) def apply_sensor(p, x): - return apply_per_node(self.sensor, p, x) + return apply_per_node(self.sensor.apply, p, x) def apply_actor(p, x): - return apply_per_node(self.actor, p, x) + return apply_per_node(self.actor.apply, p, x) def apply_critic(p, x): return apply_shared(self.critic, p, x) @@ -903,8 +893,8 @@ class PPOTrainer: self._eval_fn = build_eval_rollout_fn( env=self.env, obs_processor=self.obs_processor, - sensor_apply=lambda p, x: apply_per_node(self.sensor, p, x), - actor_apply=lambda p, x: apply_per_node(self.actor, p, x), + sensor_apply=lambda p, x: apply_per_node(self.sensor.apply, p, x), + actor_apply=lambda p, x: apply_per_node(self.actor.apply, p, x), message_passer_apply=( None if self.message_passer is None else self.message_passer.apply ),