From 9c20c28b2ad120b1debe0ed4a385e29c4b91314f Mon Sep 17 00:00:00 2001 From: Cedric Date: Wed, 6 May 2026 18:52:19 +0000 Subject: [PATCH] feat: split n flat n pad decentered --- scripts/simulate.py | 1 + .../environment/obs_processing.py | 134 +++++++++++------- .../trainers/PPOTrainer.py | 1 + tests/test_network_shapes.py | 1 + tests/test_target_direction.py | 1 + 5 files changed, 88 insertions(+), 50 deletions(-) diff --git a/scripts/simulate.py b/scripts/simulate.py index a2b121e..11b4a90 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -81,6 +81,7 @@ def main(dict_cfg: DictConfig) -> None: obs_processor = create_obs_processor( bounds_dict=training.obs_bounds.to_bounds_dict(), padding_masks=padding_masks, + segments_per_arm=env_morphology.segments_per_arm, ) # 6. Build environment diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index 96dd4ab..8b124fd 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -4,6 +4,10 @@ from typing import Dict, Tuple, Optional from brittle_star_project.environment.env_config import MorphMode +from experiment_logger import get_logger + +logger11 = get_logger() + _JOINT_SCALED_KEYS = frozenset( { "joint_position", @@ -20,12 +24,35 @@ _SEGMENT_SCALED_KEYS = frozenset( ) +def _build_joint_indices(segments_per_arm): + indices = [] + start = 0 + for segs in segments_per_arm: + # 2 joints per segment + count = segs * 2 + idx = jnp.arange(start, start + count) + indices.append(idx) + start += count + return indices + + +def _build_segment_indices(segments_per_arm): + indices = [] + start = 0 + for segs in segments_per_arm: + idx = jnp.arange(start, start + segs) + indices.append(idx) + start += segs + return indices + + def create_obs_processor( bounds_dict: Dict[str, Tuple[float, float]], num_arms: int, needed_copies: int, padding_masks: Optional[Dict] = None, morph_mode: MorphMode = MorphMode.CENTRALIZED, + segments_per_arm=[4, 4, 4, 4, 4], ): # made a set to allow O(1) search ordered_keys = frozenset( @@ -38,6 +65,8 @@ def create_obs_processor( "segment_contact", ] ) + segment_indices = _build_segment_indices(segments_per_arm) + joint_indices = _build_joint_indices(segments_per_arm) def _add_derived_features(obs: dict) -> dict: new_obs = dict(obs) @@ -69,96 +98,101 @@ def create_obs_processor( normalized[key] = arr return normalized - # ndarray (agents, features) def _pad_features(obs: dict, agent_count: int) -> dict: assert padding_masks is not None padded = {} - # arr shape = (agent_count, ...) TODO for key, arr in obs.items(): if key in _JOINT_SCALED_KEYS: - padded_arr = jnp.zeros( - (agent_count, padding_masks["target_size_2x"]), dtype=arr.dtype - ) - padded[key] = padded_arr.at[:, padding_masks["mask_2x"]].set(arr) + target_size = padding_masks["target_size_2x"] + out = jnp.zeros((agent_count, target_size), dtype=arr.dtype) + # place structured values at front, rest stays 0 + out = out.at[:, : arr.shape[1]].set(arr) + padded[key] = out + elif key in _SEGMENT_SCALED_KEYS: - padded_arr = jnp.zeros( - (agent_count, padding_masks["target_size_1x"]), dtype=arr.dtype - ) - padded[key] = padded_arr.at[:, padding_masks["mask_1x"]].set(arr) + target_size = padding_masks["target_size_1x"] + out = jnp.zeros((agent_count, target_size), dtype=arr.dtype) + out = out.at[:, : arr.shape[1]].set(arr) + padded[key] = out else: padded[key] = arr return padded - def _split_to_agents(obs: dict, agent_count: int) -> dict: - """ - Each observation type is now updated to (agent_count, feature_shape), - thus duplicating the observation for each agent - """ - + def _split_to_agents(obs: dict, morph_mode, segments_per_arm) -> dict: output = {} + num_arms = len(segments_per_arm) for key, arr in obs.items(): if key not in ordered_keys or arr.size == 0: continue + logger11.info(f"[INPUT] {key}: {arr.shape}") + if arr.ndim == 0: arr = arr.reshape(1) - output[key] = jnp.repeat(arr[None, :], agent_count, axis=0) + if morph_mode == MorphMode.CENTRALIZED: + out = arr.reshape(1, -1) + output[key] = out + continue + + if key in _SEGMENT_SCALED_KEYS: + per_agent = [jnp.take(arr, idx, axis=0) for idx in segment_indices] + out = jnp.stack(per_agent) + + elif key in _JOINT_SCALED_KEYS: + per_agent = [jnp.take(arr, idx, axis=0) for idx in joint_indices] + out = jnp.stack(per_agent) + + else: + out = jnp.repeat(arr[None, :], num_arms, axis=0) + + logger11.info(f"[OUTPUT] {key}: {out.shape}") + output[key] = out return output - # TODO: update def _flatten_features(obs: dict) -> jnp.ndarray: """ - Collapse all observations into a single array - """ + Input: + key -> (num_arms, feat_per_key) + Output: + (num_arms, total_features) + """ values = [] - for key in sorted(obs.keys()): - if key not in ordered_keys: + for key in ordered_keys: + if key not in obs: continue - # empty arrays handled in split_to_agents - v = obs[key] # (agent_count, feat) + arr = jnp.asarray(obs[key]) # (num_arms, feat) - # -------- CENTRALIZED -------- - if morph_mode == MorphMode.CENTRALIZED: - values.append(v.reshape(1, -1)) # (1, feat) + if arr.size == 0: continue - # -------- SCALE WITH SEGMENTS -------- - if key in _JOINT_SCALED_KEYS: - v = v.reshape(num_arms, -1) # (n_arms, 2) + if arr.ndim == 1: + arr = arr[:, None] - elif key in _SEGMENT_SCALED_KEYS: - v = v[:, None] # (segments, 1) + arr = arr.reshape(arr.shape[0], -1) - else: - # global key, broadcast to all nodes - v = jnp.repeat(v[None, :], num_arms, axis=0) # (num_arms, feat) + values.append(arr) - # RING + FULLY CONNECTED - v = v.reshape(num_arms, -1) - values.append(v) # (n_arms, feat) - - return jnp.concatenate(values, axis=-1) # (agent, feat) + return jnp.concatenate(values, axis=-1) # (num_arms, total_feat) def _process_single(obs_dict: dict) -> jnp.ndarray: - processed = _add_derived_features(obs_dict) # key |--> (feat-count,) - processed = _normalize_features(processed) # key |--> (feat-count,) - # TODO: split to agents - processed = _split_to_agents(processed, needed_copies) # key |--> (agents, feat-count) + processed = _add_derived_features(obs_dict) + processed = _normalize_features(processed) + # morph_mode = MorphMode.FULLY_CONNECTED + processed = _split_to_agents(processed, morph_mode, segments_per_arm) + # needed_copies = 5 if padding_masks is not None: - processed = _pad_features( - processed, agent_count=needed_copies - ) # key |--> (agents, feat-count') - for k, v in processed.items(): - print(k, v.shape) - + processed = _pad_features(processed, agent_count=needed_copies) + flat = _flatten_features(processed) # (num_arms, total_feat) + logger11.info(f"[FLATTENED FINAL] shape: {flat.shape}") + logger11.info(f"[PER AGENT] example row 0 shape: {flat[0].shape}") exit(1) return _flatten_features(processed) # (agents, feat) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 05761f5..5610fc9 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -367,6 +367,7 @@ class PPOTrainer: num_arms=self.num_arms, morph_mode=self.morph_mode, padding_masks=self.env.padding_masks, + segments_per_arm=self.segments_per_arm, ) action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) diff --git a/tests/test_network_shapes.py b/tests/test_network_shapes.py index 3cf4c14..3054510 100644 --- a/tests/test_network_shapes.py +++ b/tests/test_network_shapes.py @@ -32,6 +32,7 @@ def test_centralized_forward_pass_with_padding(): num_arms=num_arms, padding_masks=masks, morph_mode=MorphMode.CENTRALIZED, + segments_per_arm=segments_per_arm, ) global_state = obs_processor(amputated_obs) diff --git a/tests/test_target_direction.py b/tests/test_target_direction.py index f931c8a..8317fb2 100644 --- a/tests/test_target_direction.py +++ b/tests/test_target_direction.py @@ -53,6 +53,7 @@ def test_processor_converts_to_egocentric_direction(): num_arms=num_arms, padding_masks=env.padding_masks, morph_mode=MorphMode.CENTRALIZED, + segments_per_arm=segments_per_arm, ) env_state = env.reset(seed=42)