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

refactor(log): PPOTrainer

Browse source
This commit is contained in:
Tibo De Peuter 2026-04-08 19:32:48 +02:00
parent 795ae48520
commit faf31567e8
Signed by: tdpeuter
SSH key fingerprint: SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
1 changed files with 61 additions and 107 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -13,7 +13,8 @@ import numpy as np
import optax
import tqdm
from flax.training.train_state import TrainState
from torch.utils.tensorboard import SummaryWriter
from experiment_logger import get_logger
from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
@ -207,7 +208,7 @@ class PPOTrainer:
self.env = env
self.run_dir = run_dir
self.run_name = run_name
self.writer = SummaryWriter(self.run_dir)
self.logger = get_logger()
self.key = jax.random.PRNGKey(args.seed)
@ -247,16 +248,14 @@ class PPOTrainer:
self._init_random()
def _init_random(self, log: bool = True):
if log:
print(f"[RANDOM]: Setting random seed to {self.args.seed}")
def _init_random(self):
self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}")
random.seed(self.args.seed)
np.random.seed(self.args.seed)
def _init_agent(self, log: bool = True):
if log:
print("[AGENT]: Initializing agent...")
def _init_agent(self):
self.logger.info("[AGENT]: Initializing agent...")
sensor = GenericDenseLayersWithActivation()
feature_extractor = GenericDenseLayersWithActivation()
@ -267,9 +266,8 @@ class PPOTrainer:
# messenger = OneDenseLayerMLP()
return sensor, feature_extractor, actor, critic
def _init_agent_state(self, log: bool = True) -> TrainState:
if log:
print("[AGENT STATE]: Initializing agent state...")
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
@ -313,9 +311,8 @@ class PPOTrainer:
),
)
def _init_episode_stats(self, log: bool = True) -> EpisodeStatistics:
if log:
print("[EPISODE STATS]: Initializing episode stats...")
def _init_episode_stats(self) -> EpisodeStatistics:
self.logger.info("[EPISODE STATS]: Initializing episode stats...")
return EpisodeStatistics(
episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
@ -350,39 +347,29 @@ class PPOTrainer:
iteration_time_start,
loss_info,
):
metrics = {
"charts/avg_episodic_return": loss_info.avg_episodic_return,
"charts/avg_episodic_length": np.mean(
jax.device_get(episode_stats.returned_episode_lengths)
),
"charts/learning_rate": self.agent_state.opt_state[1]
.hyperparams["learning_rate"]
.item(),
"losses/value_loss": loss_info.v_loss[-1, -1].item(),
"losses/policy_loss": loss_info.pg_loss[-1, -1].item(),
"losses/entropy": loss_info.entropy_loss[-1, -1].item(),
"losses/approx_kl": loss_info.approx_kl[-1, -1].item(),
"losses/loss": loss_info.loss[-1, -1].item(),
"charts/SPS": int(global_step / (time.time() - start_time)),
"charts/SPS_update": int(
self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)
),
}
self.logger.log(metrics, step=global_step)
self.writer.add_scalar(
"charts/avg_episodic_return", loss_info.avg_episodic_return, global_step
)
self.writer.add_scalar(
"charts/avg_episodic_length",
np.mean(jax.device_get(episode_stats.returned_episode_lengths)),
global_step,
)
self.writer.add_scalar(
"charts/learning_rate",
self.agent_state.opt_state[1].hyperparams["learning_rate"].item(),
global_step,
)
self.writer.add_scalar("losses/value_loss", loss_info.v_loss[-1, -1].item(), global_step)
self.writer.add_scalar("losses/policy_loss", loss_info.pg_loss[-1, -1].item(), global_step)
self.writer.add_scalar("losses/entropy", loss_info.entropy_loss[-1, -1].item(), global_step)
self.writer.add_scalar("losses/approx_kl", loss_info.approx_kl[-1, -1].item(), global_step)
self.writer.add_scalar("losses/loss", loss_info.loss[-1, -1].item(), global_step)
self.writer.add_scalar(
"charts/SPS", int(global_step / (time.time() - start_time)), global_step
)
self.writer.add_scalar(
"charts/SPS_update",
int(self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)),
global_step,
)
def _step(
self, env_state, next_obs, next_done, is_tty: bool, iteration: int, log: bool = True
) -> tuple:
if log and not is_tty and iteration == 1:
print(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}", flush=True)
def _step(self, env_state, next_obs, next_done, is_tty: bool, iteration: int) -> tuple:
if not is_tty and iteration == 1:
self.logger.info(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}")
(
self.agent_state,
@ -394,20 +381,20 @@ class PPOTrainer:
next_env_state,
) = self._rollout(env_state, next_obs, next_done)
if log and not is_tty and iteration == 1:
print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True)
if not is_tty and iteration == 1:
self.logger.info(f">>> [HPC] First rollout completed: {time.ctime()}")
storage = self._compute_gae(storage, next_obs, next_done)
if log and not is_tty and iteration == 1:
print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True)
if not is_tty and iteration == 1:
self.logger.info(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}")
self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = (
self._ppo.update_ppo(self.agent_state, storage, self.key)
)
if log and not is_tty and iteration == 1:
print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True)
if not is_tty and iteration == 1:
self.logger.info(f">>> [HPC] First PPO update completed: {time.ctime()}")
avg_episodic_return = float(
jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item()
@ -429,77 +416,45 @@ class PPOTrainer:
def _close(self):
self.env.close()
self.writer.close()
def _save_model(self, model_path: str, log: bool = True):
if log:
print(f"[SAVE]: Saving the model to: {model_path}...")
def _save_model(self, model_path: str):
self.logger.info("[SAVE]: Saving the final model...")
with open(model_path, "wb") as f:
f.write(
flax.serialization.to_bytes(
[
vars(self.args),
[
self.agent_state.params["sensor_params"],
self.agent_state.params["actor_params"],
self.agent_state.params["critic_params"],
self.agent_state.params["feature_extractor_params"],
],
]
)
)
params = [
vars(self.args),
[
self.agent_state.params["sensor_params"],
self.agent_state.params["actor_params"],
self.agent_state.params["critic_params"],
self.agent_state.params["feature_extractor_params"],
],
]
self.logger.save_final_model(params=params)
def train(self, log: bool = True):
def train(self):
"""
Train the PPO agent for a specified number of iterations
(passed through PPOArgs in constructor).
Closes the environment at the end of training.
"""
if log:
print(f"running name: {self.run_name}")
self.logger.info(f"running name: {self.run_name}")
is_tty = sys.stdout.isatty()
if log:
print("[TRAIN]: Resetting environment...")
self.logger.info("[TRAIN]: Resetting environment...")
if not is_tty:
print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True)
if not is_tty:
self.logger.info(f">>> [HPC] Initial reset started: {time.ctime()}")
env_state = self.env.reset(seed=self.args.seed)
next_obs = _convert_obs_dict_to_array(env_state.observations)
next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_)
if log and not is_tty:
print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True)
if not is_tty:
self.logger.info(f">>> [HPC] Initial reset completed: {time.ctime()}")
global_step = 0
start_time = time.time()
if self.args.track:
import wandb
if log:
print("[TRAIN]: Initializing Weights and Biases...")
wandb.init(
project=self.args.wandb_project_name,
entity=self.args.wandb_entity,
sync_tensorboard=True,
config=vars(self.args),
name=self.run_name,
save_code=True,
)
if log:
print("[TRAIN]: Adding hyperparameters to TensorBoard...")
self.writer.add_text(
"hyperparameters",
"|param|value|\n|---|---|\n"
+ "\n".join(f"|{k}|{v}|" for k, v in vars(self.args).items()),
)
iter_bar = tqdm.tqdm(
range(1, self.args.num_iterations + 1),
disable=not is_tty,
@ -514,19 +469,18 @@ class PPOTrainer:
global_step += self.args.num_steps * self.args.num_envs
self._log(global_step, self.episode_stats, start_time, iteration_time_start, loss_info)
if log and not is_tty:
if not is_tty:
sps = int(global_step / (time.time() - start_time))
remaining_steps = self.args.total_timesteps - global_step
eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
eta_str = str(datetime.timedelta(seconds=eta_seconds))
print(
self.logger.info(
f"Iteration {iteration}/{self.args.num_iterations} | "
f"Step {global_step}/{self.args.total_timesteps} | "
f"SPS {sps} | "
f"Return {loss_info.avg_episodic_return:.4f} | "
f"ETA {eta_str}",
flush=True,
f"ETA {eta_str}"
)
if self.args.save_model: