diff --git a/src/brittle_star_project/messager.py b/src/brittle_star_project/messager.py new file mode 100644 index 0000000..7178527 --- /dev/null +++ b/src/brittle_star_project/messager.py @@ -0,0 +1,26 @@ +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 7cdc1ff..df2c571 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -24,6 +24,14 @@ from brittle_star_project.MLPs.mlps import ( Storage, ) from brittle_star_project.ppo import PPO +from enum import Enum + + +class ObsMode(Enum): + CENTRALIZED = "centralized" + ARM = "arm" + SEGMENT = "segment" + # TODO: move to config _ALLOWED_OBS_KEYS = { @@ -41,13 +49,6 @@ _ALLOWED_OBS_KEYS = { # TODO: clip scaled reward? -@jax.jit -def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray: - """Extract xy_distance_to_target for all environments.""" - # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...) - return obs_dict["xy_distance_to_target"].squeeze(-1) # shape: (num_envs,) - - @jax.jit def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: return jnp.clip(action, low, high) @@ -70,18 +71,154 @@ def _normalize_obs(obs, mean, var, eps=1e-8): return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) -@jax.jit -def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: - """Convert the raw observation dict → flat array, filtering unwanted keys.""" +from enum import Enum - def _filter_and_flatten(o: dict) -> jnp.ndarray: + +class ObsMode(Enum): + CENTRALIZED = 0 + ARM = 1 + SEGMENT = 2 + + +@jax.jit +def _convert_obs_dict_to_array(obs_dict, obs_mode, segments_per_arm): + + num_segments = sum(segments_per_arm) + num_arms = len(segments_per_arm) + + def _filter_and_flatten(o): values = [] + for key in sorted(o.keys()): - if key in _ALLOWED_OBS_KEYS: # TODO: NORMALIZATION or .. of observations?? - v = o[key] - if v.size > 0: - values.append(jnp.asarray(v).flatten()) - return jnp.concatenate(values) + if key not in _ALLOWED_OBS_KEYS: + continue + + v = o[key] + if v.size == 0: + continue + + # -------- CENTRALIZED -------- + if obs_mode == 0: + values.append(v.reshape(v.shape[0], -1)) + continue + + # -------- SPLIT TO SEGMENTS -------- + if key in _JOINT_SCALED_KEYS: + v = v.reshape(v.shape[0], num_segments, 2) + + elif key in _SEGMENT_SCALED_KEYS: + v = v[..., None] # (env, segments, 1) + + else: + # global → broadcast + if obs_mode == 2: + v = jnp.repeat(v[:, None, :], num_segments, axis=1) + else: + v = jnp.repeat(v[:, None, :], num_arms, axis=1) + values.append(v) + continue + + # -------- SEGMENT MODE -------- + if obs_mode == 2: + values.append(v) + continue + + # -------- ARM MODE -------- + # reshape (segments → arms, seg_per_arm) + v = v.reshape(v.shape[0], num_arms, -1) + + values.append(v) + + # -------- CONCAT -------- + if obs_mode == 0: + return jnp.concatenate(values, axis=-1) + else: + return jnp.concatenate(values, axis=-1) + + return jax.vmap(_filter_and_flatten)(obs_dict) + + +from enum import Enum + + +class ObsMode(Enum): + CENTRALIZED = 0 # dirty dirty code, todo, mooove outta here + ARM = 1 + SEGMENT = 2 + + +# Observation keys whose size scales with the number of joints (2 per segment). +_JOINT_SCALED_KEYS = frozenset( + { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + } +) + +# Observation keys whose size scales with the number of segments (1 per segment). +_SEGMENT_SCALED_KEYS = frozenset( + { + "segment_contact", + } +) + + +@jax.jit +def _convert_obs_dict_to_array(obs_dict, obs_mode, segments_per_arm): + + num_segments = sum(segments_per_arm) + num_arms = len(segments_per_arm) + + def _filter_and_flatten(o): + values = [] + + for key in sorted(o.keys()): + if key not in _ALLOWED_OBS_KEYS: + continue + + v = o[key] + if v.size == 0: + continue + + # -------- CENTRALIZED -------- + if obs_mode == 0: + values.append(v.reshape(v.shape[0], -1)) + continue + + # -------- SPLIT TO SEGMENTS -------- + if key in _JOINT_SCALED_KEYS: + v = v.reshape(v.shape[0], num_segments, 2) + + elif key in _SEGMENT_SCALED_KEYS: + v = v[..., None] # (env, segments, 1) + + else: + # global → broadcast + if obs_mode == 2: + v = jnp.repeat(v[:, None, :], num_segments, axis=1) + else: + v = jnp.repeat(v[:, None, :], num_arms, axis=1) + values.append(v) + continue + + # -------- SEGMENT MODE -------- + if obs_mode == 2: + values.append(v) + continue + + # -------- ARM MODE -------- + # reshape (segments → arms, seg_per_arm) + v = v.reshape(v.shape[0], num_arms, -1) + + values.append(v) + + # -------- CONCAT -------- + if obs_mode == 0: + return jnp.concatenate(values, axis=-1) + else: + return jnp.concatenate(values, axis=-1) return jax.vmap(_filter_and_flatten)(obs_dict) @@ -367,6 +504,8 @@ class PPOTrainer: ) dummy_reset = self.env.reset(seed=0) + for k, v in dummy_reset.observations.items(): + print(k, v.shape) sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0] # take first env self.obs_mean = jnp.zeros((len(sample_obs),)) self.obs_var = jnp.ones((len(sample_obs),))