Merge branch 'PufferAI:4.0' into 4.0

Author: l1onh3art88

cache_data.py

@@ -12,12 +12,12 @@
     #'impulse_wars',
     #'pacman',
     #'tetris',
-    'g2048',
+    #'g2048',
     #'moba',
-    'pong',
+    #'pong',
     #'tower_climb',
-    'grid',
-    'nmmo3',
+    #'grid',
+    #'nmmo3',
     #'snake',
     #'tripletriad'
 ])
@@ -38,7 +38,7 @@
     'train/eps',
     'train/prio_alpha',
     'train/prio_beta0',
-    #'train/horizon',
+    'train/horizon',
     'train/replay_ratio',
     'train/minibatch_size',
     'policy/hidden_size',
@@ -65,22 +65,21 @@ def pareto_idx(steps, costs, scores):
 
     return idxs
 
-def load_sweep_data(path):
+def cached_load(path, env_name, cache):
     data = {}
-    sweep_metadata = {}
     num_metrics = 0
     for fpath in glob.glob(path):
-        if 'cache.json' in fpath:
-            continue
-
-        with open(fpath, 'r') as f:
-            try:
-                exp = json.load(f)
-            except json.decoder.JSONDecodeError:
-                print(f'Skipping {fpath}')
-                continue
+        if fpath in cache:
+            exp = cache[fpath]
+        else:
+            with open(fpath, 'r') as f:
+                try:
+                    exp = json.load(f)
+                except json.decoder.JSONDecodeError:
+                    print(f'Skipping {fpath}')
+                    continue
 
-        sweep_metadata = exp.pop('sweep')
+        cache[fpath] = exp
 
         data_len = len(exp['metrics']['agent_steps'])
         if data_len > 100:
@@ -91,7 +90,7 @@ def load_sweep_data(path):
             num_metrics = len(exp['metrics'])
 
         skip = False
-        metrics = exp.pop('metrics')
+        metrics = exp['metrics']
 
         if len(metrics) != num_metrics:
             print(f'Skipping {fpath} (num_metrics={len(metrics)} != {num_metrics})')
@@ -120,62 +119,77 @@ def load_sweep_data(path):
                 breakpoint()
                 pass
 
+        sweep_metadata = exp['sweep']
+
         for k, v in pufferlib.unroll_nested_dict(exp):
             if k not in data:
                 data[k] = []
 
             data[k].append([v]*n)
 
+        for hyper in HYPERS:
+            prefix, suffix = hyper.split('/')
+            if prefix not in sweep_metadata:
+                continue
+
+            group = sweep_metadata[prefix]
+            if suffix not in group:
+                continue
+
+            param = group[suffix]
+
+            key = f'{prefix}/{suffix}_norm'
+            if key not in data:
+                data[key] = []
+
+            mmin = param['min']
+            mmax = param['max']
+            dist = param['distribution']
+            val = exp[prefix][suffix]
+
+            if 'log' in dist or 'pow2' in dist:
+                mmin = np.log(mmin)
+                mmax = np.log(mmax)
+                val = np.log(val)
+
+            norm = (val - mmin) / (mmax - mmin)
+            data[key].append([norm]*n)
+
     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]
+    steps = data['agent_steps']
+    costs = data['uptime']
+    scores = data['env/score']
+    idxs = pareto_idx(steps, costs, scores)
+    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):
-    cache_file = os.path.join(path, 'c_cache.json')
-    if not os.path.exists(cache_file):
-        data = load_sweep_data(os.path.join(path, '*.json'))
-        with open(cache_file, 'w') as f:
-            json.dump(data, f)
-
-    with open(cache_file, 'r') as f:
-        data = json.load(f)
-
-    print(f'Loaded {env_name}')
-    return data
-
 def compute_tsne():
     all_data = {}
     normed = {}
 
+    cache = {}
+    cache_file = os.path.join('cache.json')
+    if os.path.exists(cache_file):
+        cache = json.load(open(cache_file, 'r'))
+
     for env in env_names:
-        env_data = cached_sweep_load(f'logs/puffer_{env}', env)
-        sweep_metadata = env_data.pop('sweep')
-        all_data[env] = env_data
+        all_data[env] = cached_load(f'logs/puffer_{env}/*.json', env, cache)
 
+    with open(cache_file, 'w') as f:
+        json.dump(cache, f)
+
+    for env in env_names:
+        env_data = all_data[env]
         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])
-
-            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)
-
-            normed_env.append((dat - mmin) / (mmax - mmin))
+            norm_key = f'{key}_norm'
+            normed_env.append(np.array(env_data[norm_key]))
 
         normed[env] = np.stack(normed_env, axis=1)
 
@@ -192,7 +206,6 @@ def compute_tsne():
     row = 0
     for env in env_names:
         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:
             all_data[env]['tsne1'] = reduced[row:row+sz, 0].tolist()
             all_data[env]['tsne2'] = reduced[row:row+sz, 1].tolist()
@@ -203,7 +216,7 @@ def compute_tsne():
         row += sz
         print(f'Env {env} has {sz} points')
 
-    json.dump(all_data, open('all_cache.json', 'w'))
+    json.dump(all_data, open('pufferlib/ocean/constellation/default.json', 'w'))
 
 if __name__ == '__main__':
     compute_tsne()

pufferlib/__init__.py

@@ -18,11 +18,11 @@
 original_stderr = sys.stderr
 sys.stdout = open(os.devnull, 'w')
 sys.stderr = open(os.devnull, 'w')
-try:
-    import gymnasium
-    import pygame
-except ImportError:
-    pass
+#try:
+#    import gymnasium
+#    import pygame
+#except ImportError:
+#    pass
 sys.stdout.close()
 sys.stderr.close()
 sys.stdout = original_stdout

pufferlib/config/default.ini

@@ -217,13 +217,13 @@ max = 1e-4
 scale = auto
 
 [sweep.train.prio_alpha]
-distribution = logit_normal
-min = 0.1
-max = 0.99
+distribution = uniform 
+min = 0.0
+max = 1.0
 scale = auto
 
 [sweep.train.prio_beta0]
-distribution = logit_normal
-min = 0.1
-max = 0.99
+distribution = uniform 
+min = 0.0
+max = 1.0
 scale = auto

pufferlib/ocean/breakout/binding.c

@@ -3,7 +3,6 @@
 #define NUM_ATNS 1
 #define ACT_SIZES {3}
 #define OBS_TENSOR_T FloatTensor
-#define ACT_TYPE DOUBLE
 
 #define Env Breakout
 #include "vecenv.h"

pufferlib/ocean/breakout/breakout.h

@@ -40,7 +40,7 @@ typedef struct Breakout {
     Client* client;
     Log log;
     float* observations;
-    double* actions;
+    float* actions;
     float* rewards;
     float* terminals;
     int num_agents;
@@ -121,7 +121,7 @@ void init(Breakout* env) {
 void allocate(Breakout* env) {
     init(env);
     env->observations = (float*)calloc(11 + env->num_bricks, sizeof(float));
-    env->actions = (double*)calloc(1, sizeof(double));
+    env->actions = (float*)calloc(1, sizeof(float));
     env->rewards = (float*)calloc(1, sizeof(float));
     env->terminals = (float*)calloc(1, sizeof(float));
 }

pufferlib/ocean/constellation/constellation.c

@@ -889,7 +889,7 @@ int main(void) {
         "train/eps",
         "train/prio_alpha",
         "train/prio_beta0",
-        //"train/horizon",
+        "train/horizon",
         "train/replay_ratio",
         "train/minibatch_size",
         "policy/hidden_size",
@@ -978,14 +978,14 @@ int main(void) {
     int fig_env_idx = 0;
     bool fig_env_active = false;
     bool fig_x_active = false;
-    int fig_x_idx = 0;
+    int fig_x_idx = 1;
     bool fig_xscale_active = false;
     int fig_xscale_idx = 0;
     bool fig_y_active = false;
     int fig_y_idx = 2;
     bool fig_yscale_active = false;
     bool fig_z_active = false;
-    int fig_z_idx = 1;
+    int fig_z_idx = 0;
     bool fig_zscale_active = false;
     int fig_zscale_idx = 0;
     int fig_color_idx = 0;

pufferlib/ocean/torch.py

@@ -72,26 +72,43 @@ def __init__(self, env, hidden_size=512, output_size=512, **kwargs):
         self.multihot_dim = self.factors.sum()
         self.is_continuous = False
 
+        #self.map_2d = nn.Sequential(
+        #    pufferlib.pytorch.layer_init(nn.Conv2d(self.multihot_dim, 128, 5, stride=3)),
+        #    nn.ReLU(),
+        #    pufferlib.pytorch.layer_init(nn.Conv2d(128, 128, 3, stride=1)),
+        #    nn.Flatten(),
+        #)
+
         self.map_2d = nn.Sequential(
-            pufferlib.pytorch.layer_init(nn.Conv2d(self.multihot_dim, 128, 5, stride=3)),
+            nn.Conv2d(self.multihot_dim, 128, 5, stride=3),
             nn.ReLU(),
-            pufferlib.pytorch.layer_init(nn.Conv2d(128, 128, 3, stride=1)),
+            nn.Conv2d(128, 128, 3, stride=1),
             nn.Flatten(),
         )
 
+
         self.player_discrete_encoder = nn.Sequential(
             nn.Embedding(128, 32),
             nn.Flatten(),
         )
+
+        #self.proj = nn.Sequential(
+        #    pufferlib.pytorch.layer_init(nn.Linear(1817, hidden_size)),
+        #    nn.ReLU(),
+        #)
+
         self.proj = nn.Sequential(
-            pufferlib.pytorch.layer_init(nn.Linear(1817, hidden_size)),
+            nn.Linear(1817, hidden_size),
             nn.ReLU(),
         )
 
-        self.layer_norm = nn.LayerNorm(hidden_size)
-        self.actor = pufferlib.pytorch.layer_init(
-            nn.Linear(output_size, self.num_actions), std=0.01)
-        self.value_fn = pufferlib.pytorch.layer_init(nn.Linear(output_size, 1), std=1)
+        #self.layer_norm = nn.LayerNorm(hidden_size)
+        #self.actor = pufferlib.pytorch.layer_init(
+        #    nn.Linear(output_size, self.num_actions), std=0.01)
+        #self.value_fn = pufferlib.pytorch.layer_init(nn.Linear(output_size, 1), std=0.01)
+
+        self.actor = nn.Linear(output_size, self.num_actions)
+        self.value_fn = nn.Linear(output_size, 1)
 
     def forward(self, x, state=None):
         hidden = self.encode_observations(x)
@@ -120,7 +137,7 @@ def encode_observations(self, observations, state=None):
         return obs
 
     def decode_actions(self, flat_hidden):
-        flat_hidden = self.layer_norm(flat_hidden)
+        #flat_hidden = self.layer_norm(flat_hidden)
         action = self.actor(flat_hidden)
         value = self.value_fn(flat_hidden)
         return action, value

pufferlib/src/bindings.cu

@@ -218,7 +218,7 @@ void py_puff_advantage(
     kernel<<<blocks, 256>>>(
         (const precision_t*)values_ptr, (const precision_t*)rewards_ptr,
         (const precision_t*)dones_ptr,  (const precision_t*)importance_ptr,
-        (float*)advantages_ptr,
+        (precision_t*)advantages_ptr,
         gamma, lambda, rho_clip, c_clip, num_steps, horizon);
 }
 
@@ -453,10 +453,6 @@ PYBIND11_MODULE(_C, m) {
         .def("__repr__", [](const PrecisionTensor& t) { return std::string(puf_repr(&t)); })
         .def("ndim", [](const PrecisionTensor& t) { return ndim(t.shape); })
         .def("numel", [](const PrecisionTensor& t) { return numel(t.shape); });
-    py::class_<DoubleTensor>(m, "DoubleTensor")
-        .def("__repr__", [](const DoubleTensor& t) { return std::string(puf_repr(&t)); })
-        .def("ndim", [](const DoubleTensor& t) { return ndim(t.shape); })
-        .def("numel", [](const DoubleTensor& t) { return numel(t.shape); });
     py::class_<FloatTensor>(m, "FloatTensor")
         .def("__repr__", [](const FloatTensor& t) { return std::string(puf_repr(&t)); })
         .def("ndim", [](const FloatTensor& t) { return ndim(t.shape); })