drive, tower_climb fixes

Author: jsuarez5341

constellation/cache_data.py

@@ -62,7 +62,7 @@ def pareto_idx(steps, costs, scores):
 
     return idxs
 
-def cached_load(path, env_name, cache):
+def cached_load(path, env_name, cache, full_dataset=False):
     data = {}
     num_metrics = 0
     metric_keys = []
@@ -170,21 +170,22 @@ def cached_load(path, env_name, cache):
     del data['metrics/agent_steps']
 
     # Filter to pareto
-    steps = data['agent_steps']
-    costs = data['uptime']
-    scores = data['env/score']
-
-    idxs = pareto_idx(steps, costs, scores)
-    for k in data:
-        try:
-            data[k] = [data[k][i] for i in idxs]
-        except IndexError:
-            continue
+    if not full_dataset:
+        steps = data['agent_steps']
+        costs = data['uptime']
+        scores = data['env/score']
+
+        idxs = pareto_idx(steps, costs, scores)
+        for k in data:
+            try:
+                data[k] = [data[k][i] for i in idxs]
+            except IndexError:
+                continue
 
     data['sweep'] = sweep_metadata
     return data
 
-def compute_tsne():
+def compute_tsne(full_dataset=False):
     all_data = {}
     normed = {}
 
@@ -196,7 +197,7 @@ def compute_tsne():
     env_names = sorted(os.listdir('logs'))
     for env in env_names:
         print('Loading: ', env)
-        all_data[env] = cached_load(f'logs/{env}/*.json', env, cache)
+        all_data[env] = cached_load(f'logs/{env}/*.json', env, cache, full_dataset)
 
     with open(cache_file, 'w') as f:
         json.dump(cache, f)
@@ -230,7 +231,7 @@ def compute_tsne():
                 and len(v) > 0 and isinstance(v[0], (int, float))
                 and (k == 'train/max_grad_norm' or not k.endswith('_norm'))}
         all_data[env] = dat
-        print(f'Env {env} has {len(dat['env/perf'])} points')
+        print(f"Env {env} has {len(dat['env/perf'])} points")
         for k, v in dat.items():
             if 'env/perf' in k or 'score' in k:
                 print(f'{env}/{k}: min={min(v)}, max={max(v)}')
@@ -243,4 +244,8 @@ def compute_tsne():
     json.dump(all_data, open('resources/constellation/experiments.json', 'w'))
 
 if __name__ == '__main__':
-    compute_tsne()
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--full', action='store_true', help='Include full dataset (no pareto filtering)')
+    args = parser.parse_args()
+    compute_tsne(full_dataset=args.full)

ocean/drive/binding.c

@@ -3,7 +3,7 @@
 #define ACT_SIZES {7, 13}
 #define OBS_TENSOR_T FloatTensor
 
-#define MAP_BINARY_DIR "resources/drive/binaries/training"
+#define MAP_BINARY_DIR "drive_data/binaries"
 
 #define MY_VEC_INIT
 #define Env Drive

ocean/drive/dataset.py

@@ -4,14 +4,14 @@
   https://huggingface.co/datasets/daphne-cornelisse/pufferdrive_womd_train_1000
 
   uv pip install huggingface_hub
-  python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='daphne-cornelisse/pufferdrive_womd_train_1000', repo_type='dataset', local_dir='resources/drive/data')"
+  python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='daphne-cornelisse/pufferdrive_womd_train_1000', repo_type='dataset', local_dir='drive_data')"
 
 Step 1: Unzip to get folder with .json files
-    mkdir -p resources/drive/data/training
-    tar xzf resources/drive/data/pufferdrive_womd_train_1000.tar.gz --strip-components=1 -C resources/drive/data/training/
+    mkdir -p drive_data/training
+    tar xzf drive_data/pufferdrive_womd_train_1000.tar.gz --strip-components=1 -C drive_data/training/
 
 Step 2: Process to map binaries
-  python ocean/drive/dataset.py --data_folder resources/drive/data/training --output_dir resources/drive/binaries/training
+  python ocean/drive/dataset.py --data_folder drive_data/training --output_dir drive_data/binaries
 """
 
 import json