init

.gitignore

@@ -0,0 +1,166 @@
+/misc/
+/.idea/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+/generation_results

README.md

@@ -1,12 +1,32 @@
----
-title: NCERL Diverse PCG
-emoji: 📉
-colorFrom: green
-colorTo: green
-sdk: gradio
-sdk_version: 4.38.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Negatively Correlated Ensemble RL
+
+
+### Verified environment
+* Python 3.9.6
+* JPype 1.3.0
+* dtw 1.4.0
+* scipy 1.7.2
+* torch 1.8.2+cu111
+* numpy 1.20.3
+* gym 0.21.0
+* scipy 1.7.2
+* Pillow 10.0.0
+* matplotlib 3.6.3
+* pandas 1.3.2
+* sklearn 1.0.1
+
+### How to use
+
+All training are launched by running `train.py` with option and arguments. For example, execute `python train.py ncesac --lbd 0.3 --m 5` will train NCERL with hyperparameters set as $\lambda = 0.3, m=5$.
+ Plot script is `plots.py`
+
+* `python train.py gan`: to train a decoder which maps a continuous action to a game level segment.
+* `python train.py sac`: to train a standard SAC as the policy for online game level generation
+* `python train.py asyncsac`: to train a SAC with an asynchronous evaluation environment as the policy for online game level generation
+* `python train.py ncesac`: to train an NCERL based on SAC as the policy for online game level generation
+* `python train.py egsac`: to train an episodic generative SAC (see paper [*The fun facets of Mario: Multifaceted experience-driven PCG via reinforcement learning*](https://dl.acm.org/doi/abs/10.1145/3555858.3563282?casa_token=AHQWYSj_GyoAAAAA:MhwOltqfijP1NQj-c6NaTQikCnlNwyaMky07gCvTK5ZlSq063ew40awAcqEcw6S5zG9Sq9ZyDsspuaM)) as the policy for online game level generation
+* `python train.py pmoe`: to train an episodic generative SAC (see paper [*Probabilistic Mixture-of-Experts for Efficient Deep Reinforcement Learning*](https://arxiv.org/abs/2104.09122)) as the policy for online game level generation
+* `python train.py sunrise`: to train a SUNRISE (see paper [*SUNRISE: A Simple Unified Framework for Ensemble Learning in Deep Reinforcement Learning*](https://proceedings.mlr.press/v139/lee21g.html)) as the policy for online game level generation
+* `python train.py dvd`: to train a DvD-SAC (see paper [*Effective Diversity in Population Based Reinforcement Learning*](https://proceedings.neurips.cc/paper_files/paper/2020/hash/d1dc3a8270a6f9394f88847d7f0050cf-Abstract.html)) as the policy for online game level generation
+
+For the training arguments, please refer to the help `python train.py [option] --help`

analysis/generate.py

@@ -0,0 +1,71 @@
+import os
+import random
+import numpy as np
+from src.gan.gankits import *
+from src.utils.filesys import getpath
+from src.utils.img import make_img_sheet
+from src.utils.datastruct import RingQueue
+from src.olgen.olg_policy import RLGenPolicy, RandGenPolicy
+from src.smb.level import lvlhcat, save_batch
+
+
+def rand_gen_levels(n=100, h=50, dest_path=''):
+    levels = []
+    latvecs = []
+    decoder = get_decoder('models/decoder.pth', 'cuda:0')
+    init_arxv = np.load(getpath('smb/init_latvecs.npy'))
+    for _ in range(n):
+        z0 = init_arxv[random.randrange(0, len(init_arxv))]
+        z0 = torch.tensor(z0, device='cuda:0', dtype=torch.float)
+        z = torch.cat([z0, sample_latvec(h, 'cuda:0')], dim=0)
+        lvl = lvlhcat(process_onehot(decoder(z)))
+        levels.append(lvl)
+        latvecs.append(z.cpu().numpy())
+    if dest_path:
+        save_batch(levels, dest_path)
+        np.save(getpath(dest_path), np.stack(latvecs))
+    return levels, np.stack(latvecs)
+
+def generate_levels(policy, dest_folder='', batch_name='samples.lvls', n=200, h=50, parallel=64, save_img=False):
+    levels = []
+    latvecs = []
+    obs_queues = [RingQueue(policy.n) for _ in range(parallel)]
+    init_arxv = np.load(getpath('smb/init_latvecs.npy'))
+    decoder = get_decoder('models/decoder.pth', 'cuda:0')
+    while len(levels) < n:
+        veclists = [[] for _ in range(parallel)]
+        for queue, veclist in zip(obs_queues, veclists):
+            queue.clear()
+            init_latvec = init_arxv[random.randrange(0, len(init_arxv))]
+            queue.push(init_latvec)
+            veclist.append(init_latvec)
+        for _ in range(h):
+            obs = np.stack([np.concatenate(queue.to_list()) for queue in obs_queues])
+            actions = policy.step(obs)
+            for queue, veclist, action in zip(obs_queues, veclists, actions):
+                queue.push(action)
+                veclist.append(action)
+        for veclist in veclists:
+            latvecs.append(np.stack(veclist))
+            z = torch.tensor(latvecs[-1], device='cuda:0').view(-1, nz, 1, 1)
+            lvl = lvlhcat(process_onehot(decoder(z)))
+            levels.append(lvl)
+        # print(f'{len(levels)}/{n} generated')
+    if dest_folder:
+        os.makedirs(getpath(dest_folder), exist_ok=True)
+        save_batch(levels[:n], getpath(dest_folder, batch_name))
+        if save_img:
+            for i, lvl in enumerate(levels[:n]):
+                lvl.to_img(f'{dest_folder}/lvl-{i}.png')
+    return levels[:n]
+
+
+def make_samples(path, n=12, h=20, space=12):
+    plc = RLGenPolicy.from_path(path)
+    levels = generate_levels(plc, n=n, h=h)
+    imgs = [lvl.to_img() for lvl in levels]
+    make_img_sheet(imgs, ncols=1, y_margin=space, save_path=f'{path}/samples.png')
+    pass
+
+if __name__ == '__main__':
+    pass

analysis/initial_seg.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96086cfef10b8b7993278c96fe34916e08f3566655a5f419d41593db73d93468
+size 40128

analysis/tests.py

@@ -0,0 +1,213 @@
+import os
+import csv
+import time
+import random
+from src.smb.level import *
+from src.drl.me_reg import *
+from src.drl.nets import esmb_sample
+from src.utils.filesys import getpath
+from src.utils.datastruct import RingQueue
+from src.smb.asyncsimlt import AsycSimltPool
+from src.env.environments import get_padded_obs
+from src.olgen.ol_generator import VecOnlineGenerator, OnlineGenerator
+from src.drl.drl_uses import load_cfgs, load_performance
+from src.olgen.olg_policy import process_obs, RandGenPolicy, RLGenPolicy, EnsembleGenPolicy
+
+
+def evaluate_rewards(lvls, rfunc='default', dest_path='', parallel=1, eval_pool=None):
+    internal_pool = eval_pool is None
+    if internal_pool:
+        eval_pool = AsycSimltPool(parallel, rfunc_name=rfunc, verbose=False, test=True)
+    res = []
+    for lvl in lvls:
+        eval_pool.put('evaluate', (0, str(lvl)))
+        buffer = eval_pool.get()
+        for _, item in buffer:
+            res.append([sum(r) for r in zip(*item.values())])
+    if internal_pool:
+        buffer = eval_pool.close()
+    else:
+        buffer = eval_pool.get(True)
+    for _, item in buffer:
+        res.append([sum(r) for r in zip(*item.values())])
+    if len(dest_path):
+        np.save(dest_path, res)
+    return res
+
+def evaluate_mnd(lvls, refs, parallel=2):
+    eval_pool = AsycSimltPool(parallel, verbose=False, refs=[str(ref) for ref in refs])
+    # m, _ = len(lvls), len(refs)
+    res = []
+    for lvl in lvls:
+        eval_pool.put('mnd_item', str(lvl))
+        res += eval_pool.get()
+    res += eval_pool.get(wait=True)
+    res = np.array(res)
+    eval_pool.close()
+    return np.mean(res[:, 0]), np.mean(res[:, 1])
+
+def evaluate_mpd(lvls, parallel=2):
+    task_datas = [[] for _ in range(parallel)]
+    for i, (A, B) in enumerate(combinations(lvls, 2)):
+        # lvlA, lvlB = lvls[i * 2], lvls[i * 2 + 1]
+        task_datas[i % parallel].append((str(A), str(B)))
+
+    hms, dtws = [], []
+    eval_pool = AsycSimltPool(parallel, verbose=False)
+    for task_data in task_datas:
+        eval_pool.put('mpd', task_data)
+    res = eval_pool.get(wait=True)
+    for task_hms, _ in res:
+        hms += task_hms
+        # dtws += task_dtws
+    return np.mean(hms) #, np.mean(dtws)
+
+def evaluate_gen_log(path, parallel=5):
+    rfunc_name = load_cfgs(path, 'rfunc')
+    f = open(getpath(f'{path}/step_tests.csv'), 'w', newline='')
+    wrtr = csv.writer(f)
+    cols = ['step', 'r-avg', 'r-std', 'mnd-hm', 'mnd-dtw', 'mpd-hm', 'mpd-dtw', '']
+    wrtr.writerow(cols)
+    start_time = time.time()
+    for lvls, name in traverse_batched_level_files(f'{path}/gen_log'):
+        step = name[4:]
+        rewards = [sum(item) for item in evaluate_rewards(lvls, rfunc_name, parallel=parallel)]
+        r_avg, r_std = np.mean(rewards), np.std(rewards)
+        # mpd_hm, mpd_dtw = evaluate_mpd(lvls, parallel=parallel)
+        mpd = evaluate_mpd(lvls, parallel=parallel)
+        line = [step, r_avg, r_std, mpd, '']
+        wrtr.writerow(line)
+        f.flush()
+        print(
+            f'{path}: step{step} evaluated in {time.time()-start_time:.1f}s -- '
+            + '; '.join(f'{k}: {v}' for k, v in zip(cols, line))
+        )
+    f.close()
+    pass
+
+def evaluate_generator(generator, nr=200, h=50, parallel=5, dest_path=None, additional_info=None, rfunc_name='default'):
+    if additional_info is None: additional_info = {}
+    ''' Test Reward '''
+    lvls = generator.generate(nr, h)
+    rewards = [sum(item) for item in evaluate_rewards(lvls, parallel=parallel, rfunc=rfunc_name)]
+    r_avg, r_std = np.mean(rewards), np.std(rewards)
+    ''' Test MPD '''
+    # mpd, _ = evaluate_mpd(lvls, parallel=parallel)
+    mpd, *_ = evaluate_mpd(generator.generate(3000*2, h), parallel=parallel)
+    res = {
+        'r-avg': r_avg, 'r-std': r_std, 'div': mpd,
+    }
+    res.update(additional_info)
+    if dest_path:
+        with open(getpath(dest_path), 'w', newline='') as f:
+            keys = [k for k in res.keys()]
+            wrtr = csv.writer(f)
+            wrtr.writerow(keys + [''])
+            wrtr.writerow([res[k] for k in keys] + [''])
+    return res
+    pass
+
+def evaluate_jmer(training_path, n=1000, max_parallel=None, device='cuda:0'):
+    init_vecs = np.load(getpath('smb/init_latvecs.npy'))
+    try:
+        m, histlen, h, gamma, me_type = load_cfgs(training_path, 'm', 'N', 'h', 'gamma', 'me_type')
+    except KeyError:
+        return 0.
+    mereg_func = LogWassersteinExclusion(1.) if me_type == 'logw' else WassersteinExclusion(1.)
+    model = torch.load(getpath(training_path, 'policy.pth'), map_location=device)
+    model.requires_grad_(False)
+    if max_parallel is None:
+        max_parallel = min(n, 512)
+    me_regs = []
+    obs_queues = [RingQueue(histlen) for _ in range(max_parallel)]
+    while len(me_regs) < n:
+        size = min(max_parallel, n - len(me_regs))
+        mereg_vals, discount = np.zeros([size]), 1.
+        veclists = [[] for _ in range(size)]
+        for queue, veclist in zip(obs_queues, veclists):
+            queue.clear()
+            init_latvec = init_vecs[random.randrange(0, len(init_vecs))]
+            queue.push(init_latvec)
+            veclist.append(init_latvec)
+        for _ in range(h):
+            obs = np.stack([get_padded_obs(queue.to_list(), histlen) for queue in obs_queues[:size]])
+            muss, stdss, betas = model.get_intermediate(process_obs(obs, device))
+            mereg_vals += discount * mereg_func.forward(muss, stdss, betas).squeeze().cpu().numpy()
+            discount *= gamma
+            actions, _ = esmb_sample(muss, stdss, betas)
+            for queue, veclist, action in zip(obs_queues, veclists, actions.cpu().numpy()):
+                queue.push(action)
+                veclist.append(action)
+        me_regs += mereg_vals.tolist()
+    return me_regs
+
+def evaluate_baseline(*rfuncs, parallel=4):
+    nr, md, nd, h = 100, 1000, 200, 50
+    gen_policy = RandGenPolicy()
+    olgenerator = OnlineGenerator(gen_policy)
+    lvls, refs = olgenerator.generate(md, h), olgenerator.generate(nd, h)
+    divs_h, divs_js = evaluate_mnd(lvls, refs, parallel=parallel)
+    keys, vals = ['d-h', 'd-js'], [divs_h, divs_js]
+    print(f'Diversity of baseline generator: Hamming {divs_h:.2f}; TPJS {divs_js:.2f}')
+    for rfunc in rfuncs:
+        try:
+            print(f'Start to evaluate {rfunc}')
+            start_time = time.time()
+            lvls = olgenerator.generate(nr, h)
+            rewards = [sum(item) for item in evaluate_rewards(lvls, parallel=parallel, rfunc=rfunc)]
+            keys.append(rfunc)
+            vals.append(np.mean(rewards))
+            print(f'Evaluation for {rfunc} finished in {time.time()-start_time:.2f}s')
+            print(f'Evaluation results for {rfunc}: {vals[-1]:.2f}')
+        except AttributeError:
+            continue
+    with open(getpath('training_data', 'baselines.csv'), 'w', newline='') as f:
+        wrtr = csv.writer(f)
+        wrtr.writerow(keys)
+        wrtr.writerow(vals)
+
+def sample_initial():
+    playable_latvecs = np.load(getpath('smb/init_latvecs.npy'))
+    indexes = random.sample([*range(len(playable_latvecs))], 500)
+    z = playable_latvecs[indexes, :]
+
+    np.save(getpath('analysis/initial_seg.npy'), z)
+    pass
+
+def generate_levels_for_test(h=25):
+    init_set = np.load(getpath('analysis/initial_seg.npy'))
+    def _generte_one(policy, path):
+        try:
+            start = time.time()
+            generator = VecOnlineGenerator(policy, vec_num=len(init_set))
+            fd, _ = os.path.split(getpath(path))
+            os.makedirs(fd, exist_ok=True)
+            generator.re_init(init_set)
+            lvls = generator.generate(len(init_set), h, rand_init=False)
+            save_batch(lvls, path)
+            print('Save to', path, '%.2fs' % (time.time() - start))
+        except FileNotFoundError as e:
+            print(e)
+    for l, m in product(['0.0', '0.1', '0.2', '0.3', '0.4', '0.5'], [2, 3, 4, 5]):
+        for i in range(1, 6):
+            pi_path = f'training_data/varpm-fhp/l{l}_m{m}/t{i}'
+            _generte_one(RLGenPolicy.from_path(pi_path), f'test_data/varpm-fhp/l{l}_m{m}/t{i}/samples.lvls')
+            pi_path = f'training_data/varpm-lgp/l{l}_m{m}/t{i}'
+            _generte_one(RLGenPolicy.from_path(pi_path), f'test_data/varpm-lgp/l{l}_m{m}/t{i}/samples.lvls')
+    for algo in ['sac', 'egsac', 'asyncsac', 'pmoe']:
+        for i in range(1, 6):
+            pi_path = f'training_data/{algo}/fhp/t{i}'
+            _generte_one(RLGenPolicy.from_path(pi_path), f'test_data/{algo}/fhp/t{i}/samples.lvls')
+            pi_path = f'training_data/{algo}/lgp/t{i}'
+            _generte_one(RLGenPolicy.from_path(pi_path), f'test_data/{algo}/lgp/t{i}/samples.lvls')
+    for algo in ['sunrise', 'dvd']:
+        for i in range(1, 5):
+            pi_path = f'training_data/{algo}/fhp/t{i}'
+            _generte_one(EnsembleGenPolicy.from_path(pi_path), f'test_data/{algo}/fhp/t{i}/samples.lvls')
+            pi_path = f'training_data/{algo}/lgp/t{i}'
+            _generte_one(EnsembleGenPolicy.from_path(pi_path), f'test_data/{algo}/lgp/t{i}/samples.lvls')
+        pass
+
+
+if __name__ == '__main__':
+    generate_levels_for_test()

app.py

@@ -0,0 +1,46 @@
+import random
+
+import gradio as gr
+import os
+
+from src.olgen.ol_generator import VecOnlineGenerator
+from src.olgen.olg_game import MarioOnlineGenGame
+from src.olgen.olg_policy import RLGenPolicy
+from src.smb.level import save_batch
+from src.utils.filesys import getpath
+from src.utils.img import make_img_sheet
+
+import torch
+
+device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+
+def generate_and_play():
+    path = 'models/example_policy'
+    # Generate with example policy model
+    N, L = 8, 10
+    plc = RLGenPolicy.from_path(path, device)
+    generator = VecOnlineGenerator(plc, g_device=device)
+    fd, _ = os.path.split(getpath(path))
+    os.makedirs(fd, exist_ok=True)
+
+    lvls = generator.generate(N, L)
+    # save_batch(lvls, f'{path}/samples.lvls')
+    imgs = [lvl.to_img() for lvl in lvls]
+    return imgs
+    # make_img_sheet(imgs, 1, save_path=f'{path}/samples.png')
+
+    # # Play with the example policy model
+    # game = MarioOnlineGenGame(path)
+    # game.play()
+
+
+with gr.Blocks(title="NCERL Demo") as demo:
+    gallery = gr.Gallery(
+        label="Generated images", show_label=False, elem_id="gallery"
+    , columns=[3], rows=[1], object_fit="contain", height="auto")
+    btn = gr.Button("Generate levels", scale=0)
+
+    btn.click(generate_and_play, None, gallery)
+
+if __name__ == "__main__":
+    demo.launch()

generate_and_play.py

@@ -0,0 +1,27 @@
+import os
+
+from src.olgen.ol_generator import VecOnlineGenerator
+from src.olgen.olg_game import MarioOnlineGenGame
+from src.olgen.olg_policy import RLGenPolicy
+from src.smb.level import save_batch
+from src.utils.filesys import getpath
+from src.utils.img import make_img_sheet
+
+if __name__ == '__main__':
+    path = 'models/example_policy'
+    # Generate with example policy model
+    N, L = 8, 10
+    plc = RLGenPolicy.from_path(path)
+    generator = VecOnlineGenerator(plc)
+    fd, _ = os.path.split(getpath(path))
+    os.makedirs(fd, exist_ok=True)
+
+    lvls = generator.generate(N, L)
+    save_batch(lvls, f'{path}/samples.lvls')
+    imgs = [lvl.to_img() for lvl in lvls]
+    make_img_sheet(imgs, 1, save_path=f'{path}/samples.png')
+
+    # # Play with the example policy model
+    # game = MarioOnlineGenGame(path)
+    # game.play()
+    pass

models/decoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:795903ed4957a4fc8b5a349113477643f945efe272d33a276a55671084f10051
+size 1754728

models/example_policy/cfgs.json

@@ -0,0 +1 @@
+{"N": 5, "gamma": 0.9, "h": 50, "rfunc": "lgp"}

models/example_policy/policy.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95bd64a4667f1a55f73897bf1b8e9fff63d0cd2adb860ad799d180c53bc036b8
+size 2430875

models/example_policy/samples.lvls

@@ -0,0 +1,135 @@
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+---------------------------------#----------------------------------------------------------------------------------------------------------------------------------------------
+------------------------------oo-------------------------------o---------------------------------------------------------------o------------------------------------------------
+-------------------------------------------------------------------------------------------------o------------------------------------------------------------------------------
+--------------------------------------SSSSSSSSSS---------------------Q--------------------------------------------------------------QQQ-----------------------------QQQQ--------
+--------------So----------------------------------------------------------------------------------------------------------------------------------------------------------------
+----------------------------------------------o----------------------o---------------------------------------K----------------------------------------------------------------o-
+-----------#---------------------o---------------------------------------------------------------------------2------------------------------------------------------------------
+---------####--------------------------------oS------------------#---SoS-----US------------------------------U-------------------#--SSSS-----US--------tt--------##-###S-----US-
+---------####----------tt-----T------------------------TT----#--TT-----------------------TT------------B---------------TT----#T-TT---------------------tt--------##-------------
+--------########-------Tt-----T------------------------TT----TT-TT-----K----------------TTT------------B---------------TT----TT-TT---------------------tt----TT--##-------------
+-------#########--gggg-Tt---kkT------k-----kk-----gggg-TT---kTT-T#-k-k-g--k-----k-ty----TTT--ggg---k-gog----kkk---gggg-TT---kTT-T--k-k-g--k-k-----g----tt---kkg--##k-k-g--k-k---
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+;
+------------------------------------------------------------------------------------------------------------------------------------------S----------------------------S--S-----
+----------------------------------------------------------------------S-------------------------------------------------------------------------------------------S-SSSSSS------
+-----------------------------------------------------------------------------------------------------------------------------------------%%-------------------------------------
+----------------------------------S------------------------------------------------------------------------------------------------------||-------------------------------------
+----------------------------------------------------------------------------Koo---------------------------------------o------------------||----------------------------o-o-----o
+-----------------------------------------------------------------------------------------------------------------------------------------||-------------------------------------
+----------------S--Q--SSoSS--SSS--o-----------------QQoo--------------SSS----SSS%---SS-----------------------U-------SSSS-------------SSSSSSSS-------------------------SS-------
+S-------------------------------------------------------------------------------|--------------------------------o--------------------------------------------------------------
+SSS-So---------------------------------------S-S--------------------------------|-------------o--------------------------------------------------------------K------------------
+----------------------------------------------S--------o------------------------|--------------------------------S-------------------------------------------2------------oo----
+----------------Q---QS@Q----S@SSS-------------S--------2-----U-----------------S|------------US--------------U-------------------------------SS--------------U-------------%----
+----------T--------------------------------------------tt---------------#-------|------B---------------TT----#-----------------------------------------B-------------------|----
+---------TT--------------------------------K-----------tt--------------##--#-#--|------B---------------TT---TTT---------#-------------t----------------B-------------------|----
+---------TT----#---k------k----------------b---g--gggg-tt---k--------####----#--|-gggggb--k-k-k---g----TT---kT#-------###-------------t------k-g---k-gog----kkk------------|----
+---XXX-XXXXXXXXXXXXXXXXXXXXXXXXX------XXXXXXoXXXXXXXXXXXXXXXXXXXXoXXXXXXXX--oXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX%XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX--%%%%-----|---@
+---XXX-X-XXXXXXXXXXXXXXXXXXXXXXX------XXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX@--XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX|-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX---||------|----
+;
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+------------------------------oo-----------------------------------------------o------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+----------------%S---So--------------QQ--------------------------------------------------------------Q-QQ--------S--QQSSQSSSSSSS-----QQQ----------------------------------------
+----------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------
+----------------|--------------------------------------------K--------------------------------##-------------oo------------------------------oo---------------------------------
+----------------|--------------------------------------------U-------------------o-----o-----###------g-------------------------------------------------------------------------
+----------------|----------------#---SoS-----US--------------U-------------------------tt----###--#-####Q---S@S-#------------US--------------U-------------------------------o--
+----------------|------TT-----K-TT---------------------B------K--------TT----#T--------tt---###------------------------------##--------K---------------TT----TT--------t--------
+----------------|------TT-----U-TT-----K---------------B---------------TT----TT--------tt--####-----------------------------###------------------------TT----TT--------t--------
+--kk-----------g|-gggg-Tt---k-U-T#-k-k----k--------k---t----k-----gggg-TT---kTT-Tg-----tt--####----k------k-k-----ggggg----####----k-kkyk---kkk--ggg-g-TT---TTT------k-tt----#--
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX--XXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+;
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+------------------------------oo-----------------------------------------------------------------------------------------------o------------------------------------------------
+---------------------------------o-----------------------------------------------------------------------------------------------------------------o----------------------------
+----------------%----So--------------Q---------------QQQQ------------Qo--------------QQQQS-----#--------------------------------SooS--SS-----S-----SQQQ---------------o---------
+----------------|------------------------------------------------------------------------------#--------------------------------------------------------------------------------
+----------------|----------------------------------------------------------------------------###-------------K------------------------------------------------o-----------------
+----------------|----------------------------------------------------------------------------###-------------2---------------------------------o--------------------------------
+-------------oo-|----------------#--USoS-----US--#------------------------------Q-Q----QQ----###----------------------------------UQS------------------------US--------------o--
+----------------|------TT-----K-TT---------------##----t---------------t---------------------###-------B---------------TT----#T----------------------------------------tt-------
+---------------@|------TT-----U-TT-----K---------#---------------------t--------------------####-------B---------------TT----TT----------------------------------------tt-------
+---gg----------g|-gggg-Tt---k-U-T--k-k----k------#-k--kk-----k-----k-gog----kkk---or--------####---k-gog----k-k---gggg-TT---kTT--------------------kgggg--k-k-----ggggott---kkk-
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX%%%%%-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX--XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-|XX--XXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+;
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+----------------------------------------------------------------------------------------------------SSSS------------------------------------------------------------------------
+-----------------------------------------------o--------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------QS--------------------------------Q-Q------------------------------------U--SSSSSSSSo----SSS---S@S@QQ-------%---SS------------------------------QQQQ--------
+--------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------
+-----------------------------------------------------------------------------K---------------------------------------o----------|----o--------o-------------------------------o-
+-----------#-----------------------------------------------------------------2------------------g-------------------------------|------------------------------------g----------
+---------TT#-------------------------------------T#--S#S-----US--------------U---------------U--S--SS---------S-----USSS-----US-|----S-2-----US------------------######S-----US-
+---------TT----------------------------TT----TT-TT---------------------B---------------TT----#---------K------------------------|------K--K------------tt----#---##-------------
+--------TTT----T-----------------------TT----TT------------------------B---------------TT----TT---------------------------------|------B---------------Tt----TT--##-------------
+T-------TTT----T-----gg-------kg--gggg-TT---kTT-----------k--------k-gog----k-k---g----TT---kT#-------k-----k------kgggg----k---|--k-gog--k-k-----ggg--Tt----kT--##k---g--k-k---
+XXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+X-XXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+;
+--------------------------------------------------------------------------------------------------------------------------S-----------------------------------------------------
+-----------------------------------------------------------------------------------------------------------------S-SSSSS-SS-----------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+------------------------------------------------------------------------------------------------SSSSSS--------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------oo----------------
+------SSS---SS--------------------------------------SQo------------------------------Q--------------------------@SSSSSSSQSSSSSSS----SQo-----------------------SS----QQQQQSS--S@S
+------------------------------------------------------------------------------------------------------------------------------------------------S--------------#----------------
+-------------------------------------------------------------K-----------------------o-------------------------------------------------------U--S--------------#----------------
+-------------------------------------------------------------------------------------------------------o---------------g---------------------------------------#----------------
+---------S---@S------------------#--USSS-----US-----US-2-----U---------tt-----------USSS-----US--------tt--------SQ-SSSQQ----US--------2-----U----------------##Q-Q--QQQQS---o--
+-----------------------tt----TT-TT---------------------K---------------tt-------T----------------------tt----T-------------------------K--K------------------###----------------
+-----------------------Tt----TT-TT---------------------B---------------tt------------------------------tt----T-------------------------B-------------------#####----------------
+----ggk-----k-----gggg-Tt---kkTTT--k-k----k--------k-gog--k-k-----g----tt---kkg----k---g--k-----------ttt---kkk----k---------------k-gog--k-k--------------#####----------------
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+;
+-------------------------------------------------------k------------------------------------------------------------------------------------------------------------------------
+-------------------------------------S------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+-------SSSSSSSSS----------------------------------------------------------------------------------------------------------------------------------------------------------------
+------------------------------oo---------------------------------------------------------------o-----------o-o----o----------o--------------------------------------------------
+---------------------------------------------o-------------------------------------------------------------------------------------------------o--------------------------------
+----------------%S---So----------------Q--------------SSSSSSSSSS-------------------------------------SQSoSSS-SS----S--SSSS----o-----QQ---------S--------------------------------
+----------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------
+----------------|------------------------------------------------------------U-------------------------------------------------------------------------------K------------------
+----------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------
+----SS----------|----------------T--S##SS----US------------SSS---------------U------------------------QQQ----USS--Qo-----#------QSQ---SSSSSSS%S--------------U---------------o--
+----------------|------TT-----K-TT-----------------------------------------------------TT----TT-------------------------##-------------------|---------B---------------Tt----#--
+----------------|------TT-----U-TT-------------------------------------B---------------TT----TT------------------------###---#---------------|---------B---------------TT----TT-
+----------------|-gggg-Tt---k-U-TT---k----b------------------------k---b----kkk---gggg-TT---kTT----U------------------####--------or-k-------|-----k-gog----kkk---ggg--Tt----kT-
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXX--XX---------X--XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXX--XXXXXXXXXX--XXXXXXXX---XXXXX%XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXX---X---------X--XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXX--X-XXX-XXXX--XXXXXXXX---XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+;
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+----------S-------------------------------------------------------------------------------------------------------------------------------------SSS-----------------------------
+-----------------------------U-o-----------------------------------------------o------------------------------------------------------------------------------------------------
+-------------------------------------------------o---------------------------------------------o--------------------------------------------------------------------------------
+------------------------------------------SS---S%--------------------------------------------------------------------QQ-------------QQoo-----U-------QQQo------------Qo---------
+------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------
+---------#-----T--------------------------------|------------------------------------o-------------------------------------------------------K----------------------------------
+--------#------o--------------------------------|------------------------------------------------------o------------------------------------------------------------------------
+-------TT------T--------------------------------|----So2-----US---------------------USSS-----@S--------tt-----------USoS-----U------U--2-----U--###--------------------------U--
+-------TT----T---------TT----TT--------o--------|------K--K------------TT----TT-T----------------------tt----#---------K---------------K--------###--------------------B--------
+------#TT---TT---------TT----TT--------#t-------|------B---------------TT----TT------------------------tt----T---------B---------------B--------###--------------------B--------
+-----##TT---TT----gg-#-Tt---kkT---gggg-TT--#y---|-gk-gog--k-k----ggggg-TT---kTT----k------k-------g----tt---kkT-T--k-kkb--k-k------k-gob--k-k---##----k-----------gggggb----k-k-
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

plots.py

@@ -0,0 +1,733 @@
+import glob
+import json
+import os
+import re
+
+import numpy as np
+import pandas as pds
+import matplotlib
+import matplotlib.pyplot as plt
+from math import sqrt
+import torch
+from root import PRJROOT
+from sklearn.manifold import TSNE
+from itertools import product, chain
+# from src.drl.drl_uses import load_cfgs
+from src.gan.gankits import get_decoder, process_onehot
+from src.gan.gans import nz
+from src.smb.level import load_batch, hamming_dis, lvlhcat
+from src.utils.datastruct import RingQueue
+from src.utils.filesys import load_dict_json, getpath
+from src.utils.img import make_img_sheet
+from torch.distributions import Normal
+
+matplotlib.rcParams["axes.formatter.limits"] = (-5, 5)
+
+
+def print_compare_tab():
+    rand_lgp, rand_fhp, rand_divs = load_dict_json(
+        'test_data/rand_policy/performance.csv', 'lgp', 'fhp', 'diversity'
+    )
+    rand_performance = {'lgp': rand_lgp, 'fhp': rand_fhp, 'diversity': rand_divs}
+
+    def _print_line(_data, minimise=False):
+        means = _data.mean(axis=-1)
+        stds = _data.std(axis=-1)
+        max_i, min_i = np.argmax(means), np.argmin(means)
+        mean_str_content = [*map(lambda x: '%.4g' % x, _data.mean(axis=-1))]
+        std_str_content = [*map(lambda x: '$\pm$%.3g' % x, _data.std(axis=-1))]
+        if minimise:
+            mean_str_content[min_i] = r'\textbf{%s}' % mean_str_content[min_i]
+            mean_str_content[max_i] = r'\textit{%s}' % mean_str_content[max_i]
+            std_str_content[min_i] = r'\textbf{%s}' % std_str_content[min_i]
+            std_str_content[max_i] = r'\textit{%s}' % std_str_content[max_i]
+        else:
+            mean_str_content[max_i] = r'\textbf{%s}' % mean_str_content[max_i]
+            mean_str_content[min_i] = r'\textit{%s}' % mean_str_content[min_i]
+            std_str_content[max_i] = r'\textbf{%s}' % std_str_content[max_i]
+            std_str_content[min_i] = r'\textit{%s}' % std_str_content[min_i]
+        print('    &', ' & '.join(mean_str_content), r'\\')
+        print('    & &', ' & '.join(std_str_content), r'\\')
+        pass
+
+    def _print_block(_task):
+        fds = [
+            f'sac/{_task}', f'egsac/{_task}', f'asyncsac/{_task}',
+            f'pmoe/{_task}', f'dvd/{_task}', f'sunrise/{_task}',
+            f'varpm-{_task}/l0.0_m5', f'varpm-{_task}/l0.1_m5', f'varpm-{_task}/l0.2_m5',
+            f'varpm-{_task}/l0.3_m5', f'varpm-{_task}/l0.4_m5', f'varpm-{_task}/l0.5_m5'
+        ]
+        rewards, divs = [], []
+        for fd in fds:
+            rewards.append([])
+            divs.append([])
+            # print(getpath())
+            for path in glob.glob(getpath('test_data', fd, '**', 'performance.csv'), recursive=True):
+                reward, div = load_dict_json(path, 'reward', 'diversity')
+                rewards[-1].append(reward)
+                divs[-1].append(div)
+        rewards = np.array(rewards)
+        divs = np.array(divs)
+
+        print('    & \\multirow{2}{*}{Reward}')
+        _print_line(rewards)
+        print('    \\cline{2-14}')
+        print('    & \\multirow{2}{*}{Diversity}')
+        _print_line(divs)
+        print('    \\cline{2-14}')
+        print('    & \\multirow{2}{*}{G-mean}')
+        gmean = np.sqrt(rewards * divs)
+        _print_line(gmean)
+
+        print('    \\cline{2-14}')
+        print('    & \\multirow{2}{*}{N-rank}')
+        r_rank = np.zeros_like(rewards.flatten())
+        r_rank[np.argsort(-rewards.flatten())] = np.linspace(1, len(r_rank), len(r_rank))
+
+        d_rank = np.zeros_like(divs.flatten())
+        d_rank[np.argsort(-divs.flatten())] = np.linspace(1, len(r_rank), len(r_rank))
+        n_rank = (r_rank.reshape([12, 5]) + d_rank.reshape([12, 5])) / (2 * 5)
+        _print_line(n_rank, True)
+
+    print('    \\multirow{8}{*}{MarioPuzzle}')
+    _print_block('fhp')
+    print('    \\midrule')
+    print('    \\multirow{8}{*}{MultiFacet}')
+    _print_block('lgp')
+    pass
+
+def print_compare_tab_nonrl():
+    # rand_lgp, rand_fhp, rand_divs = load_dict_json(
+    #     'test_data/rand_policy/performance.csv', 'lgp', 'fhp', 'diversity'
+    # )
+    # rand_performance = {'lgp': rand_lgp, 'fhp': rand_fhp, 'diversity': rand_divs}
+
+    def _print_line(_data, minimise=False):
+        means = _data.mean(axis=-1)
+        stds = _data.std(axis=-1)
+        max_i, min_i = np.argmax(means), np.argmin(means)
+        mean_str_content = [*map(lambda x: '%.4g' % x, _data.mean(axis=-1))]
+        std_str_content = [*map(lambda x: '$\pm$%.3g' % x, _data.std(axis=-1))]
+        if minimise:
+            mean_str_content[min_i] = r'\textbf{%s}' % mean_str_content[min_i]
+            mean_str_content[max_i] = r'\textit{%s}' % mean_str_content[max_i]
+            std_str_content[min_i] = r'\textbf{%s}' % std_str_content[min_i]
+            std_str_content[max_i] = r'\textit{%s}' % std_str_content[max_i]
+        else:
+            mean_str_content[max_i] = r'\textbf{%s}' % mean_str_content[max_i]
+            mean_str_content[min_i] = r'\textit{%s}' % mean_str_content[min_i]
+            std_str_content[max_i] = r'\textbf{%s}' % std_str_content[max_i]
+            std_str_content[min_i] = r'\textit{%s}' % std_str_content[min_i]
+        print('    &', ' & '.join(mean_str_content), r'\\')
+        print('    & &', ' & '.join(std_str_content), r'\\')
+        pass
+
+    def _print_block(_task):
+        fds = [
+            f'GAN-{_task}', f'DDPM-{_task}',
+            f'varpm-{_task}/l0.0_m5', f'varpm-{_task}/l0.1_m5', f'varpm-{_task}/l0.2_m5',
+            f'varpm-{_task}/l0.3_m5', f'varpm-{_task}/l0.4_m5', f'varpm-{_task}/l0.5_m5'
+        ]
+        rewards, divs = [], []
+        for fd in fds:
+            rewards.append([])
+            divs.append([])
+            # print(getpath())
+            for path in glob.glob(getpath('test_data', fd, '**', 'performance.csv'), recursive=True):
+                reward, div = load_dict_json(path, 'reward', 'diversity')
+                rewards[-1].append(reward)
+                divs[-1].append(div)
+        rewards = np.array(rewards)
+        divs = np.array(divs)
+
+        print('    & \\multirow{2}{*}{Reward}')
+        _print_line(rewards)
+        print('    \\cline{2-10}')
+        print('    & \\multirow{2}{*}{Diversity}')
+        _print_line(divs)
+        print('    \\cline{2-10}')
+        # print('    & \\multirow{2}{*}{G-mean}')
+        # gmean = np.sqrt(rewards * divs)
+        # _print_line(gmean)
+        #
+        # print('    \\cline{2-10}')
+        # print('    & \\multirow{2}{*}{N-rank}')
+        # r_rank = np.zeros_like(rewards.flatten())
+        # r_rank[np.argsort(-rewards.flatten())] = np.linspace(1, len(r_rank), len(r_rank))
+        #
+        # d_rank = np.zeros_like(divs.flatten())
+        # d_rank[np.argsort(-divs.flatten())] = np.linspace(1, len(r_rank), len(r_rank))
+        # n_rank = (r_rank.reshape([8, 5]) + d_rank.reshape([8, 5])) / (2 * 5)
+        # _print_line(n_rank, True)
+
+    print('    \\multirow{4}{*}{MarioPuzzle}')
+    _print_block('fhp')
+    print('    \\midrule')
+    print('    \\multirow{4}{*}{MultiFacet}')
+    _print_block('lgp')
+    pass
+
+def plot_cmp_learning_curves(task, save_path='', title=''):
+    plt.style.use('seaborn')
+    colors = [plt.plot([0, 1], [-1000, -1000])[0].get_color() for _ in range(6)]
+    plt.cla()
+    plt.style.use('default')
+
+    # colors = ('#5D2CAB', '#005BD4', '#007CE4', '#0097DD', '#00ADC4', '#00C1A5')
+    def _get_algo_data(fd):
+        res = []
+        for i in range(1, 6):
+            path = getpath(fd, f't{i}', 'step_tests.csv')
+            try:
+                data = pds.read_csv(path)
+                trajectory = [
+                    [float(item['step']), float(item['r-avg']), float(item['diversity'])]
+                    for _, item in data.iterrows()
+                ]
+                trajectory.sort(key=lambda x: x[0])
+                res.append(trajectory)
+                if len(trajectory) != 26:
+                    print('Not complete (%d)/26:' % len(trajectory), path)
+            except FileNotFoundError:
+                print(path)
+        res = np.array(res)
+        # rdsum = res[:, :, 1] + res[:, :, 2]
+        gmean = np.sqrt(res[:, :, 1] * res[:, :, 2])
+        steps = res[0, :, 0]
+        # r_avgs = np.mean(res[:, :, 1], axis=0)
+        # r_stds = np.std(res[:, :, 1], axis=0)
+        # divs = np.mean(res[:, :, 2], axis=0)
+        # div_std = np.std(res[:, :, 2], axis=0)
+        _performances = {
+            'reward': (np.mean(res[:, :, 1], axis=0), np.std(res[:, :, 1], axis=0)),
+            'diversity': (np.mean(res[:, :, 2], axis=0), np.std(res[:, :, 2], axis=0)),
+            # 'rdsum': (np.mean(rdsum, axis=0), np.std(rdsum, axis=0)),
+            'gmean': (np.mean(gmean, axis=0), np.std(gmean, axis=0)),
+        }
+        # print(_performances['gmean'])
+        return steps, _performances
+
+    def _plot_criterion(_ax, _criterion):
+        i, j, k = 0, 0, 0
+        for algo, (steps, _performances) in performances.items():
+            avgs, stds = _performances[_criterion]
+            if '\lambda' in algo:
+                ls = '-'
+                _c = colors[i]
+                i += 1
+            elif algo in {'SAC', 'EGSAC', 'ASAC'}:
+                ls = ':'
+                _c = colors[j]
+                j += 1
+            else:
+                ls = '--'
+                _c = colors[j]
+                j += 1
+            _ax.plot(steps, avgs, color=_c, label=algo, ls=ls)
+            _ax.fill_between(steps, avgs - stds, avgs + stds, color=_c, alpha=0.15)
+            _ax.grid(False)
+            # plt.plot(steps, avgs, label=algo)
+            # plt.plot(_performances, label=algo)
+        pass
+        _ax.set_xlabel('Time step')
+
+    fig, ax = plt.subplots(1, 3, figsize=(9.6, 3.2), dpi=250, width_ratios=[1, 1, 1])
+    # fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(5, 4), dpi=256)
+    # fig, ax1 = plt.subplots(1, 1, figsize=(8, 3), dpi=256)
+    # ax2 = ax1.twinx()
+    # fig = plt.plot(figsize=(4, 3), dpi=256)
+    performances = {
+        'SUNRISE': _get_algo_data(f'test_data/sunrise/{task}'),
+        '$\lambda$=0.0': _get_algo_data(f'test_data/varpm-{task}/l0.0_m5'),
+        'DvD': _get_algo_data(f'test_data/dvd/{task}'),
+        '$\lambda$=0.1': _get_algo_data(f'test_data/varpm-{task}/l0.1_m5'),
+        'PMOE': _get_algo_data(f'test_data/pmoe/{task}'),
+        '$\lambda$=0.2': _get_algo_data(f'test_data/varpm-{task}/l0.2_m5'),
+        'SAC': _get_algo_data(f'test_data/sac/{task}'),
+        '$\lambda$=0.3': _get_algo_data(f'test_data/varpm-{task}/l0.3_m5'),
+        'EGSAC': _get_algo_data(f'test_data/egsac/{task}'),
+        '$\lambda$=0.4': _get_algo_data(f'test_data/varpm-{task}/l0.4_m5'),
+        'ASAC': _get_algo_data(f'test_data/asyncsac/{task}'),
+        '$\lambda$=0.5': _get_algo_data(f'test_data/varpm-{task}/l0.5_m5'),
+    }
+    # _plot_algo(*_get_algo_data(glob.glob(getpath('test_data/SAC', '**', 'step_tests.csv'))), 'SAC')
+    # _plot_algo(*_get_algo_data(glob.glob(getpath('test_data/EGSAC', '**', 'step_tests.csv'))), 'EGSAC')
+    # _plot_algo(*_get_algo_data(glob.glob(getpath('test_data/AsyncSAC', '**', 'step_tests.csv'))), 'AsyncSAC')
+    # _plot_algo(*_get_algo_data(glob.glob(getpath('test_data/SUNRISE', '**', 'step_tests.csv'))), 'SUNRISE')
+    # _plot_algo(*_get_algo_data(glob.glob(getpath('test_data/DvD-ES', '**', 'step_tests.csv'))), 'DvD-ES')
+    # _plot_algo(*_get_algo_data(glob.glob(getpath('test_data/lbd-m-crosstest/l0.04_m5', '**', 'step_tests.csv'))), 'NCESAC')
+
+
+
+    _plot_criterion(ax[0], 'reward')
+    _plot_criterion(ax[1], 'diversity')
+    # _plot_criterion(ax[2], 'rdsum')
+    _plot_criterion(ax[2], 'gmean')
+    # ax[0].set_title(f'{title} reward')
+    ax[0].set_title(f'Cumulative Reward')
+    ax[1].set_title('Diversity Score')
+    # ax[2].set_title('Summation')
+    ax[2].set_title('G-mean')
+    # plt.title(title)
+
+    lines, labels = fig.axes[-1].get_legend_handles_labels()
+    fig.suptitle(title, fontsize=14)
+    plt.tight_layout(pad=0.5)
+    if save_path:
+        plt.savefig(getpath(save_path))
+    else:
+        plt.show()
+
+    plt.cla()
+    plt.figure(figsize=(9.6, 2.4), dpi=250)
+    plt.grid(False)
+    plt.axis('off')
+    plt.yticks([1.0])
+    plt.legend(
+        lines, labels, loc='lower center', ncol=6, edgecolor='white', fontsize=15,
+        columnspacing=0.8, borderpad=0.16, labelspacing=0.2, handlelength=2.4, handletextpad=0.3
+    )
+    plt.tight_layout(pad=0.5)
+    plt.show()
+    pass
+
+def plot_crosstest_scatters(rfunc, xrange=None, yrange=None, title=''):
+    def get_pareto():
+        all_points = list(chain(*scatter_groups.values())) + cmp_points
+        res = []
+        for p in all_points:
+            non_dominated = True
+            for q in all_points:
+                if q[0] >= p[0] and q[1] >= p[1] and (q[0] > p[0] or q[1] > p[1]):
+                    non_dominated = False
+                    break
+            if non_dominated:
+                res.append(p)
+        res.sort(key=lambda item:item[0])
+        return np.array(res)
+    def _hex_color(_c):
+        return
+    scatter_groups = {}
+    all_lbd = set()
+    # Initialise
+    plt.style.use('seaborn-v0_8-dark-palette')
+    # plt.figure(figsize=(4, 4), dpi=256)
+    plt.figure(figsize=(2.5, 2.5), dpi=256)
+    plt.axes().set_axisbelow(True)
+
+    # Competitors' performances
+    cmp_folders = ['asyncsac', 'egsac', 'sac', 'sunrise', 'dvd', 'pmoe']
+    cmp_names = ['ASAC', 'EGSAC', 'SAC', 'SUNRISE', 'DvD', 'PMOE']
+    cmp_labels = ['A', 'E', 'S', 'R', 'D', 'M']
+    cmp_markers = ['2', 'x', '+', 'o', '*', 'D']
+    cmp_sizes = [42, 20, 32, 16, 24, 10, 10]
+    cmp_points = []
+    for name, folder, label, mk, s in zip(cmp_names, cmp_folders, cmp_labels, cmp_markers, cmp_sizes):
+        path_fmt = getpath('test_data', folder, rfunc, '*', 'performance.csv')
+        # print(path_fmt)
+        xs, ys = [], []
+        for path in glob.glob(path_fmt, recursive=True):
+            # print(path)
+            try:
+                x, y = load_dict_json(path, 'reward', 'diversity')
+                xs.append(x)
+                ys.append(y)
+                cmp_points.append([x, y])
+                # plt.text(x, y, label, size=7, weight='bold', va='center', ha='center', color='#202020')
+            except FileNotFoundError:
+                print(path)
+        if label in {'A', 'E', 'S'}:
+            plt.scatter(xs, ys, marker=mk, zorder=2, s=s, label=name, color='#202020')
+        else:
+            plt.scatter(
+                xs, ys, marker=mk, zorder=2, s=s, label=name, color=[0., 0., 0., 0.],
+                edgecolors='#202020', linewidths=1
+            )
+    # NCESAC performances
+    for path in glob.glob(getpath('test_data', f'varpm-{rfunc}', '**', 'performance.csv'), recursive=True):
+        try:
+            x, y = load_dict_json(path, 'reward', 'diversity')
+            key = path.split('\\')[-3]
+            _, mtxt = key.split('_')
+            ltxt, _ = key.split('_')
+            lbd = float(ltxt[1:])
+            # if mtxt in {'m2', 'm3', 'm4'}:
+            #     continue
+            all_lbd.add(lbd)
+            if key not in scatter_groups.keys():
+                scatter_groups[key] = []
+            scatter_groups[key].append([x, y])
+        except Exception as e:
+            print(path)
+            print(e)
+
+    palette = plt.get_cmap('seismic')
+    color_x = [0.2, 0.33, 0.4, 0.61, 0.67, 0.79]
+    colors = {lbd: matplotlib.colors.to_hex(c) for c, lbd in zip(palette(color_x), sorted(all_lbd))}
+    colors = {0.0: '#150080', 0.1: '#066598', 0.2: '#01E499', 0.3: '#9FD40C', 0.4: '#F3B020', 0.5: '#FA0000'}
+    for lbd in sorted(all_lbd): plt.plot([-20], [-20], label=f'$\\lambda={lbd:.1f}$', lw=6, c=colors[lbd])
+    markers = {2: 'o', 3: '^', 4: 'D', 5: 'p', 6: 'h'}
+    msizes = {2: 25, 3: 25, 4: 16, 5: 28, 6: 32}
+    for key, group in scatter_groups.items():
+        ltxt, mtxt = key.split('_')
+        l = float(ltxt[1:])
+        m = int(mtxt[1:])
+        arr = np.array(group)
+        plt.scatter(
+            arr[:, 0], arr[:, 1], marker=markers[m], s=msizes[m], color=[0., 0., 0., 0.], zorder=2,
+            edgecolors=colors[l], linewidths=1
+        )
+
+    plt.xlim(xrange)
+    plt.ylim(yrange)
+    # plt.xlabel('Task Reward')
+    # plt.ylabel('Diversity')
+    # plt.legend(ncol=2)
+    # plt.legend(
+    #     ncol=2, loc='lower left', columnspacing=1.2, borderpad=0.0,
+    #     handlelength=1, handletextpad=0.5, framealpha=0.
+    # )
+    pareto = get_pareto()
+    plt.plot(
+        pareto[:, 0], pareto[:, 1], color='black', alpha=0.18, lw=6, zorder=3,
+        solid_joinstyle='round', solid_capstyle='round'
+    )
+    # plt.plot([88, 98, 98, 88, 88], [35, 35, 0.2, 0.2, 35], color='black', alpha=0.3, lw=1.5)
+    # plt.xticks(fontsize=16)
+    # plt.yticks(fontsize=16)
+    # plt.xticks([(1+space) * (m-mlow) + 0.5 for m in ms], [f'm={m}' for m in ms])
+    plt.title(title)
+    plt.grid()
+    plt.tight_layout(pad=0.4)
+    plt.show()
+
+def plot_varpm_heat(task, name):
+    def _get_score(m, l):
+        fd = getpath('test_data', f'varpm-{task}', f'l{l}_m{m}')
+        rewards, divs = [], []
+        for i in range(5):
+            reward, div = load_dict_json(f'{fd}/t{i+1}/performance.csv', 'reward', 'diversity')
+            rewards.append(reward)
+            divs.append(div)
+        gmean = [sqrt(r * d) for r, d in zip(rewards, divs)]
+        return np.mean(rewards), np.std(rewards), \
+            np.mean(divs), np.std(divs), \
+            np.mean(gmean), np.std(gmean)
+
+    def _plot_map(avg_map, std_map, criterion):
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3), dpi=256, width_ratios=(1, 1))
+        heat1 = ax1.imshow(avg_map, cmap='spring')
+        heat2 = ax2.imshow(std_map, cmap='spring')
+        ax1.set_xlim([-0.5, 5.5])
+        ax1.set_xticks([0, 1, 2, 3, 4, 5], ['$\lambda$=0.0', '$\lambda$=0.1', '$\lambda$=0.2', '$\lambda$=0.3', '$\lambda$=0.4', '$\lambda$=0.5'])
+        ax1.set_ylim([-0.5, 3.5])
+        ax1.set_yticks([0, 1, 2, 3], ['m=5', 'm=4', 'm=3', 'm=2'])
+        ax1.set_title('Average')
+        for x, y in product([0, 1, 2, 3, 4, 5], [0, 1, 2, 3]):
+            v = avg_map[y, x]
+            s = '%.4f' % v
+            if v >= 1000: s = s[:4]
+            elif v >= 1: s = s[:5]
+            else: s = s[1:6]
+            ax1.text(x, y, s, va='center', ha='center')
+        plt.colorbar(heat1, ax=ax1, shrink=0.9)
+        ax2.set_xlim([-0.5, 5.5])
+        ax2.set_xticks([0, 1, 2, 3, 4, 5], ['$\lambda$=0.0', '$\lambda$=0.1', '$\lambda$=0.2', '$\lambda$=0.3', '$\lambda$=0.4', '$\lambda$=0.5'])
+        ax2.set_ylim([-0.5, 3.5])
+        ax2.set_yticks([0, 1, 2, 3], ['m=5', 'm=4', 'm=3', 'm=2'])
+        for x, y in product([0, 1, 2, 3, 4, 5], [0, 1, 2, 3]):
+            v = std_map[y, x]
+            s = '%.4f' % v
+            if v >= 1000: s = s[:4]
+            elif v >= 1: s = s[:5]
+            else: s = s[1:6]
+            ax2.text(x, y, s, va='center', ha='center')
+        ax2.set_title('Standard Deviation')
+        plt.colorbar(heat2, ax=ax2, shrink=0.9)
+
+        fig.suptitle(f'{name}: {criterion}', fontsize=14)
+        plt.tight_layout()
+        # plt.show()
+        plt.savefig(getpath(f'results/heat/{name}-{criterion}.png'))
+
+    r_mean_map, r_std_map, d_mean_map, d_std_map, g_mean_map, g_std_map \
+        = (np.zeros([4, 6], dtype=float) for _ in range(6))
+    ms = [2, 3, 4, 5]
+    ls = ['0.0', '0.1', '0.2', '0.3', '0.4', '0.5']
+    for i, j in product(range(4), range(6)):
+        r_mean, r_std, d_mean, d_std, g_mean, g_std = _get_score(ms[i], ls[j])
+        r_mean_map[i, j] = r_mean
+        r_std_map[i, j] = r_std
+        d_mean_map[i, j] = d_mean
+        d_std_map[i, j] = d_std
+        g_mean_map[i, j] = g_mean
+        g_std_map[i, j] = g_std
+
+    _plot_map(r_mean_map, r_std_map, 'Reward')
+    _plot_map(d_mean_map, d_std_map, 'Diversity')
+    _plot_map(g_mean_map, g_std_map,'G-mean')
+    # _plot_map(g_mean_map, g_std_map,'G-mean')
+
+def vis_samples():
+    # for l, m in product(['0.0', '0.1', '0.2', '0.3', '0.4', '0.5'], [2, 3, 4, 5]):
+    #     for i in range(1, 6):
+    #         lvls = load_batch(f'{PRJROOT}/test_data/varpm-fhp/l{l}_m{m}/t{i}/samples.lvls')
+    #         imgs = [lvl.to_img(save_path=None) for lvl in lvls[:10]]
+    #         make_img_sheet(imgs, 1, save_path=f'{PRJROOT}/test_data/varpm-fhp/l{l}_m{m}/t{i}/samples.png')
+    # for algo in ['sac', 'egsac', 'asyncsac', 'dvd', 'sunrise', 'pmoe']:
+    #     for i in range(1, 6):
+    #         lvls = load_batch(f'{PRJROOT}/test_data/{algo}/fhp/t{i}/samples.lvls')
+    #         imgs = [lvl.to_img(save_path=None) for lvl in lvls[:10]]
+    #         make_img_sheet(imgs, 1, save_path=f'{PRJROOT}/test_data/{algo}/fhp/t{i}/samples.png')
+    for i in range(1, 6):
+        lvls = load_batch(f'{PRJROOT}/test_data/DDPM-fhp/t{i}/samples.lvls')
+        imgs = [lvl.to_img(save_path=None) for lvl in lvls[:10]]
+        make_img_sheet(imgs, 1, save_path=f'{PRJROOT}/test_data/DDPM-fhp/t{i}/samples.png')
+        pass
+    pass
+
+def make_tsne(task, title, n=500, save_path=None):
+    if not os.path.exists(getpath('test_data', f'samples_dist-{task}_{n}.npy')):
+        samples = []
+        for algo in ['dvd', 'egsac', 'pmoe', 'sunrise', 'asyncsac', 'sac']:
+            for t in range(5):
+                lvls = load_batch(getpath('test_data', algo, task, f't{t+1}', 'samples.lvls'))
+                samples += lvls[:n]
+        for l in ['0.0', '0.1', '0.2', '0.3', '0.4', '0.5']:
+            for t in range(5):
+                lvls = load_batch(getpath('test_data', f'varpm-{task}', f'l{l}_m5', f't{t+1}', 'samples.lvls'))
+                samples += lvls[:n]
+        distmat = []
+        for a in samples:
+            dist_list = []
+            for b in samples:
+                dist_list.append(hamming_dis(a, b))
+            distmat.append(dist_list)
+        distmat = np.array(distmat)
+        np.save(getpath('test_data', f'samples_dist-{task}_{n}.npy'), distmat)
+
+    labels = (
+        '$\lambda$=0.0', '$\lambda$=0.1', '$\lambda$=0.2', '$\lambda$=0.3', '$\lambda$=0.4',
+        '$\lambda$=0.5', 'DvD', 'EGSAC', 'PMOE', 'SUNRISE', 'ASAC', 'SAC'
+    )
+    tsne = TSNE(learning_rate='auto', n_components=2, metric='precomputed')
+    print(np.load(getpath('test_data', f'samples_dist-{task}_{n}.npy')).shape)
+    data = np.load(getpath('test_data', f'samples_dist-{task}_{n}.npy'))
+    embx = np.array(tsne.fit_transform(data))
+
+    plt.style.use('seaborn-dark-palette')
+    plt.figure(figsize=(5, 5), dpi=384)
+    colors = [plt.plot([-1000, -1100], [0, 0])[0].get_color() for _ in range(6)]
+    for i in range(6):
+        x, y = embx[i*n*5:(i+1)*n*5, 0], embx[i*n*5:(i+1)*n*5, 1]
+        plt.scatter(x, y, s=10, label=labels[i], marker='x', c=colors[i])
+    for i in range(6, 12):
+        x, y = embx[i*n*5:(i+1)*n*5, 0], embx[i*n*5:(i+1)*n*5, 1]
+        plt.scatter(x, y, s=8, linewidths=0, label=labels[i], c=colors[i-6])
+    # plt.scatter(embx[100:200, 0], embx[100:200, 1], c=colors[1], s=12, linewidths=0, label='Killer')
+    # plt.scatter(embx[200:, 0], embx[200:, 1], c=colors[2], s=12, linewidths=0, label='Collector')
+    # for i in range(4):
+    #     plt.text(embx[i+100, 0], embx[i+100, 1], str(i+1))
+    #     plt.text(embx[i+200, 0], embx[i+200, 1], str(i+1))
+    #     pass
+    # for emb, lb, c in zip(embs, labels,colors):
+    #     plt.scatter(emb[:,0], emb[:,1], c=c, label=lb, alpha=0.15, linewidths=0, s=7)
+
+    # xspan = 1.08 * max(abs(embx[:, 0].max()), abs(embx[:, 0].min()))
+    # yspan = 1.08 * max(abs(embx[:, 1].max()), abs(embx[:, 1].min()))
+
+    xrange = [1.05 * embx[:, 0].min(), 1.05 * embx[:, 0].max()]
+    yrange = [1.05 * embx[:, 1].min(), 1.25 * embx[:, 1].max()]
+
+    plt.xlim(xrange)
+    plt.ylim(yrange)
+    plt.xticks([])
+    plt.yticks([])
+    # plt.legend(ncol=6, handletextpad=0.02, labelspacing=0.05, columnspacing=0.16)
+    # plt.xticks([-xspan, -0.5 * xspan, 0, 0.5 * xspan, xspan], [''] * 5)
+    # plt.yticks([-yspan, -0.5 * yspan, 0, 0.6 * yspan, yspan], [''] * 5)
+    plt.title(title)
+    plt.legend(loc='upper center', ncol=6, fontsize=9, handlelength=.5, handletextpad=0.5, columnspacing=0.3, framealpha=0.)
+    plt.tight_layout(pad=0.2)
+
+    if save_path:
+        plt.savefig(getpath(save_path))
+    else:
+        plt.show()
+
+def _prob_fmt(p, digitals=3, threshold=0.001):
+    fmt = '%.' + str(digitals) + 'f'
+    if p < threshold:
+        return '$\\approx 0$'
+    else:
+        txt = '$%s$' % ((fmt % p)[1:])
+        if txt == '$.000$':
+            txt = '$1.00$'
+        return txt
+
+def _g_fmt(v, digitals=4):
+    fmt = '%.' + str(digitals) + 'g'
+    txt = (fmt % v)
+    lack = digitals - len(txt.replace('-', '').replace('.', ''))
+    if lack > 0 and '.' not in txt:
+        txt += '.'
+    return txt + '0' * lack
+    pass
+
+def print_selection_prob(path, h=15, runs=2):
+    s0 = 0
+    model = torch.load(getpath(f'{path}/policy.pth'), map_location='cpu')
+    model.requires_grad_(False)
+    model.to('cpu')
+    n = 11
+    # n = load_cfgs(path, 'N')
+    # print(model.m)
+
+    init_vec = np.load(getpath('analysis/initial_seg.npy'))[s0]
+    decoder = get_decoder(device='cpu')
+    obs_buffer = RingQueue(n)
+    for r in range(runs):
+        for _ in range(h): obs_buffer.push(np.zeros([nz]))
+        obs_buffer.push(init_vec)
+        level_latvecs = [init_vec]
+        probs = np.zeros([model.m, h])
+        # probs = []
+        selects = []
+        for t in range(h):
+            # probs.append([])
+            obs = torch.tensor(np.concatenate(obs_buffer.to_list(), axis=-1), dtype=torch.float).view([1, -1])
+            muss, stdss, betas = model.get_intermediate(torch.tensor(obs))
+            i = torch.multinomial(betas.squeeze(), 1).item()
+            # print(i)
+            mu, std = muss[0][i], stdss[0][i]
+            action = Normal(mu, std).rsample([1]).squeeze().numpy()
+            # print(action)
+            # print(mu)
+            # print(std)
+            # print(action.numpy())
+            obs_buffer.push(action)
+            level_latvecs.append(action)
+            # i = torch.multinomial(betas.squeeze(), 1).item()
+            # print(i)
+            probs[:, t] = betas.squeeze().numpy()
+            selects.append(i)
+            pass
+        onehots = decoder(torch.tensor(level_latvecs).view(-1, nz, 1, 1))
+        segs = process_onehot(onehots)
+        lvl = lvlhcat(segs)
+        lvl.to_img(f'figures/gen_process/run{r}-01.png')
+        txts = [[_prob_fmt(p) for p in row] for row in probs]
+        for t, i in enumerate(selects):
+            txts[i][t] = r'$\boldsymbol{%s}$' % txts[i][t][1:-1]
+        for i, txt in enumerate(txts):
+            print(f'    & $\\beta_{i+1}$ &', ' & '.join(txt), r'\\')
+        print(r'\midrule')
+
+    pass
+
+def calc_selection_freqs(task, n):
+    def _count_one_init():
+        counts = np.zeros([model.m])
+        # init_vec = np.load(getpath('analysis/initial_seg.npy'))
+        obs_buffer = RingQueue(n)
+        for _ in range(runs):
+            for _ in range(h): obs_buffer.push(np.zeros([len(init_vecs), nz]))
+            obs_buffer.push(init_vecs)
+            # level_latvecs = [init_vec]
+            for _ in range(h):
+                obs = np.concatenate(obs_buffer.to_list(), axis=-1)
+                obs = torch.tensor(obs, device='cuda:0', dtype=torch.float)
+                muss, stdss, betas = model.get_intermediate(obs)
+                selects = torch.multinomial(betas.squeeze(), 1).squeeze()
+                mus = muss[[*range(len(init_vecs))], selects, :]
+                stds = stdss[[*range(len(init_vecs))], selects, :]
+                actions = Normal(mus, stds).rsample().squeeze().cpu().numpy()
+                obs_buffer.push(actions)
+                for i in selects:
+                    counts[i] = counts[i] + 1
+        return counts
+        # onehots = decoder(torch.tensor(level_latvecs).view(-1, nz, 1, 1))
+        pass
+    pass
+    init_vecs = np.load(getpath('analysis/initial_seg.npy'))
+    freqs = [[] for _ in range(30)]
+    start_line = 0
+    for l in ('0.0', '0.1', '0.2', '0.3', '0.4', '0.5'):
+        print(r'    \midrule')
+        for t, m in product(range(1, 6), (2, 3, 4, 5)):
+            path = getpath(f'test_data/varpm-{task}/l{l}_m{m}/t{t}')
+            model = torch.load(getpath(f'{path}/policy.pth'), map_location='cuda:0')
+            model.requires_grad_(False)
+            freq = np.zeros([m])
+            # n = load_cfgs(path, 'N')
+            runs, h = 100, 25
+            freq += _count_one_init()
+            freq /= (len(init_vecs) * runs * h)
+            freq = np.sort(freq)[::-1]
+            i = start_line + t - 1
+            freqs[i] += freq.tolist()
+            print(freqs[i])
+        start_line += 5
+        print(freqs)
+    with open(getpath(f'analysis/select_freqs-{task}.json'), 'w') as f:
+        json.dump(freqs, f)
+
+def print_selection_freq():
+    # task, n = 'lgp', 5
+    task, n = 'fhp', 11
+    if not os.path.exists(getpath(f'analysis/select_freqs-{task}.json')):
+        calc_selection_freqs(task, n)
+    with open(getpath(f'analysis/select_freqs-{task}.json'), 'r') as f:
+        freqs = json.load(f)
+    lbds = ['0.0', '0.1', '0.2', '0.3', '0.4', '0.5']
+    for i, row_data in enumerate(freqs):
+        if i % 5 == 0:
+            print(r'  \midrule')
+            print(r'  \multirow{5}{*}{$%s$}' % lbds[i//5])
+        txt = ' & '.join(map(_prob_fmt, row_data))
+        print(f'  & {i%5+1} &', txt, r'\\')
+
+def print_individual_performances(task):
+    for m, l in product((2, 3, 4, 5), ('0.0', '0.1', '0.2', '0.3', '0.4', '0.5')):
+        values = []
+        if l == '0.0':
+            print(r'  \midrule')
+            print(r'  \multirow{6}{*}{%d}' % m)
+        for t in range(1, 6):
+            path = f'test_data/varpm-{task}/l{l}_m{m}/t{t}/performance.csv'
+            reward, diversity = load_dict_json(path, 'reward', 'diversity')
+            values.append([reward, diversity])
+        values.sort(key=lambda item: -item[0])
+        values = [*chain(*values)]
+        txts = [_g_fmt(v) for v in values]
+        print('  &', f'${l}$ & ', ' & '.join(txts), r'\\')
+    pass
+
+if __name__ == '__main__':
+    # print_selection_prob('test_data/varpm-fhp/l0.5_m5/t5')
+    # print_selection_prob('test_data/varpm-fhp/l0.1_m5/t5')
+    # print_selection_freq()
+    # print_compare_tab_nonrl()
+    # print_individual_performances('fhp')
+    # print('\n\n')
+    # print_individual_performances('lgp')
+
+    # plot_cmp_learning_curves('fhp', save_path='results/learning_curves/fhp.png', title='MarioPuzzle')
+    # plot_cmp_learning_curves('lgp', save_path='results/learning_curves/lgp.png', title='MultiFacet')
+
+    # plot_crosstest_scatters('fhp', title='MarioPuzzle')
+    # plot_crosstest_scatters('lgp', title='MultiFacet')
+    # # plot_crosstest_scatters('fhp', yrange=(0, 2500), xrange=(20, 70), title='MarioPuzzle')
+    # plot_crosstest_scatters('lgp', yrange=(0, 1500), xrange=(20, 50), title='MultiFacet')
+    # plot_crosstest_scatters('lgp', yrange=(0, 800), xrange=(44, 48), title=' ')
+
+
+    # plot_varpm_heat('fhp', 'MarioPuzzle')
+    # plot_varpm_heat('lgp', 'MultiFacet')
+
+    vis_samples()
+
+    # make_tsne('fhp', 'MarioPuzzle', n=100)
+    # make_tsne('lgp', 'MultiFacet', n=100)
+    pass
+

pyproject.toml

@@ -0,0 +1,21 @@
+[tool.poetry]
+name = "ncerl"
+version = "0.1.0"
+description = ""
+authors = ["Ziqi Wang"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.9"
+JPype1 = "1.3.0"
+dtw = "1.4.0"
+torch = "1.8.1"
+numpy = "^2.0.0"
+pillow = "10.0.0"
+matplotlib = "3.6.3"
+pandas = "1.3.2"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

root.py

@@ -0,0 +1,3 @@
+import os
+
+PRJROOT = os.path.dirname(os.path.realpath(__file__))

smb/img/README.md

@@ -0,0 +1 @@
+This folder contain the game graphics. All the graphics file has been modified by hand from the original files that was grabbed from: https://www.spriters-resource.com/nes/supermariobros/. Except for the font.gif file as it is the same as file from the MarioAI framework.