diff --git a/configs/architecture/decentralized.yaml b/configs/architecture/decentralized.yaml index 4b9fffb..70b0fb2 100644 --- a/configs/architecture/decentralized.yaml +++ b/configs/architecture/decentralized.yaml @@ -26,7 +26,7 @@ critic: activation: "tanh" # Synchronous message-passing rounds per control step -message_passing_steps: 1 +message_passing_steps: 4 # Connectivity topology (e.g., ring, fully_connected) -topology_type: "ring" +topology_type: "fully_connected" diff --git a/configs/environment/dir_loc_further.yaml b/configs/environment/dir_loc_further.yaml new file mode 100644 index 0000000..236f0c5 --- /dev/null +++ b/configs/environment/dir_loc_further.yaml @@ -0,0 +1,2 @@ +simulation_time: 50000.0 +target_distance: 3.0 \ No newline at end of file diff --git a/configs/experiment/long_2arm.yaml b/configs/experiment/long_2arm.yaml new file mode 100644 index 0000000..ae4f4b0 --- /dev/null +++ b/configs/experiment/long_2arm.yaml @@ -0,0 +1,6 @@ +# Testing chicken dinner 4 but further distance. + +exp_name: "long2arm" +seed: 123 +torch_deterministic: true +cuda: true \ No newline at end of file diff --git a/configs/morphology/2_arms_decentralized.yaml b/configs/morphology/2_arms_decentralized.yaml new file mode 100644 index 0000000..9fe3160 --- /dev/null +++ b/configs/morphology/2_arms_decentralized.yaml @@ -0,0 +1,6 @@ +# 2 Arms Morphology Configuration + +segments_per_arm: [4, 0, 4, 0, 0] +use_p_control: true +use_torque_control: false +morph_mode: FULLY_CONNECTED \ No newline at end of file diff --git a/configs/morphology/5_arms_damaged.yaml b/configs/morphology/5_arms_damaged.yaml new file mode 100644 index 0000000..2afd643 --- /dev/null +++ b/configs/morphology/5_arms_damaged.yaml @@ -0,0 +1,6 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 0, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/5_arms_full_fullconnected.yaml b/configs/morphology/5_arms_full_fullconnected.yaml new file mode 100644 index 0000000..ab08e59 --- /dev/null +++ b/configs/morphology/5_arms_full_fullconnected.yaml @@ -0,0 +1,7 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 4, 4, 4] +use_p_control: true +use_torque_control: false +morph_mode: FULLY_CONNECTED diff --git a/configs/ppo/chickendinnerwinner.yaml b/configs/ppo/chickendinnerwinner.yaml new file mode 100644 index 0000000..cf5e3d6 --- /dev/null +++ b/configs/ppo/chickendinnerwinner.yaml @@ -0,0 +1,16 @@ +anneal_lr: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.001 +gae_lambda: 0.95 +gamma: 0.99 +learning_rate: 0.0001 +max_grad_norm: 0.5 +norm_adv: true +num_envs: 32 +num_minibatches: 32 +num_steps: 64 +target_kl: 0.02 +total_timesteps: 12288000 +update_epochs: 4 +vf_coef: 1.0 \ No newline at end of file diff --git a/configs/ppo/smoke_test.yaml b/configs/ppo/smoke_test.yaml index 40e4c07..4b7bb84 100644 --- a/configs/ppo/smoke_test.yaml +++ b/configs/ppo/smoke_test.yaml @@ -2,9 +2,9 @@ # Lower timestep count for quick iterations/testing. learning_rate: 0.0005 -total_timesteps: 65536 -num_envs: 512 -num_steps: 128 +total_timesteps: 1024 +num_envs: 32 +num_steps: 32 anneal_lr: true gamma: 0.99 gae_lambda: 0.95 diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md deleted file mode 100644 index 92f23a1..0000000 --- a/experiments/debug-experiment-10042026/used_variables.md +++ /dev/null @@ -1,107 +0,0 @@ -## Default envconfig -task: Task = Task.DIRECTED_LOCOMOTION -simulation_time: float = 500.0 -num_physics_steps_per_control_step: int = 10 -time_scale: int = 2 -camera_ids: list[int] = field(default_factory=lambda: [0, 1]) -render_size: tuple[int, int] = (480, 640) -joint_randomization_noise_scale: float = 0.0 -target_distance: float = 3.0 -light_perlin_noise_scale: int = 0 - - -## Default ppoargs -seed: int = 1 -torch_deterministic: bool = True -cuda: bool = True -track: bool = False -checkpoint_frequency: int = 100 -learning_rate: float = 2.5e-4 -anneal_lr: bool = True -gamma: float = 0.99 -gae_lambda: float = 0.95 -update_epochs: int = 4 -norm_adv: bool = True -clip_vloss: bool = True -max_grad_norm: float = 0.5 -target_kl: float | None = None -batch_size: int = 0 -minibatch_size: int = 0 -num_iterations: int = 0 - -## Used config file: (hpc/debug.yaml) -exp_name: "debug-experiment" -seed: 42 -track: true -wandb_project_name: "Let's-find-that-bug" -wandb_entity: "SEL3-2026-Groep-4" -run_dir: "/data/gent/465/vsc46589" -num_envs: 32 -num_steps: 32 -num_minibatches: 32 -total_timesteps: 409600 -num_arms: 2 -cuda: true - -ent_coef: 0.005 -vf_coef: 1.0 -clip_coef: 0.2 - -anneal_lr: true -learning_rate: 0.0003 - -## Arena config: -size: tuple[float, float] = (10.0, 5.0) -sand_ground_color: bool = True -attach_target: bool = True -wall_height: float = 1.5 -wall_thickness: float = 0.1 - -## Morphology: -num_segments_per_arm: int = 4 -use_p_control: bool = True -use_torque_control: bool = False - -## MLPs: -### Sensor & Feature_extractor: -Both with 3 layers of 300 neurons per layer. - -class GenericDenseLayersWithActivation(nn.Module): - layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64]) - activation: Callable = nn.tanh - - @nn.compact - def __call__(self, x): - for size in self.layer_sizes: - x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x) - x = self.activation(x) - return x - -### Actor: -class Actor(nn.Module): - action_dim: int - @nn.compact - def __call__(self, x): - mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x) - log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,)) - return mean, log_std - -### Critic: -class OneDenseLayerMLP(nn.Module): - @nn.compact - def __call__(self, x): - return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x) - -### Observations: -_ALLOWED_OBS_KEYS = { - "joint_position", - "joint_velocity", - "joint_actuator_force", - "actuator_force", - "disk_position", - "disk_rotation", - "disk_linear_velocity", - "disk_angular_velocity", - "unit_xy_direction_to_target", - "xy_distance_to_target", -} \ No newline at end of file diff --git a/scripts/simulate.py b/scripts/simulate.py index a2b121e..a56e970 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -17,12 +17,14 @@ import numpy as np from omegaconf import DictConfig, OmegaConf import yaml +import jax.numpy as jnp + from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory from brittle_star_project.configs.main_config import BrittleStarConfig from brittle_star_project.configs.register_configs import register_configs from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks from brittle_star_project.environment.obs_processing import create_obs_processor -from brittle_star_project.environment.env_config import MorphologyConfig +from brittle_star_project.environment.env_config import MorphMode, MorphologyConfig from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs from brittle_star_project.evaluation.policy import PolicyAgent @@ -32,6 +34,7 @@ from brittle_star_project.evaluation.video import ( create_evaluation_dir, save_evaluation_metadata, ) +from brittle_star_project.MLPs.adjancency_builder import build_adjacency @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") @@ -78,9 +81,34 @@ def main(dict_cfg: DictConfig) -> None: segments_per_arm=env_morphology.segments_per_arm, reference_segments_per_arm=training.morphology.segments_per_arm, ) + + segs_per_arm = jnp.array(env_morphology.segments_per_arm) + + needed_copies = 0 + agent_indices = [0, 1, 2, 3, 4] + match env_morphology.morph_mode: + case MorphMode.CENTRALIZED: + needed_copies = 1 + case MorphMode.FULLY_CONNECTED | MorphMode.RING: + agent_mask = segs_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item() + case MorphMode.SEGMENT: + agent_mask = segs_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item() + needed_copies = (segs_per_arm.sum() + jnp.where(segs_per_arm > 0, 1, 0).sum()).item() + + num_arms = jnp.where(segs_per_arm > 0, 1, 0).sum().item() + obs_processor = create_obs_processor( bounds_dict=training.obs_bounds.to_bounds_dict(), padding_masks=padding_masks, + needed_copies=needed_copies, + num_arms=num_arms, + morph_mode=env_morphology.morph_mode, + segments_per_arm=env_morphology.segments_per_arm, + agent_indices=agent_indices, ) # 6. Build environment @@ -104,11 +132,24 @@ def main(dict_cfg: DictConfig) -> None: state0 = env.reset(seed=seed) # Calculate the action dimension the model was trained with - trained_action_dim = sum(training.morphology.segments_per_arm) * 2 + trained_action_dim = raw_env.action_space.shape[0] // needed_copies # 7. Load policy + message_passing_steps = (metadata.get("architecture", {}) or {}).get("message_passing_steps") + if message_passing_steps is None: + message_passing_steps = 4 + message_passing_steps = int(message_passing_steps) + + adj_matrix = None + if env_morphology.morph_mode != MorphMode.CENTRALIZED: + adj_matrix = build_adjacency(env_morphology.segments_per_arm, env_morphology.morph_mode) + policy = PolicyAgent.from_checkpoint( - model_path, action_dim=trained_action_dim, obs_processor=obs_processor + model_path, + action_dim=trained_action_dim, + obs_processor=obs_processor, + message_passing_steps=message_passing_steps, + adj_matrix=adj_matrix, ) # Convert the JAX boolean mask to a numpy array for easy indexing diff --git a/scripts/simulate.sh b/scripts/simulate.sh new file mode 100755 index 0000000..1e6901d --- /dev/null +++ b/scripts/simulate.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash + +path=$1 + +uv run simulate.py \ + simulation.model_path="$path"/final_model.flax \ + simulation.record_video=True \ + simulation.video_output_path=../vids/simulation.mp4 \ + simulation.max_steps=10000 diff --git a/src/brittle_star_project/MLPs/__init__.py b/src/brittle_star_project/MLPs/__init__.py new file mode 100644 index 0000000..8f9c7f7 --- /dev/null +++ b/src/brittle_star_project/MLPs/__init__.py @@ -0,0 +1,19 @@ +from .mlps import ( + GenericDenseLayersWithActivation, + OneDenseLayerMLP, + Actor, + MessagePasser, + AgentParams, + Storage, +) +from .adjancency_builder import build_adjacency + +__all__ = [ + "GenericDenseLayersWithActivation", + "OneDenseLayerMLP", + "Actor", + "MessagePasser", + "AgentParams", + "Storage", + "build_adjacency", +] diff --git a/src/brittle_star_project/MLPs/adjancency_builder.py b/src/brittle_star_project/MLPs/adjancency_builder.py new file mode 100644 index 0000000..878a02c --- /dev/null +++ b/src/brittle_star_project/MLPs/adjancency_builder.py @@ -0,0 +1,67 @@ +from brittle_star_project.environment.env_config import MorphMode +import jax.numpy as jnp + + +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 diff --git a/src/brittle_star_project/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py index 5e2deb5..2c5989f 100644 --- a/src/brittle_star_project/MLPs/mlps.py +++ b/src/brittle_star_project/MLPs/mlps.py @@ -1,11 +1,11 @@ from dataclasses import dataclass, fields, field -import flax import flax.linen as nn import jax.numpy as jnp import jax.tree_util from typing import Sequence, Callable from flax.linen.initializers import constant, orthogonal +from flax.core import FrozenDict # semi generic so we can easily make a config for it in experiments @@ -37,30 +37,56 @@ class Actor(nn.Module): return mean, log_std +class MessagePasser(nn.Module): + hidden_dim: int + num_propagation_steps: int + adj_matrix: jnp.ndarray + + @nn.compact + def __call__(self, x: jnp.ndarray): + for _ in range(self.num_propagation_steps): + # (n_nodes, feat) + messages = nn.Dense(self.hidden_dim)(x) + messages = nn.tanh(messages) + + # note: if mean is wanted: adj_matrix / (adj.sum(axis=-1, keepdims=True) + 1e-8) + agg = self.adj_matrix + aggregated = agg @ messages + + x_concat = jnp.concatenate([x, aggregated], axis=-1) + + gate = nn.sigmoid(nn.Dense(self.hidden_dim)(x_concat)) + candidate = nn.tanh(nn.Dense(self.hidden_dim)(x_concat)) + x = gate * x + (1 - gate) * candidate + + return x + + @jax.tree_util.register_dataclass @dataclass class AgentParams: - sensor_params: flax.core.FrozenDict - actor_params: flax.core.FrozenDict - critic_params: flax.core.FrozenDict - feature_extractor_params: flax.core.FrozenDict + sensor_params: FrozenDict | dict + actor_params: FrozenDict | dict + critic_params: FrozenDict | dict + feature_extractor_params: FrozenDict | dict + message_passer_params: FrozenDict | dict @jax.tree_util.register_dataclass @dataclass class Storage: - obs: jnp.array - actions: jnp.array - logprobs: jnp.array - dones: jnp.array - values: jnp.array - advantages: jnp.array - returns: jnp.array - rewards: jnp.array + obs: jnp.ndarray + actions: jnp.ndarray + logprobs: jnp.ndarray + dones: jnp.ndarray + values: jnp.ndarray + advantages: jnp.ndarray + returns: jnp.ndarray + rewards: jnp.ndarray - raw_actions: jnp.ndarray = None # before clipping - means: jnp.ndarray = None # policy mean - stds: jnp.ndarray = None # policy std + raw_actions: jnp.ndarray | None = None # before clipping + means: jnp.ndarray | None = None # policy mean + stds: jnp.ndarray | None = None # policy std def replace(self, **kwargs) -> "Storage": fs = fields(self) diff --git a/src/brittle_star_project/dataclasses/EpisodeStatistics.py b/src/brittle_star_project/dataclasses/EpisodeStatistics.py index 0b2832f..ff2982e 100644 --- a/src/brittle_star_project/dataclasses/EpisodeStatistics.py +++ b/src/brittle_star_project/dataclasses/EpisodeStatistics.py @@ -4,7 +4,7 @@ import jax.numpy as jnp @flax.struct.dataclass class EpisodeStatistics: - episode_returns: jnp.array - episode_lengths: jnp.array - returned_episode_returns: jnp.array - returned_episode_lengths: jnp.array + episode_returns: jnp.ndarray + episode_lengths: jnp.ndarray + returned_episode_returns: jnp.ndarray + returned_episode_lengths: jnp.ndarray diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index 77d106c..7a9a9eb 100644 --- a/src/brittle_star_project/environment/__init__.py +++ b/src/brittle_star_project/environment/__init__.py @@ -1,4 +1,4 @@ -from .env_config import ArenaConfig, EnvConfig, MorphologyConfig +from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, MorphMode from .env_types import Backend, Task from .env_wrapper import BrittleStarEnv from .factory import BrittleStarEnvFactory @@ -13,6 +13,7 @@ __all__ = [ "Task", "BrittleStarEnv", "BrittleStarEnvFactory", + "MorphMode", "create_obs_processor", "compute_padding_masks", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 22b8967..33e8bdf 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -1,10 +1,18 @@ from __future__ import annotations from dataclasses import dataclass, field +from enum import Enum from .env_types import Task +class MorphMode(Enum): + CENTRALIZED = 0 + FULLY_CONNECTED = 1 + RING = 2 + SEGMENT = 3 + + @dataclass class MorphologyConfig: """Brittle star morphology configuration. @@ -20,6 +28,7 @@ class MorphologyConfig: segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4]) use_p_control: bool = True use_torque_control: bool = False + morph_mode: MorphMode = MorphMode.CENTRALIZED @property def num_arms(self) -> int: diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index b18ee8c..f38d52b 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -2,6 +2,12 @@ import jax 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 + +logger = get_logger() + _JOINT_SCALED_KEYS = frozenset( { "joint_position", @@ -18,9 +24,53 @@ _SEGMENT_SCALED_KEYS = frozenset( ) +def _build_joint_indices(segments_per_arm, indices_mlp): + indices = [] + start = 0 + for i, segs in enumerate(segments_per_arm): + # 2 joints per segment + 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, indices_mlp): + indices = [] + start = 0 + 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 + + def create_obs_processor( - bounds_dict: Dict[str, Tuple[float, float]], padding_masks: Optional[Dict] = None + 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], + agent_indices=[0, 1, 2, 3, 4], ): + # made a set to allow O(1) search + ordered_keys = frozenset( + [ + "disk_z_tilt", + "joint_actuator_force", + "joint_position", + "joint_velocity", + "robot_direction_to_target", + "segment_contact", + ] + ) + 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) if "disk_rotation" in new_obs: @@ -51,41 +101,91 @@ def create_obs_processor( normalized[key] = arr return normalized - def _pad_features(obs: dict) -> dict: - padded = {} + def _split_to_agents(obs: dict, morph_mode) -> dict: + output = {} + num_agents = needed_copies # IMPORTANT: number of MLPs + + segs_per_arm = 4 + joints_per_segment = 2 + joints_per_arm = segs_per_arm * joints_per_segment for key, arr in obs.items(): - if key in _JOINT_SCALED_KEYS: - padded_arr = jnp.zeros(padding_masks["target_size_2x"], dtype=arr.dtype) - padded[key] = padded_arr.at[padding_masks["mask_2x"]].set(arr) - elif key in _SEGMENT_SCALED_KEYS: - padded_arr = jnp.zeros(padding_masks["target_size_1x"], dtype=arr.dtype) - padded[key] = padded_arr.at[padding_masks["mask_1x"]].set(arr) + if arr.size == 0: + continue + + if arr.ndim == 0: + arr = arr.reshape(1) + + if key in _SEGMENT_SCALED_KEYS: + per_agent = [] + for i, _ in enumerate(agent_indices): + idx = segment_indices[i] + taken = jnp.take(arr, idx, axis=0) + pad_len = segs_per_arm - taken.shape[0] + padded = jnp.pad(taken, [(9, pad_len)] + [(0, 0)] * (taken.ndim - 1)) + + per_agent.append(padded.reshape(-1)) + arr = jnp.stack(per_agent) + elif key in _JOINT_SCALED_KEYS: + per_agent = [] + for i, _ in enumerate(agent_indices): + idx = joint_indices[i] + taken = jnp.take(arr, idx, axis=0) + pad_len = joints_per_arm - taken.shape[0] + padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1)) + + per_agent.append(padded.reshape(-1)) + arr = jnp.stack(per_agent) else: - padded[key] = arr - return padded + arr = jnp.repeat(arr[None, :], num_agents, axis=0) + + if morph_mode == MorphMode.CENTRALIZED: + output[key] = arr.reshape(1, -1) + elif key in _JOINT_SCALED_KEYS: + output[key] = arr.reshape(num_agents, -1) + elif key in _SEGMENT_SCALED_KEYS: + output[key] = arr[:, None] + else: + output[key] = arr + + return output def _flatten_features(obs: dict) -> jnp.ndarray: - ordered_keys = [ - "disk_z_tilt", - "joint_actuator_force", - "joint_position", - "joint_velocity", - "robot_direction_to_target", - "segment_contact", - ] + """ + Input: + key -> (num_arms, feat_per_key) + + Output: + (num_arms, total_features) + """ values = [] + for key in ordered_keys: - if key in obs: - arr = jnp.asarray(obs[key]).flatten() - if arr.size > 0: - values.append(arr) - return jnp.concatenate(values) + if key not in obs: + continue + + arr = jnp.asarray(obs[key]) # (num_arms, feat) + + if arr.size == 0: + continue + + if arr.ndim == 1: + arr = arr[:, None] + + arr = arr.reshape(arr.shape[0], -1) + + values.append(arr) + + return jnp.concatenate(values, axis=-1) # (num_arms, total_feat) def _process_single(obs_dict: dict) -> jnp.ndarray: processed = _add_derived_features(obs_dict) processed = _normalize_features(processed) - if padding_masks is not None: - processed = _pad_features(processed) - return _flatten_features(processed) + processed = _split_to_agents(processed, morph_mode) + flat = _flatten_features(processed) # (num_arms, total_feat) + + logger.debug(f"[FLATTENED FINAL] shape: {flat.shape}") + logger.debug(f"[PER AGENT] example row 0 shape: {flat[0].shape}") + + return flat # (agents, feat) return jax.jit(jax.vmap(_process_single)) diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py index ea00713..3216dfe 100644 --- a/src/brittle_star_project/environment/padded_obs_wrapper.py +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -30,6 +30,12 @@ def compute_padding_masks( mask_2x = [] for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)): + if not isinstance(actual, int): + actual = actual.item() + + if not isinstance(ref, int): + ref = ref.item() + if not (0 <= actual <= ref): raise ValueError( f"Invalid amputation at arm {arm_idx}: " diff --git a/src/brittle_star_project/evaluation/checkpoint.py b/src/brittle_star_project/evaluation/checkpoint.py index d10046b..9b868d5 100644 --- a/src/brittle_star_project/evaluation/checkpoint.py +++ b/src/brittle_star_project/evaluation/checkpoint.py @@ -4,6 +4,8 @@ import yaml from dataclasses import dataclass from pathlib import Path +from collections.abc import Mapping + import flax from omegaconf import OmegaConf @@ -32,15 +34,19 @@ def load_params(path: Path) -> dict: sensor_params = None actor_params = None + message_passer_params = None # Extract params from restored checkpoint - if isinstance(restored, dict): + if isinstance(restored, Mapping): params_sub = restored.get("params", {}) sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params") actor_params = restored.get("actor_params") or params_sub.get("actor_params") + message_passer_params = restored.get("message_passer_params") or params_sub.get( + "message_passer_params" + ) elif isinstance(restored, (list, tuple)) and len(restored) >= 2: params_part = restored[1] - if isinstance(params_part, dict): + if isinstance(params_part, Mapping): sensor_params = params_part.get("0", params_part.get(0)) actor_params = params_part.get("1", params_part.get(1)) elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2: @@ -53,6 +59,7 @@ def load_params(path: Path) -> dict: return { "sensor_params": sensor_params, "actor_params": actor_params, + "message_passer_params": message_passer_params, } diff --git a/src/brittle_star_project/evaluation/evaluate_mjx.py b/src/brittle_star_project/evaluation/evaluate_mjx.py index 0193d99..317fa98 100644 --- a/src/brittle_star_project/evaluation/evaluate_mjx.py +++ b/src/brittle_star_project/evaluation/evaluate_mjx.py @@ -52,6 +52,7 @@ def build_eval_rollout_fn( obs_processor: Callable, sensor_apply: Callable, actor_apply: Callable, + message_passer_apply: Callable | None = None, action_low: jnp.ndarray, action_high: jnp.ndarray, reward_fn: Callable, @@ -67,6 +68,9 @@ def build_eval_rollout_fn( returned by ``create_obs_processor``. sensor_apply: The sensor network's ``apply`` method (JIT-compiled). actor_apply: The actor network's ``apply`` method (JIT-compiled). + message_passer_apply: Optional message-passing module apply method. + When provided, it is applied between the sensor and actor, using + ``params["message_passer_params"]``. action_low: Per-joint action lower bound (JAX array, shape ``(action_dim,)``). action_high: Per-joint action upper bound (JAX array, shape ``(action_dim,)``). reward_fn: Shaped reward function with signature @@ -101,10 +105,14 @@ def build_eval_rollout_fn( obs = obs_processor(state.observations) hidden = sensor_apply(params["sensor_params"], obs) + if message_passer_apply is not None: + mp_params = params["message_passer_params"] + hidden = jax.vmap(lambda x: message_passer_apply(mp_params, x))(hidden) mean, _log_std = actor_apply(params["actor_params"], hidden) # Deterministic action: use the actor mean, no exploration noise. - action = jnp.clip(mean, action_low, action_high) + flat_mean = mean.reshape(mean.shape[0], -1) + action = jnp.clip(flat_mean, action_low, action_high) next_state = step_1(state=state, action=action) shaped_reward = reward_fn(state, next_state) diff --git a/src/brittle_star_project/evaluation/policy.py b/src/brittle_star_project/evaluation/policy.py index 75b6a22..5b12d07 100644 --- a/src/brittle_star_project/evaluation/policy.py +++ b/src/brittle_star_project/evaluation/policy.py @@ -24,10 +24,17 @@ class PolicyAgent: *, sensor_params: Any, actor_params: Any, + message_passer_params: Any | None = None, + message_passing_steps: int | None = None, + adj_matrix: Any | None = None, action_dim: int, obs_processor: Any, ) -> None: - from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation + from brittle_star_project.MLPs.mlps import ( + Actor, + GenericDenseLayersWithActivation, + MessagePasser, + ) # Infer layer sizes from params try: @@ -45,7 +52,8 @@ class PolicyAgent: key = f"Dense_{idx}" if key not in dense_params: break - layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1])) + + layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[-1])) idx += 1 if not layer_sizes: @@ -53,11 +61,31 @@ class PolicyAgent: self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes) self._actor = Actor(action_dim=action_dim) - self._sensor_apply = jax.jit(self._sensor.apply) - self._actor_apply = jax.jit(self._actor.apply) + + self._message_passer = None + if message_passer_params is not None and not ( + isinstance(message_passer_params, dict) and len(message_passer_params) == 0 + ): + if message_passing_steps is None or adj_matrix is None: + raise ValueError( + "Checkpoint contains message_passer_params but PolicyAgent was not given " + "message_passing_steps and adj_matrix. Pass these when constructing the agent " + "so decentralized evaluation matches training." + ) + + hidden_dim = int(layer_sizes[-1]) + self._message_passer = MessagePasser( + hidden_dim=hidden_dim, + num_propagation_steps=int(message_passing_steps), + adj_matrix=jnp.asarray(adj_matrix), + ) + self._message_passer.apply = jax.jit(self._message_passer.apply) + self._sensor.apply = jax.jit(self._sensor.apply) + self._actor.apply = jax.jit(self._actor.apply) self._params = { "sensor_params": sensor_params, "actor_params": actor_params, + "message_passer_params": message_passer_params, } self._obs_processor = obs_processor @@ -67,6 +95,9 @@ class PolicyAgent: *, sensor_params: Any, actor_params: Any, + message_passer_params: Any | None = None, + message_passing_steps: int | None = None, + adj_matrix: Any | None = None, action_dim: int, obs_processor: Any, ) -> "PolicyAgent": @@ -74,14 +105,24 @@ class PolicyAgent: return cls( sensor_params=sensor_params, actor_params=actor_params, + message_passer_params=message_passer_params, + message_passing_steps=message_passing_steps, + adj_matrix=adj_matrix, action_dim=action_dim, obs_processor=obs_processor, ) - def set_params(self, *, sensor_params: Any, actor_params: Any) -> None: + def set_params( + self, + *, + sensor_params: Any, + actor_params: Any, + message_passer_params: Any | None = None, + ) -> None: """Update parameters for evaluation without rebuilding the model.""" self._params["sensor_params"] = sensor_params self._params["actor_params"] = actor_params + self._params["message_passer_params"] = message_passer_params @classmethod def from_checkpoint( @@ -90,6 +131,8 @@ class PolicyAgent: *, action_dim: int, obs_processor: Any, + message_passing_steps: int | None = None, + adj_matrix: Any | None = None, ) -> "PolicyAgent": """Load params from .flax and construct the agent.""" params = load_params(model_path) @@ -97,15 +140,38 @@ class PolicyAgent: return cls( sensor_params=params["sensor_params"], actor_params=params["actor_params"], + message_passer_params=params.get("message_passer_params"), + message_passing_steps=message_passing_steps, + adj_matrix=adj_matrix, action_dim=action_dim, 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)[0] - hidden = self._sensor_apply(self._params["sensor_params"], obs) - mean, _log_std = self._actor_apply(self._params["actor_params"], hidden) + obs = self._obs_processor(batched_obs) + + hidden = self._apply_per_node(self._sensor, self._params["sensor_params"], obs) + + if self._message_passer is not None: + mp_params = self._params.get("message_passer_params") + if mp_params is None or (isinstance(mp_params, dict) and len(mp_params) == 0): + raise ValueError( + "PolicyAgent has a message passer but message_passer_params are missing/empty." + ) + 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) return np.asarray(mean, dtype=np.float32).ravel() diff --git a/src/brittle_star_project/evaluation/rollout.py b/src/brittle_star_project/evaluation/rollout.py index 79c86ed..982a564 100644 --- a/src/brittle_star_project/evaluation/rollout.py +++ b/src/brittle_star_project/evaluation/rollout.py @@ -114,18 +114,20 @@ def rollout_viewer( episode_return = 0.0 observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None reached_target = _target_reached(state=state) steps = 0 with mujoco.viewer.launch_passive(model, data) as viewer: step_iter = range(int(max_steps)) if max_steps is not None else itertools.count() - for _step_idx in step_iter: + for _ in step_iter: if not viewer.is_running(): break step_start = time.time() obs_dict = observations or {} + action = policy.act(observations=obs_dict) if action_mask is not None: action = action[action_mask] diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index 1ea37a8..fb27ccd 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -1,14 +1,27 @@ from functools import partial -import flax import jax import jax.numpy as jnp +from jax import debug +from flax.core import FrozenDict +from experiment_logger import get_logger +from brittle_star_project.utils import logged_jit + +logger = get_logger() # Chose to use a class as it seemed the easiest way to integrate the CleanRL code style # with our need to seperate concerns class PPO: - def __init__(self, args, sensor, actor, critic, feature_extractor, message_passer=None): + def __init__( + self, + args, + sensor_apply, + actor_apply, + critic_apply, + feature_extractor_apply, + message_passer=None, + ): self.args = args if not message_passer: @@ -18,10 +31,10 @@ class PPO: partial( ppo_loss, args=args, - sensor_apply=sensor.apply, - actor_apply=actor.apply, - critic_apply=critic.apply, - feature_extractor_apply=feature_extractor.apply, + sensor_apply=sensor_apply, + actor_apply=actor_apply, + critic_apply=critic_apply, + feature_extractor_apply=feature_extractor_apply, message_passer=message_passer, ), has_aux=True, @@ -29,8 +42,14 @@ class PPO: # This PPO class should be initialized only once, # or this function will need to recompile - @partial(jax.jit, static_argnums=0) + @partial(logged_jit, static_argnums=0) def update_ppo(self, agent_state, storage, key): + debug.callback(logger.debug, f"[PPO] storage.obs shape: {storage.obs.shape}") + debug.callback(logger.debug, f"[PPO] storage.actions shape: {storage.actions.shape}") + debug.callback(logger.debug, f"[PPO] storage.logprobs shape: {storage.logprobs.shape}") + debug.callback(logger.debug, f"[PPO] storage.advantages shape: {storage.advantages.shape}") + debug.callback(logger.debug, f"[PPO] storage.returns shape: {storage.returns.shape}") + args = self.args ppo_loss_grad_fn = self.ppo_loss_grad_fn @@ -49,6 +68,16 @@ class PPO: shuffled_storage = jax.tree.map(convert_data, flatten_storage) def update_minibatch(agent_state, minibatch): + debug.callback(logger.debug, f"[PPO] minibatch.obs: {minibatch.obs.shape}") + debug.callback(logger.debug, f"[PPO] minibatch.actions: {minibatch.actions.shape}") + debug.callback( + logger.debug, f"[PPO] minibatch.logprobs: {minibatch.logprobs.shape}" + ) + debug.callback( + logger.debug, f"[PPO] minibatch.advantages: {minibatch.advantages.shape}" + ) + debug.callback(logger.debug, f"[PPO] minibatch.returns: {minibatch.returns.shape}") + (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn( agent_state.params, minibatch.obs, @@ -58,19 +87,12 @@ class PPO: minibatch.returns, ) agent_state = agent_state.apply_gradients(grads=grads) - return agent_state, ( - loss, - pg_loss, - v_loss, - entropy_loss, - approx_kl, - grads, - ) + return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl) agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage) return (agent_state, key), metrics - (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan( + (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl) = jax.lax.scan( update_epoch, (agent_state, key), (), length=args.update_epochs ) return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key @@ -84,27 +106,45 @@ that are now not in the same scope """ -@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4)) +@partial(logged_jit, static_argnums=(0, 1, 2, 3, 4)) def get_action_and_value( sensor_apply, actor_apply, message_passer, critic_apply, feature_extractor_apply, - params: flax.core.FrozenDict, + params: FrozenDict, x: jnp.ndarray, action: jnp.ndarray, ): hidden_sensor = sensor_apply(params["sensor_params"], x) hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x) - hidden_sensor = message_passer(hidden_sensor) + + # only apply message passing in decentralized context + if message_passer is not None: + hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor) + + debug.callback(logger.debug, f"[SHAPE] hidden_sensor: {hidden_sensor.shape}") + debug.callback(logger.debug, f"[SHAPE] hidden_critic: {hidden_critic.shape}") + mean, log_std = actor_apply(params["actor_params"], hidden_sensor) + + debug.callback(logger.debug, f"[SHAPE] mean: {mean.shape}") + debug.callback(logger.debug, f"[SHAPE] log_std: {log_std.shape}") + debug.callback(logger.debug, f"[SHAPE] action: {action.shape}") + log_std = jnp.clip(log_std, -5, 2) std = jnp.exp(log_std) - logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) - entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1) + logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)) + debug.callback(logger.debug, f"[SHAPE] logprob pre-sum: {logprob.shape}") + + logprob = logprob.sum(axis=(-2, -1)) + debug.callback(logger.debug, f"[SHAPE] logprob final: {logprob.shape}") + + entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(axis=(-2, -1)) value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1) + debug.callback(logger.debug, f"[SHAPE] value: {value.shape}") return logprob, entropy, value @@ -150,7 +190,7 @@ def ppo_loss( return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl)) -def identity(hidden): +def identity(_, hidden): """ Used for seamless jax integration, avoids having branching inside jitted function, diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 7e7f7b3..76bb128 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -3,13 +3,14 @@ import random import time from dataclasses import asdict, dataclass from functools import partial +from typing import Any, Optional import jax import jax.numpy as jnp import numpy as np import optax +import flax.linen as nn from flax.training.train_state import TrainState -from typing import Any from experiment_logger import get_logger @@ -26,24 +27,21 @@ from brittle_star_project.MLPs.mlps import ( Actor, AgentParams, GenericDenseLayersWithActivation, + MessagePasser, OneDenseLayerMLP, Storage, ) +from brittle_star_project.MLPs.adjancency_builder import build_adjacency from brittle_star_project.ppo import PPO +from brittle_star_project.environment import MorphMode +from brittle_star_project.utils import logged_jit from brittle_star_project.environment.env_types import Backend # 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 +@logged_jit def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: return jnp.clip(action, low, high) @@ -54,39 +52,54 @@ def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> fl return float(explained_var) -@jax.jit +@logged_jit def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate): frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations return learning_rate * frac def _get_action_and_value_noise( - sensor: GenericDenseLayersWithActivation, - feature_extractor: GenericDenseLayersWithActivation, - actor: Actor, - critic: OneDenseLayerMLP, + sensor: nn.Module, + feature_extractor: nn.Module, + actor: nn.Module, + critic: nn.Module, + message_passer: Optional[nn.Module], agent_state: TrainState, next_obs: jnp.ndarray, - key: jax.random.PRNGKey, + key, action_low, action_high, ): - hidden = sensor.apply(agent_state.params["sensor_params"], next_obs) - hidden_critic = feature_extractor.apply( - agent_state.params["feature_extractor_params"], next_obs + # (B, n_nodes, feat) + hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs) + + if message_passer is not None: + params = agent_state.params["message_passer_params"] + # (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension + hidden = jax.vmap(lambda x: message_passer.apply(params, x))(hidden) + + hidden_critic = apply_shared( + feature_extractor, agent_state.params["feature_extractor_params"], next_obs ) - mean, log_std = actor.apply(agent_state.params["actor_params"], hidden) + mean, log_std = apply_per_node(actor, 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) std = jnp.exp(log_std) - raw_action = mean + noise * std - clipped_action = _clip_action(raw_action, action_low, action_high) - logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) - value = critic.apply(agent_state.params["critic_params"], hidden_critic) - return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key + raw_action = mean + noise * std + flat_action = raw_action.reshape( + raw_action.shape[0], -1 + ) # concat the per agent, keep the envs dim (batch, agent * action) + flat_clipped_action = _clip_action(flat_action, action_low, action_high) + + logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum( + axis=(-2, -1) + ) + value = apply_shared(critic, agent_state.params["critic_params"], hidden_critic) + + return flat_clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key def _step_once( @@ -94,31 +107,59 @@ def _step_once( _, env_step_fn, num_envs: int, - sensor: GenericDenseLayersWithActivation, - feature_extractor: GenericDenseLayersWithActivation, - actor: Actor, - critic: OneDenseLayerMLP, + sensor: nn.Module, + feature_extractor: nn.Module, + actor: nn.Module, + critic: nn.Module, + message_passer: Optional[nn.Module], action_low, action_high, ): agent_state, episode_stats, obs, done, key, env_state, terminated_any, truncated_any = carry - clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise( - sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high + flat_clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise( + sensor, + feature_extractor, + actor, + critic, + message_passer, + agent_state, + obs, + key, + action_low, + action_high, ) + logger = get_logger() + logger.debug(f"[_step_once] raw_action: {raw_action.shape}") + logger.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}") + + # Supporting signals (often where mismatch originates) + logger.debug(f"[_step_once] logprob: {logprob.shape}") + logger.debug(f"[_step_once] value: {value.shape}") + logger.debug(f"[_step_once] mean: {mean.shape}") + logger.debug(f"[_step_once] std: {std.shape}") + + key, reset_key = jax.random.split(key) + reset_rngs = jax.random.split(reset_key, num_envs) + + # ---- ENV STEP ---- key, reset_key = jax.random.split(key) reset_rngs = jax.random.split(reset_key, num_envs) episode_stats, env_state, (next_obs, reward, next_done, terminated, truncated) = env_step_fn( episode_stats, env_state, - clipped_action, + flat_clipped_action, reset_rngs, ) terminated_any = terminated_any | terminated truncated_any = truncated_any | truncated + logger.debug(f"[_step_once] next_obs: {next_obs.shape}") + logger.debug(f"[_step_once] reward: {reward.shape}") + logger.debug(f"[_step_once] next_done: {next_done.shape}") + storage = Storage( obs=obs, actions=raw_action, @@ -231,6 +272,25 @@ 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: + batch_size = x.shape[0] + x_flattened = x.reshape(batch_size, -1) + return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened) + + def _rollout_jit( agent_state, episode_stats, @@ -241,10 +301,11 @@ def _rollout_jit( max_steps, step_env_fn, num_envs: int, - sensor: GenericDenseLayersWithActivation, - feature_extractor: GenericDenseLayersWithActivation, - actor: Actor, - critic: OneDenseLayerMLP, + sensor: nn.Module, + feature_extractor: nn.Module, + actor: nn.Module, + critic: nn.Module, + message_passer: Optional[nn.Module], action_low, action_high, ): @@ -270,6 +331,7 @@ def _rollout_jit( feature_extractor=feature_extractor, actor=actor, critic=critic, + message_passer=message_passer, env_step_fn=step_env_fn, num_envs=num_envs, action_low=action_low, @@ -321,9 +383,10 @@ def _compute_gae_jit( feature_extractor, critic, ): - next_value = critic.apply( + next_value = apply_shared( + critic, agent_state.params["critic_params"], - feature_extractor.apply(agent_state.params["feature_extractor_params"], next_obs), + apply_shared(feature_extractor, agent_state.params["feature_extractor_params"], next_obs), ).squeeze(-1) advantages = jnp.zeros((num_envs,)) @@ -357,7 +420,11 @@ class TrainingMeasurements: class PPOTrainer: def __init__( - self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str + self, + cfg: BrittleStarConfig, + env: BrittleStarJaxEnvWrapper, + run_dir: str, + run_name: str, ): self.cfg = cfg self.ppo = cfg.ppo @@ -375,24 +442,49 @@ class PPOTrainer: self.key = jax.random.PRNGKey(self.experiment.seed) + self.morph_mode = self.cfg.morphology.morph_mode + + self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32) + self.num_segments = self.segments_per_arm.sum().item() + self.num_arms = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item() + + self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}") + self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode) + + ( + self.sensor, + self.message_passer, + self.actor, + self.feature_extractor, + self.critic, + self.needed_copies, + self.agent_indices, + ) = self._init_agent() + + self.sensor.apply = logged_jit(self.sensor.apply) + self.feature_extractor.apply = logged_jit(self.feature_extractor.apply) + self.actor.apply = logged_jit(self.actor.apply) + self.critic.apply = logged_jit(self.critic.apply) + # Build the centralized observation processor: derive -> normalize -> pad -> flatten. self.obs_processor = create_obs_processor( bounds_dict=self.cfg.obs_bounds.to_bounds_dict(), + needed_copies=self.needed_copies, + num_arms=self.num_arms, + morph_mode=self.morph_mode, padding_masks=self.env.padding_masks, + segments_per_arm=self.segments_per_arm, + agent_indices=self.agent_indices, ) - self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent() - self.sensor.apply = jax.jit(self.sensor.apply) - self.feature_extractor.apply = jax.jit(self.feature_extractor.apply) - self.actor.apply = jax.jit(self.actor.apply) - self.critic.apply = jax.jit(self.critic.apply) + self.logger.debug(f"needed copies = {self.needed_copies}") action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32) self._action_low = action_low self._action_high = action_high - self._rollout_jit = jax.jit( + self._rollout_jit = logged_jit( partial( _rollout_jit, max_steps=self.ppo.num_steps, @@ -407,11 +499,12 @@ class PPOTrainer: feature_extractor=self.feature_extractor, actor=self.actor, critic=self.critic, + message_passer=self.message_passer, action_low=action_low, action_high=action_high, ) ) - self._compute_gae_jit = jax.jit( + self._compute_gae_jit = logged_jit( partial( _compute_gae_jit, num_envs=self.ppo.num_envs, @@ -422,7 +515,30 @@ class PPOTrainer: ) ) - self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor) + def apply_sensor(p, x): + return apply_per_node(self.sensor, p, x) + + def apply_actor(p, x): + return apply_per_node(self.actor, p, x) + + def apply_critic(p, x): + return apply_shared(self.critic, p, x) + + def apply_feature(p, x): + return apply_shared(self.feature_extractor, p, x) + + def apply_message_passer(p, x): + assert self.message_passer is not None + return jax.vmap(lambda x_in: self.message_passer.apply(p, x_in))(x) + + self._ppo = PPO( + self.ppo, + apply_sensor, + apply_actor, + apply_critic, + apply_feature, + apply_message_passer if self.message_passer is not None else None, + ) self.agent_state = self._init_agent_state() @@ -440,33 +556,136 @@ 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 = ( + self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum() + ).item() + # scale actor output with size of model --> more models ==> less actions needed per model + actor = Actor(action_dim=self.env.single_action_space.shape[0] // needed_copies) sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) + message_passer: Optional[nn.Module] = ( + MessagePasser( + hidden_dim=300, + num_propagation_steps=self.cfg.architecture.message_passing_steps or 4, + adj_matrix=self.adj, + ) + if self.morph_mode != MorphMode.CENTRALIZED + else None + ) + feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) - actor = Actor(action_dim=self.env.single_action_space.shape[0]) critic = OneDenseLayerMLP() - return sensor, feature_extractor, actor, critic + 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...") - self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split( - self.key, 5 + self.key, sensor_key, actor_key, critic_key, feature_extractor_key, message_passer_key = ( + jax.random.split(self.key, 6) ) dummy_reset = self.env.reset(seed=0) + + for k, v in dummy_reset.observations.items(): + self.logger.debug(k, v.shape) + sample_obs = self.obs_processor(dummy_reset.observations)[0] # take first env - sensor_params = self.sensor.init(sensor_key, sample_obs) - feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs) - actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs)) - critic_params = self.critic.init( - critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs) + + self.logger.debug(f"[_init_agent_state] sample_obs: {sample_obs.shape}") + self.obs_mean = jnp.zeros((sample_obs.shape[-1],)) + self.obs_var = jnp.ones((sample_obs.shape[-1],)) + self.obs_count = 1e-4 + self.logger.debug(f"[_init_agent_state] obs_mean: {self.obs_mean.shape}") + self.logger.debug(f"[_init_agent_state] obs_var: {self.obs_var.shape}") + + self.logger.debug(f"[_init_agent_state]: Needed copies: {self.needed_copies}") + sensor_keys = jax.random.split(sensor_key, self.needed_copies) + actor_keys = jax.random.split(actor_key, self.needed_copies) + + # (needed_copies, X) + sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys) + self.logger.debug( + f"[_init_agent_state] sensor_params: {jax.tree.map(lambda x: x.shape, sensor_params)}" + ) + + single_sensor_param = jax.tree.map(lambda x: x[0], sensor_params) + self.logger.debug( + f"[_init_agent_state] single_sensor_param: { + jax.tree.map(lambda x: x.shape, single_sensor_param) + }" + ) + + sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs) + self.logger.debug( + f"[_init_agent_state] sensor_params_sample shape: {sensor_params_sample.shape}" + ) + + actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys) + self.logger.debug( + f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}" + ) + + message_passer_params = {} + if self.morph_mode != MorphMode.CENTRALIZED: + assert self.message_passer is not None, "decentralized modes require a message passer" + + message_passer_params = self.message_passer.init( + message_passer_key, + self.sensor.apply(single_sensor_param, sample_obs), + ) + self.logger.debug( + f"[_init_agent_state] message_passer_params: { + jax.tree.map(lambda x: x.shape, message_passer_params) + }" + ) + + flat_obs = sample_obs.reshape(-1) # BECAUSE 1 centralized critic + self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}") + + feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs) + self.logger.debug( + f"[_init_agent_state] feature_extractor_params: { + jax.tree.map(lambda x: x.shape, feature_extractor_params) + }" + ) + + critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs) + self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}") + + critic_params = self.critic.init(critic_key, critic_input) + self.logger.debug( + f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}" ) return TrainState.create( apply_fn=None, params=asdict( - AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params) + AgentParams( + sensor_params, + actor_params, + critic_params, + feature_extractor_params, + message_passer_params, + ) ), tx=optax.chain( optax.clip_by_global_norm(self.ppo.max_grad_norm), @@ -569,7 +788,7 @@ class PPOTrainer: def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple: if iteration == 1: self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}") - + self.logger.debug(f"[_step] next_obs (in): {next_obs.shape}") ( self.agent_state, self.episode_stats, @@ -581,12 +800,12 @@ class PPOTrainer: terminated_any, truncated_any, ) = self._rollout(env_state, next_obs, next_done) - + self.logger.debug(f"[_step] next_obs (post-rollout): {next_obs.shape}") if iteration == 1: self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}") storage = self._compute_gae(storage, next_obs, next_done) - + self.logger.debug(f"[_step] storage.obs (post-gae): {storage.obs.shape}") if iteration == 1: self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}") @@ -684,8 +903,11 @@ class PPOTrainer: self._eval_fn = build_eval_rollout_fn( env=self.env, obs_processor=self.obs_processor, - sensor_apply=self.sensor.apply, - actor_apply=self.actor.apply, + sensor_apply=lambda p, x: apply_per_node(self.sensor, p, x), + actor_apply=lambda p, x: apply_per_node(self.actor, p, x), + message_passer_apply=( + None if self.message_passer is None else self.message_passer.apply + ), action_low=self._action_low, action_high=self._action_high, reward_fn=reward_fn, @@ -718,7 +940,10 @@ class PPOTrainer: self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}") env_state = self.env.reset(seed=self.experiment.seed) + next_obs = self.obs_processor(env_state.observations) + self.logger.debug(f"[train] next_obs: {next_obs.shape}") + next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_) self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}") diff --git a/src/brittle_star_project/utils/__init__.py b/src/brittle_star_project/utils/__init__.py new file mode 100644 index 0000000..ed72e5a --- /dev/null +++ b/src/brittle_star_project/utils/__init__.py @@ -0,0 +1,3 @@ +from .logged_jit import logged_jit + +__all__ = ["logged_jit"] diff --git a/src/brittle_star_project/utils/logged_jit.py b/src/brittle_star_project/utils/logged_jit.py new file mode 100644 index 0000000..3d29ba6 --- /dev/null +++ b/src/brittle_star_project/utils/logged_jit.py @@ -0,0 +1,17 @@ +import jax +from experiment_logger import get_logger + + +def logged_jit(fn, **jit_kwargs): + logger = get_logger() + name = getattr(fn, "__name__", getattr(fn, "__qualname__", repr(fn))) + + def decorator(func): + def traced_func(*args, **kwargs): + logger.debug(f"[JIT] Compiling {name}...") + return func(*args, **kwargs) + + jitted = jax.jit(traced_func, **jit_kwargs) + return jitted + + return decorator(fn) diff --git a/tests/test_adjacency.py b/tests/test_adjacency.py new file mode 100644 index 0000000..27c8b43 --- /dev/null +++ b/tests/test_adjacency.py @@ -0,0 +1,98 @@ +import jax.numpy as jnp +import numpy as np + +from brittle_star_project.MLPs import build_adjacency +from brittle_star_project.environment.env_config import MorphMode + + +def assert_symmetric(adj): + assert jnp.all(adj == adj.T) + + +def test_centralized(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.CENTRALIZED) + + assert adj.shape == (1, 1) + assert adj[0, 0] == 1 + + +def test_fully_connected(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.FULLY_CONNECTED) + + assert adj.shape == (5, 5) + assert jnp.all(adj == 1) + + +def test_ring(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.RING) + + assert adj.shape == (5, 5) + assert_symmetric(adj) + + # each node should connect to itself + 2 neighbors + for node in range(5): + assert adj[node, node] == 1 + assert jnp.sum(adj[node]) == 3 + neighbor1 = (node - 1) % 5 + neighbor2 = (node + 1) % 5 + assert adj[neighbor1, node] == 1 + assert adj[node, neighbor2] == 1 # Symmetrical + + +def test_segment_structure(): + segments = [4, 4, 4, 4, 4] + adj = build_adjacency(segments, MorphMode.SEGMENT) + + num_arms = 5 + num_segments = sum(segments) + num_nodes = num_arms + num_segments + + assert adj.shape == (num_nodes, num_nodes) + + # --- ring connectivity --- + for i in range(num_arms): + assert adj[i, i] == 1 + assert adj[i, (i - 1) % num_arms] == 1 + assert adj[i, (i + 1) % num_arms] == 1 + + # --- segment chain checks --- + offset = num_arms + for arm in range(5): + for i in range(4): + node = offset + arm * 4 + i + + # self + assert adj[node, node] == 1 + + # chain neighbors + if i > 0: + assert adj[node, node - 1] == 1 + if i < 3: + assert adj[node, node + 1] == 1 + + # --- ring ↔ segment connections --- + for arm in range(5): + first_seg = num_arms + arm * 4 + assert adj[arm, first_seg] == 1 + assert adj[first_seg, arm] == 1 + + save_adj(adj) + + +def save_adj(adj, name="adjacency_debug.txt"): + a = np.array(adj) + + with open(name, "w") as f: + f.write("\nAdjacency matrix:\n") + f.write(" " + " ".join([f"{i:2d}" for i in range(a.shape[0])]) + "\n") + + for i, row in enumerate(a): + line = f"{i:2d} " + " ".join(["█" if x > 0 else "." for x in row]) + f.write(line + "\n") + + +def test_no_isolated_nodes(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.SEGMENT) + + # no node should be completely isolated + assert jnp.all(jnp.sum(adj, axis=0) > 0) diff --git a/tests/test_network_shapes.py b/tests/test_network_shapes.py index e1a7a1a..7bec744 100644 --- a/tests/test_network_shapes.py +++ b/tests/test_network_shapes.py @@ -1,5 +1,6 @@ import jax import jax.numpy as jnp +from brittle_star_project.environment.env_config import MorphMode from brittle_star_project.environment.padded_obs_wrapper import ( compute_padding_masks, ) @@ -20,14 +21,24 @@ def test_centralized_forward_pass_with_padding(): "segment_contact": jnp.zeros((batch_size, 14)), } + segments_per_arm = jnp.array((4, 0, 4, 2, 4)) + num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item() + # 2. Process and Pad Observation - masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4)) - obs_processor = create_obs_processor(bounds_dict={}, padding_masks=masks) + masks = compute_padding_masks(segments_per_arm=list(segments_per_arm)) + obs_processor = create_obs_processor( + bounds_dict={}, + needed_copies=1, + num_arms=num_arms, + padding_masks=masks, + morph_mode=MorphMode.CENTRALIZED, + segments_per_arm=segments_per_arm, + ) global_state = obs_processor(amputated_obs) - # 40 + 40 + 20 = 100 dimensions - assert global_state.shape == (batch_size, 100), ( - f"Expected global state shape (2, 100), got {global_state.shape}" + # 40 + 40 + 20 + padding = 145 dimensions + assert global_state.shape == (batch_size, 1, 145), ( + f"Expected global state shape (2, 1, 145), got {global_state.shape}" ) # 4. Initialize dummy networks (40 actuators for the max morphology output) @@ -45,9 +56,11 @@ def test_centralized_forward_pass_with_padding(): action_mean, action_log_std = actor.apply(actor_params, global_state) value = critic.apply(critic_params, global_state) - assert action_mean.shape == (batch_size, 40), f"Actor mean shape mismatch: {action_mean.shape}" + assert action_mean.shape == (batch_size, 1, 40), ( + f"Actor mean shape mismatch: {action_mean.shape}" + ) assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}" - assert value.shape == (batch_size, 1) or value.shape == (batch_size,), ( + assert value.shape == (batch_size, 1, 1) or value.shape == (batch_size,), ( f"Critic value shape mismatch: {value.shape}" ) diff --git a/tests/test_obs_processor.py b/tests/test_obs_processor.py new file mode 100644 index 0000000..17c5c8f --- /dev/null +++ b/tests/test_obs_processor.py @@ -0,0 +1,119 @@ +import jax +import jax.numpy as jnp + +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.environment.env_config import MorphMode, ObservationBoundsConfig + + +obs_bounds = ObservationBoundsConfig().to_bounds_dict() + +""" +Test for obs_processor. + +Centralized: 40 features per agent: + disk_z_tilt → scalar → reshaped to (1,) → 1 feat + joint_actuator_force → 8 joints padded to 8 → 8 feat + joint_position → 8 joints padded to 8 → 8 feat + joint_velocity → 8 joints padded to 8 → 8 feat + robot_direction_to_target→ (x, y) → 2 feat + segment_contact → 4 segs, pre-padded by 9 → 13 feat (9 leading + 4) +""" + +NUM_ARMS = 5 +SEGS_PER_ARM = 4 # healthy segments per arm +JOINTS_PER_SEG = 2 # from _build_joint_indices: segs * 2 + +SEGS_HEALTHY = [4, 4, 4, 4, 4] +SEGS_DAMAGED = [4, 4, 4, 4, 0] # arm 4 fully disabled +SEGS_DAMAGED_2 = [4, 0, 4, 2, 4] # arm 3 fully disabled +AGENT_INDICES = [0, 1, 2, 3, 4] + +FEAT_PER_AGENT = 1 + 8 + 8 + 8 + 2 + 13 # = 40 + + +def make_obs(segs_per_arm: list[int]) -> dict: + total_segs = sum(segs_per_arm) + total_joints = JOINTS_PER_SEG * total_segs + + return { + "actuator_force": jnp.ones(total_joints), + "disk_angular_velocity": jnp.zeros(3), + "disk_linear_velocity": jnp.zeros(3), + "disk_position": jnp.zeros(3), + "disk_rotation": jnp.array([0.1, 0.1, 0.5]), # (roll, pitch, yaw) + "joint_actuator_force": jnp.full(total_joints, 1.0), + "joint_position": jnp.full(total_joints, 0.5), + "joint_velocity": jnp.full(total_joints, 2.0), + "segment_contact": jnp.ones(total_segs), + "tendon_position": jnp.zeros(0), + "tendon_velocity": jnp.zeros(0), + "unit_xy_direction_to_target": jnp.array([1.0, 0.0]), + "xy_distance_to_target": jnp.array([3.5]), + } + + +def batch_obs(obs: dict): + return jax.tree_util.tree_map(lambda x: x[None, :], obs) + + +def make_processor(morph_mode: MorphMode, needed_copies: int, segments_per_arm: list[int]): + return create_obs_processor( + bounds_dict=obs_bounds, + num_arms=NUM_ARMS, + needed_copies=needed_copies, + morph_mode=morph_mode, + segments_per_arm=segments_per_arm, + agent_indices=AGENT_INDICES, + ) + + +def test_centralized_no_damage(): + proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_HEALTHY) + obs = make_obs(SEGS_HEALTHY) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, 1, 188) + + # TODO: more? + + +def test_centralized_damaged_1_arm(): + proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED) + obs = make_obs(SEGS_DAMAGED) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, 1, 188) + + +def test_centralized_damaged_2_arms(): + proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED_2) + obs = make_obs(SEGS_DAMAGED_2) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, 1, 188) + + +def test_decentralized_fully_connected_no_damage(): + proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_HEALTHY) + obs = make_obs(SEGS_HEALTHY) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT) + + +def test_decentralized_fully_connected_damaged_1_arm(): + proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_DAMAGED) + obs = make_obs(SEGS_DAMAGED) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT) diff --git a/tests/test_target_direction.py b/tests/test_target_direction.py index a81cdf1..8317fb2 100644 --- a/tests/test_target_direction.py +++ b/tests/test_target_direction.py @@ -1,6 +1,7 @@ import jax.numpy as jnp from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.environment.env_config import MorphMode from brittle_star_project.environment.env_types import Backend from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper from brittle_star_project.environment.obs_processing import create_obs_processor @@ -43,8 +44,16 @@ def test_processor_converts_to_egocentric_direction(): cfg = BrittleStarConfig() env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX) + segments_per_arm = jnp.array((4, 4, 4, 4, 4)) + num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item() + obs_processor = create_obs_processor( - bounds_dict=cfg.obs_bounds.to_bounds_dict(), padding_masks=env.padding_masks + bounds_dict=cfg.obs_bounds.to_bounds_dict(), + needed_copies=1, + 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) @@ -66,10 +75,13 @@ def test_processor_converts_to_egocentric_direction(): processed_2 = obs_processor(dummy_obs_2) # Find the indices of the elements that changed - diff_array = jnp.abs(processed_1[0] - processed_2[0]) + diff_array = jnp.abs(processed_1[0, 0] - processed_2[0, 0]) changed_indices = jnp.where(diff_array > 1e-4)[0] - local_target = processed_1[0, changed_indices] + # (143,) + local_target = processed_1[0, 0, changed_indices] + + # (2,) expected_local_target = jnp.array([0.0, -1.0]) assert jnp.sum(jnp.abs(local_target - expected_local_target)) < 1e-4, (