aggregate_runs.py

import sys
import argparse
import pandas as pd

from tqdm import tqdm


parser = argparse.ArgumentParser(description="Aggregate multiple runs of a model")
parser.add_argument("path_list", type=str, nargs="+", default=[],
                    help="Paths to the progress.csv files")
args = parser.parse_args()


if len(args.path_list) < 2:
    print("You need to give at least two paths to progress.csv files,
          so that we have something to aggregate")
    sys.exit(1)

data_csv = []
for csv_path in args.path_list:
    data_csv.append(pd.read_csv(csv_path))

res = data_csv[0].copy()

visited_rooms = set()
def format_rooms(rooms):
    # Input rooms are strings: '[1;2;3]'
    rooms = rooms.strip("][")
    if rooms:
        rooms = rooms.split(";")
        return rooms
    return []

rooms = format_rooms(data_csv[0].rooms.values[-1])
for r in rooms:
    visited_rooms.add(r)

for data_cv in tqdm(data_csv[1:], desc="files"):
    base_reward = res.rewtotal.values[-1]
    base_tcount = res.tcount.values[-1]
    highest_reward = res.best_ret.values[-1]

    for row in tqdm(data_cv.values, desc="rows", leave=False):
        reward = row[43]
        if pd.isna(reward):
            continue

        tcount = row[32]
        rooms = format_rooms(row[31])
        for r in rooms:
            visited_rooms.add(r)
        highest_reward = max(row[2], highest_reward)

        total_reward = base_reward + reward
        total_tcount = base_tcount + tcount
        frame = list(row)
        frame[2] = highest_reward
        frame[8] = len(visited_rooms)
        frame[43] = total_reward
        frame[32] = total_tcount
        df = pd.DataFrame([frame], columns=list(data_csv[0].columns))
        res = res.append(df, ignore_index=True)

res.to_csv("aggregated_progress.csv")