tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

other: backup, finding bug in conversion from dict to array

Browse source
This commit is contained in:
Robin Meersman 2026-04-30 15:58:46 +02:00
parent e85fc09ed6
commit f8eb43004f
1 changed files with 13 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

26
src/brittle_star_project/trainers/PPOTrainer.py
View file

@ -130,12 +130,12 @@ def _normalize_obs(obs, mean, var, eps=1e-8):
return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0)
@jax.jit def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, segments_per_arm: jnp.ndarray):
def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, segments_per_arm): num_segments = segments_per_arm.sum()
num_segments = sum(segments_per_arm) num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum()
num_arms = sum(1 for s in segments_per_arm if s > 0)
def _filter_and_flatten(o): @jax.jit
def _filter_and_flatten(o) -> jnp.ndarray:
values = [] values = []
for key in sorted(o.keys()): for key in sorted(o.keys()):
@ -147,13 +147,13 @@ def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, segments_per_arm
continue continue
# -------- CENTRALIZED -------- # -------- CENTRALIZED --------
if morph_mode == 0: if morph_mode == MorphMode.CENTRALIZED:
values.append(v.reshape(v.shape[0], -1)) values.append(v.reshape(v.shape[0], -1))
continue continue
# -------- SPLIT TO SEGMENTS -------- # -------- SPLIT TO SEGMENTS --------
if key in _JOINT_SCALED_KEYS: if key in _JOINT_SCALED_KEYS:
if morph_mode == 3: if morph_mode == MorphMode.SEGMENT:
# special case: segment lvl and scale with joints, # special case: segment lvl and scale with joints,
# needs to split logic for center mlps and arms # needs to split logic for center mlps and arms
# center: # center:
@ -185,7 +185,7 @@ def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, segments_per_arm
continue continue
# -------- SEGMENT MODE -------- # -------- SEGMENT MODE --------
if morph_mode == 3: if morph_mode == MorphMode.SEGMENT:
values.append(v) values.append(v)
continue continue
@ -193,7 +193,7 @@ def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, segments_per_arm
v = v.reshape(v.shape[0], num_arms, -1) v = v.reshape(v.shape[0], num_arms, -1)
values.append(v) values.append(v)
return jnp.concatenate(values, axis=-1) return jnp.concatenate(values, axis=0)
return jax.vmap(_filter_and_flatten)(obs_dict) return jax.vmap(_filter_and_flatten)(obs_dict)
@ -247,7 +247,7 @@ def _get_action_and_value_noise(
return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
# TODO: update to work with extra dimension + message passing # TODO: update to work vectorized (sensor, actor, message passer) + message passing
def _step_once( def _step_once(
carry, carry,
_, _,
@ -349,7 +349,7 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn, morph_mode,
def apply_per_node(net, params, x): def apply_per_node(net, params, x):
# x: (batch, nodes, feat) # x: (batch, nodes, feat)
return jax.vmap(lambda node_x: net.apply(params, node_x), in_axes=1, out_axes=1)(x) return jax.vmap(net.apply, in_axes=(0, 1), out_axes=1)(params, x)
# TODO: update to work with extra dimension + message passing # TODO: update to work with extra dimension + message passing
@ -359,7 +359,6 @@ def _rollout_jit(
env_state, env_state,
next_obs, next_obs,
next_done, next_done,
adj_matrix,
key, key,
max_steps, max_steps,
step_env_fn, step_env_fn,
@ -370,6 +369,7 @@ def _rollout_jit(
message_passer: nn.Module, message_passer: nn.Module,
action_low, action_low,
action_high, action_high,
adj_matrix,
): ):
(agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan( (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
partial( partial(
@ -470,6 +470,7 @@ class PPOTrainer:
self.key = jax.random.PRNGKey(self.experiment.seed) self.key = jax.random.PRNGKey(self.experiment.seed)
self.morph_mode = self.cfg.morphology.morph_mode self.morph_mode = self.cfg.morphology.morph_mode
self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32)
self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}") self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}")
self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode) self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode)
@ -487,7 +488,6 @@ class PPOTrainer:
self.feature_extractor.apply = jax.jit(self.feature_extractor.apply) self.feature_extractor.apply = jax.jit(self.feature_extractor.apply)
self.actor.apply = jax.jit(self.actor.apply) self.actor.apply = jax.jit(self.actor.apply)
self.critic.apply = jax.jit(self.critic.apply) self.critic.apply = jax.jit(self.critic.apply)
self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32)
action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) 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) action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32)