Files
go2_rl_gym/tools/logs_merge.py
2026-02-09 22:44:57 +08:00

123 lines
5.1 KiB
Python

import time
import yaml
import argparse
import pandas as pd
from tqdm import tqdm
from pathlib import Path
from collections import defaultdict
from tensorboard.compat.proto import event_pb2
from tensorboard.backend.event_processing import event_file_loader
PATH_PARENT = Path(__file__).parent.resolve()
BASE_COLUMNS = [
'it', 'benchmark',
'lin_vel_err',
'ang_vel_err',
'dof_limits',
'dof_power',
'orientation_stability',
'torque_smoothness',
'flat', 'wave', 'obstacle',
'slope_fd', 'slope_bd',
'stairs_fd', 'stairs_bd',
'terrain_level'
]
def fast_read(event_file_path, tag_names):
loader = event_file_loader.RawEventFileLoader(event_file_path)
tag_data = defaultdict(dict)
for raw_event in loader.Load():
event = event_pb2.Event.FromString(raw_event)
if event.HasField('summary'):
for value in event.summary.value:
if value.tag in tag_names:
tag_data[event.step][value.tag] = value.simple_value
return pd.DataFrame(tag_data).T
class Collector:
def __init__(self, log_dirs):
self.log_dirs = Path(log_dirs)
assert self.log_dirs.exists(), f"Log directory {log_dirs} does not exist."
assert self.log_dirs.is_dir(), f"{log_dirs} is not a directory."
alg_name = self.log_dirs.parent.name
date_str = self.log_dirs.name
self.output_dir = PATH_PARENT / f"{alg_name}_{date_str}"
if self.output_dir.exists():
s = input(f"[Warning] Output directory {self.output_dir} already exists, press Enter to continue and overwrite or type 'q' to quit...")
if s.lower() == 'q':
exit(0)
self.output_dir.mkdir(parents=True, exist_ok=True)
self.output_csv = self.output_dir / f"{alg_name}_{date_str}.csv"
self.datas = defaultdict(list)
self.output_tb = self.output_dir / "tb.csv"
if self.output_tb.exists():
print(f"Loading existing tensorboard data from {self.output_tb}")
self.tb_df = pd.read_csv(self.output_tb)
else:
start_time = time.time()
print(f"Start reading tensorboard events at {time.ctime(start_time)}")
self.tb_df = fast_read(str(self.log_dirs.glob("events.out.tfevents.*").__next__()), [
'Terrain/terrain_level_all', 'Episode/terrain_level_all',
'RoboGauge/benchmark'
])
print(f"Finished reading tensorboard events in {time.time() - start_time:.2f} seconds.")
self.tb_df.to_csv(self.output_tb, index=False)
print(f"Saved tensorboard data to {self.output_tb}")
def collect(self):
robogauge_results_path = self.log_dirs / "robogauge_results"
results = list(robogauge_results_path.glob("*.yaml"))
results = sorted(results, key=lambda x: int(x.stem.split("_")[-1]))
for result in tqdm(results):
it = int(result.stem.split("_")[-1])
with open(result, 'r', encoding='utf-8') as f:
data = yaml.safe_load(f)
self.datas['it'].append(it)
self.datas['benchmark'].append(float(data['benchmark_score']))
for metric_name in [
'lin_vel_err',
'ang_vel_err',
'dof_limits',
'dof_power',
'orientation_stability',
'torque_smoothness'
]:
self.datas[f'{metric_name}_mean'].append(float(data['summary'][metric_name]['mean'].split(' ')[0]))
self.datas[f'{metric_name}_mean@25'].append(float(data['summary'][metric_name]['mean@25'].split(' ')[0]))
self.datas[f'{metric_name}_mean@50'].append(float(data['summary'][metric_name]['mean@50'].split(' ')[0]))
for terrain_name in [
'flat',
'wave',
'obstacle',
'slope_fd',
'slope_bd',
'stairs_fd',
'stairs_bd',
]:
if data['robust_score'][terrain_name] is None:
self.datas[f'{terrain_name}_mean'].append(0.0)
self.datas[f'{terrain_name}_mean@25'].append(0.0)
self.datas[f'{terrain_name}_mean@50'].append(0.0)
continue
self.datas[f'{terrain_name}_mean'].append(float(data['robust_score'][terrain_name]['mean']))
self.datas[f'{terrain_name}_mean@25'].append(float(data['robust_score'][terrain_name]['mean@25']))
self.datas[f'{terrain_name}_mean@50'].append(float(data['robust_score'][terrain_name]['mean@50']))
self.datas['terrain_level'].append(float(self.tb_df[self.tb_df['step'] == it]['value'].values[0]))
df = pd.DataFrame(self.datas)
df.to_csv(self.output_csv, index=False)
print(f"Saved merged results to {self.output_csv}")
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--log-dirs")
args = parser.parse_args()
collector = Collector(args.log_dirs)
# collector.collect()