mirror of
https://github.com/wassname/cr-dv3.git
synced 2026-09-09 11:20:12 +08:00
127 lines
4.3 KiB
Python
127 lines
4.3 KiB
Python
"""Simple script to plot results of a crafter run with dreamerv3-cr."""
|
|
import pathlib, os
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
import collections
|
|
import json
|
|
import warnings
|
|
import argparse
|
|
|
|
def main():
|
|
runs = []
|
|
|
|
home = os.path.expanduser("~")
|
|
default_filename = f'{home}/logdir/crafter-dv3-cr_1/stats.jsonl'
|
|
|
|
parser = argparse.ArgumentParser(description='Process some arguments')
|
|
parser.add_argument('--filename', default=default_filename,
|
|
help='The path to stats.jsonl file')
|
|
args = parser.parse_args()
|
|
|
|
filename = args.filename
|
|
print(f"Filename: {filename}")
|
|
budget=1e6
|
|
rewards, lengths, achievements = load_stats(pathlib.Path(filename), budget)
|
|
task, method, seed = pathlib.Path(filename).parts[-5:-2]
|
|
|
|
print(f'Run length {sum(lengths)}: {filename}')
|
|
runs.append(dict(
|
|
task=task,
|
|
method=method,
|
|
seed=str(id),
|
|
xs=np.cumsum(lengths).tolist(),
|
|
reward=rewards,
|
|
length=lengths,
|
|
**achievements,
|
|
))
|
|
|
|
scores, tasks, percents, methods = print_summary(runs, budget, verbose=True)
|
|
|
|
if len(rewards) < 100:
|
|
print('>>> Not plotting reward curve until at least 100 episodes saved out')
|
|
else:
|
|
plt.plot(np.cumsum(lengths).tolist(), pd.Series(rewards).rolling(window=100).mean().values)
|
|
plt.xlabel('Steps')
|
|
plt.ylabel('Reward')
|
|
plt.show()
|
|
|
|
def load_stats(filename, budget):
|
|
steps = 0
|
|
rewards = []
|
|
lengths = []
|
|
achievements = collections.defaultdict(list)
|
|
for line in filename.read_text().split('\n'):
|
|
if not line.strip():
|
|
continue
|
|
episode = json.loads(line)
|
|
steps += episode['length']
|
|
if steps > budget:
|
|
break
|
|
lengths.append(episode['length'])
|
|
for key, value in episode.items():
|
|
if key.startswith('achievement_'):
|
|
achievements[key].append(value)
|
|
unlocks = int(np.sum([(v[-1] >= 1) for v in achievements.values()]))
|
|
health = -0.9
|
|
rewards.append(unlocks + health)
|
|
return rewards, lengths, achievements
|
|
|
|
|
|
def print_summary(runs, budget, verbose):
|
|
episodes = np.array([len(x['length']) for x in runs])
|
|
rewards = np.array([np.mean(x['reward']) for x in runs])
|
|
lengths = np.array([np.mean(x['length']) for x in runs])
|
|
percents, methods, seeds, tasks = compute_success_rates(
|
|
runs, budget, sortby=0)
|
|
scores = np.squeeze(compute_scores(percents))
|
|
print(f'Score: {np.mean(scores):10.2f} ± {np.std(scores):.2f}')
|
|
print(f'Reward: {np.mean(rewards):10.2f} ± {np.std(rewards):.2f}')
|
|
print(f'Length: {np.mean(lengths):10.2f} ± {np.std(lengths):.2f}')
|
|
print(f'Episodes: {np.mean(episodes):10.2f} ± {np.std(episodes):.2f}')
|
|
if verbose:
|
|
for task, percent in zip(tasks, np.squeeze(percents).T):
|
|
name = task[len('achievement_'):].replace('_', ' ').title()
|
|
print(f'{name:<20} {np.mean(percent):6.2f}%')
|
|
return scores, tasks, percents, methods
|
|
|
|
|
|
def compute_success_rates(runs, budget=1e6, sortby=None):
|
|
methods = sorted(set(run['method'] for run in runs))
|
|
seeds = sorted(set(run['seed'] for run in runs))
|
|
tasks = sorted(key for key in runs[0] if key.startswith('achievement_'))
|
|
percents = np.empty((len(methods), len(seeds), len(tasks)))
|
|
percents[:] = np.nan
|
|
for run in runs:
|
|
episodes = (np.array(run['xs']) <= budget).sum()
|
|
i = methods.index(run['method'])
|
|
j = seeds.index(run['seed'])
|
|
for key, values in run.items():
|
|
if key in tasks:
|
|
k = tasks.index(key)
|
|
percent = 100 * (np.array(values[:episodes]) >= 1).mean()
|
|
if np.isnan(percent):
|
|
print(percent)
|
|
percents[i, j, k] = percent
|
|
if isinstance(sortby, (str, int)):
|
|
if isinstance(sortby, str):
|
|
sortby = methods.index(sortby)
|
|
order = np.argsort(-np.nanmean(percents[sortby], 0), -1)
|
|
percents = percents[:, :, order]
|
|
tasks = np.array(tasks)[order].tolist()
|
|
return percents, methods, seeds, tasks
|
|
|
|
|
|
def compute_scores(percents):
|
|
# Geometric mean with an offset of 1%.
|
|
assert (0 <= percents).all() and (percents <= 100).all()
|
|
if (percents <= 1.0).all():
|
|
print('Warning: The input may not be in the right range.')
|
|
with warnings.catch_warnings(): # Empty seeds become NaN.
|
|
warnings.simplefilter('ignore', category=RuntimeWarning)
|
|
scores = np.exp(np.nanmean(np.log(1 + percents), -1)) - 1
|
|
return scores
|
|
|
|
if __name__ == "__main__":
|
|
main()
|