8a0629ded6
Closes #952 - `SamplingConfig` supports `batch_size=None`. #1077 - tests and examples are covered by `mypy`. #1077 - `NetBase` is more used, stricter typing by making it generic. #1077 - `utils.net.common.Recurrent` now receives and returns a `RecurrentStateBatch` instead of a dict. #1077 --------- Co-authored-by: Michael Panchenko <m.panchenko@appliedai.de>
140 lines
4.8 KiB
Python
Executable File
140 lines
4.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import csv
|
|
import os
|
|
import re
|
|
from collections import defaultdict
|
|
from os import PathLike
|
|
from re import Pattern
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import tqdm
|
|
from tensorboard.backend.event_processing import event_accumulator
|
|
|
|
|
|
def find_all_files(root_dir: str | PathLike[str], pattern: str | Pattern[str]) -> list:
|
|
"""Find all files under root_dir according to relative pattern."""
|
|
file_list = []
|
|
for dirname, _, files in os.walk(root_dir):
|
|
for f in files:
|
|
absolute_path = os.path.join(dirname, f)
|
|
if re.match(pattern, absolute_path):
|
|
file_list.append(absolute_path)
|
|
return file_list
|
|
|
|
|
|
def group_files(file_list: list[str], pattern: str | Pattern[str]) -> dict[str, list]:
|
|
res = defaultdict(list)
|
|
for f in file_list:
|
|
match = re.search(pattern, f)
|
|
key = match.group() if match else ""
|
|
res[key].append(f)
|
|
return res
|
|
|
|
|
|
def csv2numpy(csv_file: str) -> dict[Any, np.ndarray]:
|
|
csv_dict = defaultdict(list)
|
|
with open(csv_file) as f:
|
|
for row in csv.DictReader(f):
|
|
for k, v in row.items():
|
|
csv_dict[k].append(eval(v))
|
|
return {k: np.array(v) for k, v in csv_dict.items()}
|
|
|
|
|
|
def convert_tfevents_to_csv(
|
|
root_dir: str | PathLike[str],
|
|
refresh: bool = False,
|
|
) -> dict[str, list]:
|
|
"""Recursively convert test/reward from all tfevent file under root_dir to csv.
|
|
|
|
This function assumes that there is at most one tfevents file in each directory
|
|
and will add suffix to that directory.
|
|
|
|
:param bool refresh: re-create csv file under any condition.
|
|
"""
|
|
tfevent_files = find_all_files(root_dir, re.compile(r"^.*tfevents.*$"))
|
|
print(f"Converting {len(tfevent_files)} tfevents files under {root_dir} ...")
|
|
result = {}
|
|
with tqdm.tqdm(tfevent_files) as t:
|
|
for tfevent_file in t:
|
|
t.set_postfix(file=tfevent_file)
|
|
output_file = os.path.join(os.path.split(tfevent_file)[0], "test_reward.csv")
|
|
if os.path.exists(output_file) and not refresh:
|
|
with open(output_file) as f:
|
|
content = list(csv.reader(f))
|
|
if content[0] == ["env_step", "reward", "time"]:
|
|
for i in range(1, len(content)):
|
|
content[i] = list(map(eval, content[i]))
|
|
result[output_file] = content
|
|
continue
|
|
ea = event_accumulator.EventAccumulator(tfevent_file)
|
|
ea.Reload()
|
|
initial_time = ea._first_event_timestamp
|
|
content = [["env_step", "reward", "time"]]
|
|
for test_reward in ea.scalars.Items("test/reward"):
|
|
content.append(
|
|
[
|
|
round(test_reward.step, 4),
|
|
round(test_reward.value, 4),
|
|
round(test_reward.wall_time - initial_time, 4),
|
|
],
|
|
)
|
|
with open(output_file, "w") as f:
|
|
csv.writer(f).writerows(content)
|
|
result[output_file] = content
|
|
return result
|
|
|
|
|
|
def merge_csv(
|
|
csv_files: dict[str, list],
|
|
root_dir: str | PathLike[str],
|
|
remove_zero: bool = False,
|
|
) -> None:
|
|
"""Merge result in csv_files into a single csv file."""
|
|
assert len(csv_files) > 0
|
|
if remove_zero:
|
|
for v in csv_files.values():
|
|
if v[1][0] == 0:
|
|
v.pop(1)
|
|
sorted_keys = sorted(csv_files.keys())
|
|
sorted_values = [csv_files[k][1:] for k in sorted_keys]
|
|
content = [
|
|
[
|
|
"env_step",
|
|
"reward",
|
|
"reward:shaded",
|
|
*["reward:" + os.path.relpath(f, root_dir) for f in sorted_keys],
|
|
],
|
|
]
|
|
for rows in zip(*sorted_values, strict=True):
|
|
array = np.array(rows)
|
|
assert len(set(array[:, 0])) == 1, (set(array[:, 0]), array[:, 0])
|
|
line = [rows[0][0], round(array[:, 1].mean(), 4), round(array[:, 1].std(), 4)]
|
|
line += array[:, 1].tolist()
|
|
content.append(line)
|
|
output_path = os.path.join(root_dir, f"test_reward_{len(csv_files)}seeds.csv")
|
|
print(f"Output merged csv file to {output_path} with {len(content[1:])} lines.")
|
|
with open(output_path, "w") as f:
|
|
csv.writer(f).writerows(content)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--refresh",
|
|
action="store_true",
|
|
help="Re-generate all csv files instead of using existing one.",
|
|
)
|
|
parser.add_argument(
|
|
"--remove-zero",
|
|
action="store_true",
|
|
help="Remove the data point of env_step == 0.",
|
|
)
|
|
parser.add_argument("--root-dir", type=str)
|
|
args = parser.parse_args()
|
|
|
|
csv_files = convert_tfevents_to_csv(args.root_dir, args.refresh)
|
|
merge_csv(csv_files, args.root_dir, args.remove_zero)
|