diff --git a/scripts/analysis/explore_tensorboard.py b/scripts/analysis/explore_tensorboard.py index c9e9e5b..0ed1027 100644 --- a/scripts/analysis/explore_tensorboard.py +++ b/scripts/analysis/explore_tensorboard.py @@ -21,13 +21,14 @@ except ImportError: print("Error: Missing dependency. Please run: pip install tensorboard") sys.exit(1) + def explore_run(log_dir): """ Extracts and displays a summary of scalar metrics from a TensorBoard log directory. """ - print(f"\n{'='*20} Exploring Run {'='*20}") + print(f"\n{'=' * 20} Exploring Run {'=' * 20}") print(f"Directory: {log_dir}") - print(f"{'='*55}\n") + print(f"{'=' * 55}\n") if not os.path.exists(log_dir): print(f"Error: Directory '{log_dir}' does not exist.") @@ -35,17 +36,20 @@ def explore_run(log_dir): # Initialize EventAccumulator # size_guidance=0 loads all data points for each tag. - ea = event_accumulator.EventAccumulator(log_dir, size_guidance={ - event_accumulator.SCALARS: 0, - event_accumulator.TENSORS: 0, - }) - + ea = event_accumulator.EventAccumulator( + log_dir, + size_guidance={ + event_accumulator.SCALARS: 0, + event_accumulator.TENSORS: 0, + }, + ) + print("Loading event files (this may take a moment for large runs)...") ea.Reload() tags = ea.Tags() - scalar_tags = tags.get('scalars', []) - + scalar_tags = tags.get("scalars", []) + if not scalar_tags: print("No scalar metrics found in this directory.") return None @@ -60,43 +64,48 @@ def explore_run(log_dir): events = ea.Scalars(tag) if not events: continue - + values = [e.value for e in events] last_event = events[-1] data[tag] = values - - summary.append({ - "Metric": tag, - "Steps": len(events), - "Last Value": f"{last_event.value:.4f}", - "Max": f"{max(values):.4f}", - "Min": f"{min(values):.4f}" - }) + + summary.append( + { + "Metric": tag, + "Steps": len(events), + "Last Value": f"{last_event.value:.4f}", + "Max": f"{max(values):.4f}", + "Min": f"{min(values):.4f}", + } + ) # Display summary table formatted manually - summary = sorted(summary, key=lambda x: x['Metric']) + summary = sorted(summary, key=lambda x: x["Metric"]) print(f"{'Metric':<30} {'Steps':>10} {'Last':>12} {'Max':>12} {'Min':>12}") print("-" * 80) for row in summary: - print(f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} {row['Max']:>12} {row['Min']:>12}") + print( + f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} " + f"{row['Max']:>12} {row['Min']:>12}" + ) # Calculate and display global metadata - if 'charts/SPS' in data: - sps_events = ea.Scalars('charts/SPS') + if "charts/SPS" in data: + sps_events = ea.Scalars("charts/SPS") if len(sps_events) > 1: total_duration_hours = (sps_events[-1].wall_time - sps_events[0].wall_time) / 3600 print(f"\nTotal Recorded Duration: {total_duration_hours:.2f} hours") - + # Estimate completion if total_timesteps is available in hyperparameters try: - hp_tags = [t for t in tags.get('tensors', []) if 'hyperparameters' in t] + hp_tags = [t for t in tags.get("tensors", []) if "hyperparameters" in t] if hp_tags: hp_event = ea.Tensors(hp_tags[0])[0] - hp_text = hp_event.tensor_proto.string_val[0].decode('utf-8') - if 'total_timesteps' in hp_text: - for line in hp_text.split('\n'): - if 'total_timesteps' in line: - target = int(line.split('|')[2].strip()) + hp_text = hp_event.tensor_proto.string_val[0].decode("utf-8") + if "total_timesteps" in hp_text: + for line in hp_text.split("\n"): + if "total_timesteps" in line: + target = int(line.split("|")[2].strip()) current = ea.Scalars(scalar_tags[0])[-1].step percent = (current / target) * 100 print(f"Progress: {current:,} / {target:,} steps ({percent:.1f}%)") @@ -105,31 +114,30 @@ def explore_run(log_dir): return data + def main(): - parser = argparse.ArgumentParser(description="Clean, reproducible TensorBoard exploration tool.") + parser = argparse.ArgumentParser(description="Reproducible TensorBoard exploration tool.") parser.add_argument("log_dir", help="Path to the TensorBoard run directory.") - parser.add_argument("--csv", help="Optional: Path to save all scalar data as a CSV.", default=None) - + parser.add_argument("--csv", help="Optional: Path to export scalar data to CSV.", default=None) + args = parser.parse_args() - + scalar_data = explore_run(args.log_dir) - + if args.csv and scalar_data: # Reloading for wall_time and steps ea = event_accumulator.EventAccumulator(args.log_dir).Reload() - with open(args.csv, mode='w', newline='') as f: + with open(args.csv, mode="w", newline="") as f: writer = csv.DictWriter(f, fieldnames=["tag", "step", "value", "wall_time"]) writer.writeheader() for tag in scalar_data.keys(): for e in ea.Scalars(tag): - writer.writerow({ - "tag": tag, - "step": e.step, - "value": e.value, - "wall_time": e.wall_time - }) - + writer.writerow( + {"tag": tag, "step": e.step, "value": e.value, "wall_time": e.wall_time} + ) + print(f"\nData exported to: {args.csv}") + if __name__ == "__main__": main() diff --git a/scripts/export_hpc_requirements.py b/scripts/export_hpc_requirements.py index 078d54e..9ff3de7 100644 --- a/scripts/export_hpc_requirements.py +++ b/scripts/export_hpc_requirements.py @@ -53,17 +53,17 @@ def main() -> None: dep_dict: dict[str, str] = {} for dep in data.get("project", {}).get("dependencies", []): dep_dict[normalise(pkg_name(dep))] = dep - + # Add cuda extras (takes precedence for HPC) optional_deps = data.get("project", {}).get("optional-dependencies", {}) for group in ["cuda"]: for dep in optional_deps.get(group, []): dep_dict[normalise(pkg_name(dep))] = dep - + deps = list(dep_dict.values()) final_deps: list[str] = [] - print(f"Checking dependencies against HPC module list...", file=sys.stderr) + print("Checking dependencies against HPC module list...", file=sys.stderr) for dep in deps: name = normalise(pkg_name(dep)) # Smart check: if the package name is a substring of any loaded module name diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 1420e19..a693286 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -55,10 +55,11 @@ def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]: with open(path, "r") as f: if path.endswith(".yaml") or path.endswith(".yml"): import yaml + config_dict = yaml.safe_load(f) else: config_dict = json.load(f) - + morphology = MorphologyConfig(**config_dict.get("morphology", {})) arena = ArenaConfig(**config_dict.get("arena", {})) env = EnvConfig(**config_dict.get("env", {})) diff --git a/src/train.py b/src/train.py index f8c3515..21c2b9b 100644 --- a/src/train.py +++ b/src/train.py @@ -52,10 +52,12 @@ def train(args: PPOArgs): 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 - + # Try to get git short hash try: - git_hash = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() + git_hash = ( + subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() + ) except Exception: git_hash = "none" @@ -64,8 +66,9 @@ def train(args: PPOArgs): if args.run_dir is None: args.run_dir = f"runs/{run_name}" - + import os + os.makedirs(args.run_dir, exist_ok=True) if args.track: @@ -348,14 +351,14 @@ def train(args: PPOArgs): remaining_steps = 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}/{args.num_iterations} | " f"Step {global_step}/{args.total_timesteps} | " f"SPS {sps} | " f"Return {avg_episodic_return:.4f} | " - f"ETA {eta_str}", - flush=True + f"ETA {eta_str}", + flush=True, ) if args.save_model: @@ -388,7 +391,7 @@ def train(args: PPOArgs): def main() -> None: 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) @@ -397,7 +400,7 @@ def main() -> None: for key, value in config.items(): if hasattr(temp_args, key): setattr(temp_args, key, value) - + # Re-parse CLI to ensure they OVERRIDE the yaml args = tyro.cli(PPOArgs, default=temp_args) else: