diff --git a/configs/hpc/smoke_test.yaml b/configs/hpc/smoke_test.yaml index 720f829..c892e3e 100644 --- a/configs/hpc/smoke_test.yaml +++ b/configs/hpc/smoke_test.yaml @@ -1,5 +1,5 @@ # Minimal config to verify HPC setup is functional. -# Run with: python src/train.py --config-path configs/hpc/smoke_test.yaml +# Run with: python experiments/train.py --config-path configs/hpc/smoke_test.yaml exp_name: "hpc_smoke_test" seed: 0 track: false # Test WandB integration diff --git a/experiments/PPOTrainer.py b/experiments/PPOTrainer.py index aea69d2..cfadb54 100644 --- a/experiments/PPOTrainer.py +++ b/experiments/PPOTrainer.py @@ -1,4 +1,6 @@ +import datetime import random +import sys import time from dataclasses import asdict, dataclass from functools import partial @@ -200,11 +202,12 @@ class LossInfo: class PPOTrainer: - def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_name: str): + def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str): self.args = args self.env = env + self.run_dir = run_dir self.run_name = run_name - self.writer = SummaryWriter(f"runs/{self.run_name}") + self.writer = SummaryWriter(self.run_dir) self.key = jax.random.PRNGKey(args.seed) @@ -244,11 +247,17 @@ class PPOTrainer: self._init_random() - def _init_random(self): + def _init_random(self, log: bool = True): + if log: + print(f"[RANDOM]: Setting random seed to {self.args.seed}") + random.seed(self.args.seed) np.random.seed(self.args.seed) - def _init_agent(self): + def _init_agent(self, log: bool = True): + if log: + print("[AGENT]: Initializing agent...") + sensor = GenericDenseLayersWithActivation() feature_extractor = GenericDenseLayersWithActivation() actor = Actor( @@ -258,7 +267,10 @@ class PPOTrainer: # messenger = OneDenseLayerMLP() return sensor, feature_extractor, actor, critic - def _init_agent_state(self) -> TrainState: + def _init_agent_state(self, log: bool = True) -> TrainState: + if log: + print("[AGENT STATE]: Initializing agent state...") + self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split( self.key, 5 ) @@ -301,7 +313,10 @@ class PPOTrainer: ), ) - def _init_episode_stats(self) -> EpisodeStatistics: + def _init_episode_stats(self, log: bool = True) -> EpisodeStatistics: + if log: + print("[EPISODE STATS]: Initializing episode stats...") + return EpisodeStatistics( episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32), episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), @@ -335,6 +350,7 @@ class PPOTrainer: iteration_time_start, loss_info, ): + self.writer.add_scalar( "charts/avg_episodic_return", loss_info.avg_episodic_return, global_step ) @@ -353,7 +369,6 @@ class PPOTrainer: 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 ) @@ -363,7 +378,10 @@ class PPOTrainer: global_step, ) - def _step(self, env_state, next_obs, next_done) -> tuple: + def _step(self, env_state, next_obs, next_done, is_tty: bool, iteration: int) -> tuple: + if not is_tty and iteration == 1: + print(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}", flush=True) + ( self.agent_state, self.episode_stats, @@ -374,12 +392,21 @@ class PPOTrainer: next_env_state, ) = self._rollout(env_state, next_obs, next_done) + if not is_tty and iteration == 1: + print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True) + storage = self._compute_gae(storage, next_obs, next_done) + if not is_tty and iteration == 1: + print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True) + 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 not is_tty and iteration == 1: + print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True) + avg_episodic_return = float( jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)) ) @@ -402,7 +429,10 @@ class PPOTrainer: self.env.close() self.writer.close() - def _save_model(self, model_path: str): + def _save_model(self, model_path: str, log: bool = True): + if log: + print(f"[SAVE]: Saving the model to: {model_path}...") + with open(model_path, "wb") as f: f.write( flax.serialization.to_bytes( @@ -418,21 +448,38 @@ class PPOTrainer: ) ) - def train(self): + def train(self, log: bool = True): """ 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}") + + is_tty = sys.stdout.isatty() + if log: + print("[TRAIN]: Resetting environment...") + + if not is_tty: + print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True) + 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) + 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, @@ -442,21 +489,47 @@ class PPOTrainer: 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()), ) - for _ in tqdm.tqdm(range(self.args.num_iterations)): + iter_bar = tqdm.tqdm( + range(1, self.args.num_iterations + 1), + disable=not sys.stdout.isatty(), + ) + for iteration in iter_bar: iteration_time_start = time.time() env_state, next_obs, next_done, loss_info = self._step(env_state, next_obs, next_done) + + if not is_tty and iteration == 1: + print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True) + 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 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( + 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, + ) + if self.args.save_model: - model_path = f"runs/{self.run_name}/{self.args.exp_name}.cleanrl_model" + model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model" self._save_model(model_path=model_path) self._close() diff --git a/experiments/_train_backup.py b/experiments/_train_backup.py index 2cc38ab..a46e500 100644 --- a/experiments/_train_backup.py +++ b/experiments/_train_backup.py @@ -81,7 +81,6 @@ def train(args: PPOArgs): run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" # args.num_iterations = args.total_timesteps // args.batch_size args.num_iterations = 5 - run_name = f"{args.exp_name}__seed_{args.seed}__{int(time.time())}" print(f"running name: {run_name}") if args.run_dir is None: @@ -385,8 +384,6 @@ def train(args: PPOArgs): ) if args.save_model: - model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model" - save_model(model_path, agent_state, args) model_path = f"{args.run_dir}/{args.exp_name}.cleanrl_model" with open(model_path, "wb") as f: f.write( @@ -408,12 +405,6 @@ def train(args: PPOArgs): writer.close() print("Saving loss plot...") - simple_plot( - list(range(len(returns))), - returns, - show_window=True, - filename=f"runs/{run_name}/{args.exp_name}_losses.png", - ) plt.plot(losses) plt.title("PPO Loss, mean over minibatches") plt.savefig(f"{args.run_dir}/{args.exp_name}_losses.png") diff --git a/experiments/train.py b/experiments/train.py index 2e4bca7..596dad9 100644 --- a/experiments/train.py +++ b/experiments/train.py @@ -1,7 +1,10 @@ +import subprocess import time import torch import tyro +import yaml +import os from brittle_star_project.dataclasses import PPOArgs from PPOTrainer import PPOTrainer @@ -14,16 +17,53 @@ def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs) +def parse_args() -> PPOArgs: + temp_args = tyro.cli(PPOArgs) + + if temp_args.env_config_path is not None: + with open(temp_args.env_config_path, "r") as f: + config = yaml.safe_load(f) + if config: + # parse PPOArgs with defaults from yaml. + for key, value in config.items(): + if hasattr(temp_args, key): + setattr(temp_args, key, value) + + # Reparse CLI to ensure they OVERRIDE the yaml + args = tyro.cli(PPOArgs, default=temp_args) + else: + args = temp_args + return args + + +def get_git_hash() -> str: + try: + return ( + subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() + ) + except subprocess.CalledProcessError | UnicodeDecodeError: + return "none" + + if __name__ == "__main__": - args = tyro.cli(PPOArgs) + args = parse_args() args.batch_size = args.num_envs * args.num_steps args.minibatch_size = args.batch_size // args.num_minibatches args.num_iterations = args.total_timesteps // args.batch_size - run_name = f"{args.exp_name}__seed_{args.seed}__{int(time.time())}" - env = make_env(args.config_path, args.num_envs) + + git_hash = get_git_hash() + run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" + if args.run_dir is None: + run_dir = f"runs/{run_name}" + else: + run_dir = args.run_dir + + os.makedirs(run_dir, exist_ok=True) + + env = make_env(args.env_config_path, args.num_envs) torch.backends.cudnn.deterministic = args.torch_deterministic - ppo_trainer = PPOTrainer(args, env, run_name) + ppo_trainer = PPOTrainer(args, env, run_dir, run_name) ppo_trainer.train() diff --git a/uv.lock b/uv.lock index 7fbaacb..fd45cfc 100644 --- a/uv.lock +++ b/uv.lock @@ -35,6 +35,9 @@ dependencies = [ ] [package.optional-dependencies] +analysis = [ + { name = "tensorboard" }, +] cuda = [ { name = "jax", extra = ["cuda13"] }, ] @@ -64,12 +67,13 @@ requires-dist = [ { name = "protobuf", specifier = ">=5.0.0" }, { name = "pyopengl", specifier = ">=3.1.10" }, { name = "pyopengl-accelerate", specifier = ">=3.1.10" }, + { name = "tensorboard", marker = "extra == 'analysis'" }, { name = "torch", specifier = ">=2.4.0" }, { name = "tyro", specifier = ">=1.0.10" }, { name = "wandb", specifier = "==0.24.2" }, { name = "warp-lang" }, ] -provides-extras = ["cuda"] +provides-extras = ["cuda", "analysis"] [package.metadata.requires-dev] dev = [