diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index 8b124fd..a57f00b 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -24,25 +24,27 @@ _SEGMENT_SCALED_KEYS = frozenset( ) -def _build_joint_indices(segments_per_arm): +def _build_joint_indices(segments_per_arm, indices_mlp): indices = [] start = 0 - for segs in segments_per_arm: + for i, segs in enumerate(segments_per_arm): # 2 joints per segment - count = segs * 2 - idx = jnp.arange(start, start + count) - indices.append(idx) - start += count + if i in indices_mlp: + count = segs * 2 + idx = jnp.arange(start, start + count) + indices.append(idx) + start += count return indices -def _build_segment_indices(segments_per_arm): +def _build_segment_indices(segments_per_arm, indices_mlp): indices = [] start = 0 - for segs in segments_per_arm: - idx = jnp.arange(start, start + segs) - indices.append(idx) - start += segs + for i, segs in enumerate(segments_per_arm): + if i in indices_mlp: + idx = jnp.arange(start, start + segs) + indices.append(idx) + start += segs return indices @@ -53,6 +55,7 @@ def create_obs_processor( padding_masks: Optional[Dict] = None, morph_mode: MorphMode = MorphMode.CENTRALIZED, segments_per_arm=[4, 4, 4, 4, 4], + agent_indices=[0, 1, 2, 3, 4], ): # made a set to allow O(1) search ordered_keys = frozenset( @@ -65,8 +68,8 @@ def create_obs_processor( "segment_contact", ] ) - segment_indices = _build_segment_indices(segments_per_arm) - joint_indices = _build_joint_indices(segments_per_arm) + segment_indices = _build_segment_indices(segments_per_arm, agent_indices) + joint_indices = _build_joint_indices(segments_per_arm, agent_indices) def _add_derived_features(obs: dict) -> dict: new_obs = dict(obs) @@ -120,34 +123,64 @@ def create_obs_processor( padded[key] = arr return padded - def _split_to_agents(obs: dict, morph_mode, segments_per_arm) -> dict: - output = {} - num_arms = len(segments_per_arm) + def _split_to_agents(obs: dict, morph_mode) -> dict: + total = 0 + for k, v in obs.items(): + if hasattr(v, "shape"): + size = v.size + logger11.info(f"[RAW] {k}: shape={v.shape}, size={size}") + total += size + else: + logger11.info(f"[RAW] {k}: non-array") + logger11.info(f"[RAW TOTAL FEATURES]: {total}") + output = {} + num_agents = needed_copies # IMPORTANT: number of MLPs 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) - + # -------- CENTRALIZED -------- if morph_mode == MorphMode.CENTRALIZED: - out = arr.reshape(1, -1) - output[key] = out + output[key] = arr.reshape(1, -1) continue + # -------- SEGMENTS -------- if key in _SEGMENT_SCALED_KEYS: - per_agent = [jnp.take(arr, idx, axis=0) for idx in segment_indices] + per_agent = [] + + for i, agent_id in enumerate(agent_indices): + idx = segment_indices[i] + taken = jnp.take(arr, idx, axis=0) # (segs, ...) + logger11.info(f"WHY {taken.shape}") + # pad to 4 + pad_len = 4 - taken.shape[0] + padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1)) + + per_agent.append(padded.reshape(-1)) + out = jnp.stack(per_agent) + # -------- JOINTS -------- elif key in _JOINT_SCALED_KEYS: - per_agent = [jnp.take(arr, idx, axis=0) for idx in joint_indices] + per_agent = [] + + for i, agent_id in enumerate(agent_indices): + idx = joint_indices[i] + taken = jnp.take(arr, idx, axis=0) # (joint_n, ...) + # pad to 8 + pad_len = 8 - taken.shape[0] + + padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1)) + per_agent.append(padded.reshape(-1)) + out = jnp.stack(per_agent) + # -------- GLOBAL -------- else: - out = jnp.repeat(arr[None, :], num_arms, axis=0) + out = jnp.repeat(arr[None, :], num_agents, axis=0) logger11.info(f"[OUTPUT] {key}: {out.shape}") output[key] = out @@ -185,15 +218,10 @@ def create_obs_processor( def _process_single(obs_dict: dict) -> jnp.ndarray: 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) + processed = _split_to_agents(processed, morph_mode) 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) return jax.jit(jax.vmap(_process_single)) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 5610fc9..18f30b7 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -353,6 +353,7 @@ class PPOTrainer: self.feature_extractor, self.critic, self.needed_copies, + self.agent_indices, ) = self._init_agent() self.sensor.apply = logged_jit(self.sensor.apply) @@ -368,6 +369,7 @@ class PPOTrainer: morph_mode=self.morph_mode, padding_masks=self.env.padding_masks, segments_per_arm=self.segments_per_arm, + agent_indices=self.agent_indices, ) action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) @@ -436,13 +438,18 @@ class PPOTrainer: def _init_agent(self): self.logger.info("[AGENT]: Initializing agent...") - + agent_indices = [0, 1, 2, 3, 4] match self.morph_mode: case MorphMode.CENTRALIZED: needed_copies = 1 case MorphMode.FULLY_CONNECTED | MorphMode.RING: + agent_mask = self.segments_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] needed_copies = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item() case MorphMode.SEGMENT: + agent_mask = self.segments_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item() needed_copies = ( self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum() ).item() @@ -462,7 +469,15 @@ class PPOTrainer: feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) critic = OneDenseLayerMLP() - return sensor, message_passer, actor, feature_extractor, critic, needed_copies + return ( + sensor, + message_passer, + actor, + feature_extractor, + critic, + needed_copies, + agent_indices, + ) def _init_agent_state(self) -> TrainState: self.logger.info("[AGENT STATE]: Initializing agent state...")