From 07ee32fd3baa794b55ece634f43ddec1f84583f9 Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Thu, 7 May 2026 19:53:49 +0200 Subject: [PATCH] fix(simulate): actor action dimension updated to new way of calculating (as in training) --- .../used_variables.md | 107 ------------------ scripts/simulate.py | 30 ++++- .../environment/obs_processing.py | 63 ++++------- src/brittle_star_project/evaluation/policy.py | 27 ++++- .../evaluation/rollout.py | 4 +- 5 files changed, 71 insertions(+), 160 deletions(-) delete mode 100644 experiments/debug-experiment-10042026/used_variables.md 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 11b4a90..5564aa8 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 @@ -78,10 +80,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 @@ -105,7 +131,7 @@ 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 policy = PolicyAgent.from_checkpoint( diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index 216b008..550624d 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -6,6 +6,8 @@ from brittle_star_project.environment.env_config import MorphMode from experiment_logger import get_logger +logger = get_logger() + _JOINT_SCALED_KEYS = frozenset( { "joint_position", @@ -55,8 +57,6 @@ def create_obs_processor( segments_per_arm=[4, 4, 4, 4, 4], agent_indices=[0, 1, 2, 3, 4], ): - logger = get_logger() - # made a set to allow O(1) search ordered_keys = frozenset( [ @@ -87,13 +87,6 @@ def create_obs_processor( return new_obs - def _prune_features(obs: dict) -> dict: - pruned = {} - for key, arr in obs.items(): - if key in ordered_keys: - pruned[key] = arr - return pruned - def _normalize_features(obs: dict) -> dict: normalized = {} for key, arr in obs.items(): @@ -108,35 +101,11 @@ def create_obs_processor( normalized[key] = arr return normalized - def _pad_features(obs: dict, agent_count: int) -> dict: - assert padding_masks is not None - - padded = {} - - for key, arr in obs.items(): - if key in _JOINT_SCALED_KEYS: - 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: - 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, morph_mode) -> dict: - key_to_agents = {} + output = {} num_agents = needed_copies # IMPORTANT: number of MLPs - for key, arr in obs.items(): - # TODO Should this still be here? - if arr.size == 0: + if key not in ordered_keys or arr.size == 0: continue logger.debug(f"[INPUT] {key}: {arr.shape}") @@ -144,17 +113,22 @@ def create_obs_processor( if arr.ndim == 0: arr = arr.reshape(1) + # -------- CENTRALIZED -------- if morph_mode == MorphMode.CENTRALIZED: - key_to_agents[key] = arr.reshape(1, -1) + output[key] = arr.reshape(1, -1) + # TODO: padding for centralized continue + # -------- SEGMENTS -------- if key in _SEGMENT_SCALED_KEYS: per_agent = [] - for agent_id in agent_indices: - taken = jnp.take(arr, agent_id, axis=0) # (segs, ...) + for i, agent_id in enumerate(agent_indices): + idx = segment_indices[i] + taken = jnp.take(arr, idx, axis=0) # (segs, ...) logger.debug(f"WHY {taken.shape}") - # pad to 4 + + # pad to 4 (segments per arm?) pad_len = 4 - taken.shape[0] padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1)) @@ -162,11 +136,13 @@ def create_obs_processor( out = jnp.stack(per_agent) + # -------- JOINTS -------- elif key in _JOINT_SCALED_KEYS: per_agent = [] - for agent_id in agent_indices: - taken = jnp.take(arr, agent_id, axis=0) # (joint_n, ...) + for i, _ 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] @@ -180,9 +156,9 @@ def create_obs_processor( out = jnp.repeat(arr[None, :], num_agents, axis=0) logger.debug(f"[OUTPUT] {key}: {out.shape}") - key_to_agents[key] = out + output[key] = out - return key_to_agents + return output def _flatten_features(obs: dict) -> jnp.ndarray: """ @@ -214,7 +190,6 @@ def create_obs_processor( def _process_single(obs_dict: dict) -> jnp.ndarray: processed = _add_derived_features(obs_dict) - processed = _prune_features(processed) processed = _normalize_features(processed) processed = _split_to_agents(processed, morph_mode) flat = _flatten_features(processed) # (num_arms, total_feat) diff --git a/src/brittle_star_project/evaluation/policy.py b/src/brittle_star_project/evaluation/policy.py index 5aa2d6e..b26903e 100644 --- a/src/brittle_star_project/evaluation/policy.py +++ b/src/brittle_star_project/evaluation/policy.py @@ -45,7 +45,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,8 +54,8 @@ 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._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, @@ -79,11 +80,25 @@ 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)[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) + + # TODO: message passing + hidden = self._apply_per_node(self._sensor, self._params["sensor_params"], obs) + + # hidden = jax.vmap(...) + 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]