Static native

cache_data.py

@@ -4,18 +4,20 @@
 import glob
 import os
 
+import pufferlib
+
 
 env_names = sorted([
     'breakout',
     #'impulse_wars',
     #'pacman',
     #'tetris',
-    #'g2048',
+    'g2048',
     #'moba',
-    #'pong',
+    'pong',
     #'tower_climb',
-    #'grid',
-    #'nmmo3',
+    'grid',
+    'nmmo3',
     #'snake',
     #'tripletriad'
 ])
@@ -36,19 +38,21 @@
     'train/eps',
     'train/prio_alpha',
     'train/prio_beta0',
-    'train/horizon',
+    #'train/horizon',
     'train/replay_ratio',
     'train/minibatch_size',
     'policy/hidden_size',
     'vec/total_agents',
 ]
 
-ALL_KEYS = [
+METRICS = [
     'agent_steps',
-    'cost',
-    'environment/score',
-    'environment/perf'
-] + HYPERS
+    'uptime',
+    'env/score',
+    'env/perf',
+]
+
+ALL_KEYS = HYPERS + METRICS
 
 def pareto_idx(steps, costs, scores):
     idxs = []
@@ -63,62 +67,77 @@ def pareto_idx(steps, costs, scores):
 
 def load_sweep_data(path):
     data = {}
-    keys = None
+    sweep_metadata = {}
+    num_metrics = 0
     for fpath in glob.glob(path):
         if 'cache.json' in fpath:
             continue
 
         with open(fpath, 'r') as f:
-            exp = json.load(f)
-
-        if not data:
-            for kk in exp.keys():
-                if kk == 'data':
-                    for k, v in exp[kk][-1].items():
-                        data[k] = []
-                else:
-                    data[kk] = []
-
-        discard = False
-        for kk in list(data.keys()):
-            if kk not in exp and kk not in exp['data'][-1]:
-                discard = True
+            try:
+                exp = json.load(f)
+            except json.decoder.JSONDecodeError:
+                print(f'Skipping {fpath}')
+                continue
+
+        sweep_metadata = exp.pop('sweep')
+
+        data_len = len(exp['metrics']['agent_steps'])
+        if data_len > 100:
+            print(f'Skipping {fpath} (len={data_len})')
+            continue
+
+        if num_metrics == 0:
+            num_metrics = len(exp['metrics'])
+
+        skip = False
+        metrics = exp.pop('metrics')
+
+        if len(metrics) != num_metrics:
+            print(f'Skipping {fpath} (num_metrics={len(metrics)} != {num_metrics})')
+            continue
+
+        n = len(metrics['agent_steps'])
+        for k, v in metrics.items():
+            if len(v) != n:
+                skip = True
+                break
+
+            if k not in data:
+                data[k] = []
+
+            if np.isnan(v).any():
+                skip = True
                 break
 
-        if discard:
+        if skip:
+            print(f'Skipping {fpath} (bad data)')
             continue
 
-        for kk in list(data.keys()):
-            if kk in exp:
-                v = exp[kk]
-                sweep_key = f'sweep/{kk}/distribution'
-                if sweep_key in data and exp[sweep_key] == 'logit_normal':
-                    v = 1 - v
-                elif kk in ('train/vtrace_rho_clip', 'train/vtrace_c_clip'):
-                    v = max(v, 0.1)
+        for k, v in metrics.items():
+            data[k].append(v)
+            if len(data[k]) != len(data['SPS']):
+                breakpoint()
+                pass
 
-                data[kk].append(v)
-            else:
-                data[kk].append(exp['data'][-1][kk])
+        for k, v in pufferlib.unroll_nested_dict(exp):
+            if k not in data:
+                data[k] = []
 
-    steps = data['agent_steps']
-    costs = data['cost']
-    scores = data['environment/score']
+            data[k].append([v]*n)
 
-    idxs = pareto_idx(steps, costs, scores)
+    for k, v in data.items():
+        data[k] = [item for sublist in v for item in sublist]
 
+    #steps = data['agent_steps']
+    #costs = data['uptime']
+    #scores = data['env/score']
+    #idxs = pareto_idx(steps, costs, scores)
     # Filter to pareto
-    for k in data:
-        data[k] = [data[k][i] for i in idxs]
-
-    # Monkey patch: Cap performance
-    data['environment/perf'] = [min(e, 1.0) for e in data['environment/perf']]
-
-    # Monkey patch: Adjust steps by frameskip if present
-    if 'env/frameskip' in data:
-        skip = data['env/frameskip']
-        data['agent_steps'] = [n*m for n, m in zip(data['agent_steps'], skip)]
-
+    #for k in data:
+    #    data[k] = [data[k][i] for i in idxs]
+
+    data['sweep'] = sweep_metadata
     return data
 
 def cached_sweep_load(path, env_name):
@@ -135,31 +154,32 @@ def cached_sweep_load(path, env_name):
     return data
 
 def compute_tsne():
-    data = {name: cached_sweep_load(f'experiments/logs/puffer_{name}', name) for name in env_names}
+    all_data = {}
+    normed = {}
 
-    flat = []
-    flat_mmin = []
-    flat_mmax = []
     for env in env_names:
-        flat.append(np.stack([data[env][hyper] for hyper in HYPERS], axis=1))
-        flat_mmin.append(np.stack([data[env][f'sweep/{hyper}/min'] for hyper in HYPERS], axis=1))
-        flat_mmax.append(np.stack([data[env][f'sweep/{hyper}/max'] for hyper in HYPERS], axis=1))
+        env_data = cached_sweep_load(f'logs/puffer_{env}', env)
+        sweep_metadata = env_data.pop('sweep')
+        all_data[env] = env_data
+
+        normed_env = []
+        for key in HYPERS:
+            prefix, suffix = key.split('/')
+            mmin = sweep_metadata[prefix][suffix]['min']
+            mmax = sweep_metadata[prefix][suffix]['max']
+            dat = np.array(env_data[key])
 
-    flat_distribution = [data[env][f'sweep/{hyper}/distribution'] for env in env_names for hyper in HYPERS]
+            dist = sweep_metadata[prefix][suffix]['distribution']
+            if 'log' in dist or 'pow2' in dist:
+                mmin = np.log(mmin)
+                mmax = np.log(mmax)
+                dat = np.log(dat)
 
-    flat = np.concatenate(flat, axis=0)
-    flat_mmin = np.concatenate(flat_mmin, axis=0).min(axis=0)
-    flat_mmax = np.concatenate(flat_mmax, axis=0).max(axis=0)
+            normed_env.append((dat - mmin) / (mmax - mmin))
 
-    normed = flat.copy()
-    for i in range(len(HYPERS)):
-        dist = flat_distribution[i]
-        if 'log' in dist or 'pow2' in dist:
-            flat_mmin[i] = np.log(flat_mmin[i])
-            flat_mmax[i] = np.log(flat_mmax[i])
-            normed[:, i] = np.log(flat[:, i])
+        normed[env] = np.stack(normed_env, axis=1)
 
-        normed[:, i] = (normed[:, i] - flat_mmin[i]) / (flat_mmax[i] - flat_mmin[i])
+    normed = np.concatenate(list(normed.values()), axis=0)
 
     from sklearn.manifold import TSNE
     proj = TSNE(n_components=2)
@@ -171,25 +191,19 @@ def compute_tsne():
 
     row = 0
     for env in env_names:
-        '''
-        for i, hyper in enumerate(HYPERS):
-            sz = len(data[env][hyper])
-            data[env][hyper] = normed[row:row+sz, i].tolist()
-        '''
-        sz = len(data[env]['agent_steps'])
-
-        data[env] = {k: v for k, v in data[env].items() if k in ALL_KEYS}
+        sz = len(all_data[env]['agent_steps'])
+        #all_data[env] = {k: v for k, v in all_data[env].items()}
         if reduced is not None:
-            data[env]['tsne1'] = reduced[row:row+sz, 0].tolist()
-            data[env]['tsne2'] = reduced[row:row+sz, 1].tolist()
+            all_data[env]['tsne1'] = reduced[row:row+sz, 0].tolist()
+            all_data[env]['tsne2'] = reduced[row:row+sz, 1].tolist()
         else:
-            data[env]['tsne1'] = np.random.rand(sz).tolist()
-            data[env]['tsne2'] = np.random.rand(sz).tolist()
+            all_data[env]['tsne1'] = np.random.rand(sz).tolist()
+            all_data[env]['tsne2'] = np.random.rand(sz).tolist()
 
         row += sz
         print(f'Env {env} has {sz} points')
 
-    json.dump(data, open('all_cache.json', 'w'))
+    json.dump(all_data, open('all_cache.json', 'w'))
 
 if __name__ == '__main__':
     compute_tsne()