pneuma-godot/aggregator.py

39 lines
1.2 KiB
Python

from glob import glob
import tensorboard_reducer as tbr
from policy import policies
strict_steps=False
overwrite = True
reduce_ops = ("mean", "min", "max")
for policy_name in policies:
print(f"Calculating {policy_name}")
# Define glob pattern to find the directories
input_event_dirs = sorted(glob(f"tnsrbrd/sb3/{policy_name}_[0-9]"))
print(f"Found {input_event_dirs}")
# Define output dir
events_output_dir = f"tnsrbrd/results/{policy_name}"
# Load TB events
events_dict = tbr.load_tb_events(input_event_dirs, strict_steps=False, verbose=True)
# Get number of steps and other params
n_scalars = len(events_dict)
n_steps, n_events = list(events_dict.values())[0].shape
print(f"Loaded {n_events} TensorBoard runs with {n_scalars} scalars and {n_steps} steps each")
print(", ".join(events_dict))
# Reduce events
reduced_events = tbr.reduce_events(events_dict, reduce_ops, verbose=True)
for op in reduce_ops:
print(f"Writing '{op}' reduction to '{events_output_dir}-{op}'")
tbr.write_tb_events(reduced_events, events_output_dir, overwrite)
print(f"Reduction of {policy_name} done.")