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

other(ppo training loop): code cleanup + enabled debug logging of model parameters

Browse source
This commit is contained in:
Robin Meersman 2026-05-04 21:29:27 +02:00
parent dd2612e2ec
commit 791dcf231c
2 changed files with 26 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

2
scripts/train.py
View file

@ -42,7 +42,7 @@ def main(dict_cfg: DictConfig):
base_dir=os.path.dirname(run_dir), base_dir=os.path.dirname(run_dir),
) )
logger = get_logger() logger = get_logger()
logger.set_level(logging.DEBUG) logger.set_level(logging.INFO)
logger.info(f"Hydra-initialized run: {run_name}") logger.info(f"Hydra-initialized run: {run_name}")
logger.info(f"Output directory: {run_dir}") logger.info(f"Output directory: {run_dir}")

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

@ -643,21 +643,18 @@ class PPOTrainer:
sensor_keys = jax.random.split(sensor_key, self.needed_copies) sensor_keys = jax.random.split(sensor_key, self.needed_copies)
actor_keys = jax.random.split(actor_key, self.needed_copies) actor_keys = jax.random.split(actor_key, self.needed_copies)
# note: assumed only 1 message passer needed for now # (needed_copies, X)
# message_passer_keys = jax.random.split(message_passer_key, self.needed_copies)
# (needed_copies, 175)
sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys) sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys)
# self.logger.debug( self.logger.debug(
# f"[_init_agent_state] sensor_params: {jax.tree.map(lambda x: x.shape, sensor_params)}" 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) single_sensor_param = jax.tree.map(lambda x: x[0], sensor_params)
# self.logger.debug( self.logger.debug(
# f"[_init_agent_state] single_sensor_param: { f"[_init_agent_state] single_sensor_param: {
# jax.tree.map(lambda x: x.shape, single_sensor_param) jax.tree.map(lambda x: x.shape, single_sensor_param)
# }" }"
# ) )
sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs) sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs)
self.logger.debug( self.logger.debug(
@ -665,47 +662,42 @@ class PPOTrainer:
) )
actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys) actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys)
# self.logger.debug( self.logger.debug(
# f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}" f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}"
# ) )
message_passer_params = {} message_passer_params = {}
if self.morph_mode != MorphMode.CENTRALIZED: if self.morph_mode != MorphMode.CENTRALIZED:
assert self.message_passer is not None, "MessagePasser is None" assert self.message_passer is not None, "MessagePasser is None"
# message_passer_params = jax.vmap(
# lambda k: self.message_passer.init(
# k, self.sensor.apply(single_sensor_param, sample_obs), self.adj
# )
# )(message_passer_keys)
message_passer_params = self.message_passer.init( message_passer_params = self.message_passer.init(
message_passer_key, message_passer_key,
self.sensor.apply(single_sensor_param, sample_obs), self.sensor.apply(single_sensor_param, sample_obs),
self.adj, self.adj,
) )
# self.logger.debug( self.logger.debug(
# f"[_init_agent_state] message_passer_params: { f"[_init_agent_state] message_passer_params: {
# jax.tree.map(lambda x: x.shape, message_passer_params) jax.tree.map(lambda x: x.shape, message_passer_params)
# }" }"
# ) )
flat_obs = sample_obs.reshape(-1) # BECAUSE 1 centralized critic flat_obs = sample_obs.reshape(-1) # BECAUSE 1 centralized critic
self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}") self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}")
feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs) feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs)
# self.logger.debug( self.logger.debug(
# f"[_init_agent_state] feature_extractor_params: { f"[_init_agent_state] feature_extractor_params: {
# jax.tree.map(lambda x: x.shape, feature_extractor_params) jax.tree.map(lambda x: x.shape, feature_extractor_params)
# }" }"
# ) )
critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs) critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs)
self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}") self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}")
critic_params = self.critic.init(critic_key, critic_input) critic_params = self.critic.init(critic_key, critic_input)
# self.logger.debug( self.logger.debug(
# f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}" f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}"
# ) )
return TrainState.create( return TrainState.create(
apply_fn=None, apply_fn=None,