add src.env

.gitignore

@@ -125,11 +125,11 @@ celerybeat.pid
 *.sage.py
 
 # Environments
-.env
+# .env
 .venv
-env/
+# env/
 venv/
-ENV/
+# ENV/
 env.bak/
 venv.bak/
 

src/env/environments.py

@@ -0,0 +1,399 @@
+import random
+import time
+
+import gym
+import numpy as np
+from typing import Tuple, List, Dict, Callable, Optional
+
+from stable_baselines3.common.env_util import make_vec_env
+from stable_baselines3.common.vec_env import SubprocVecEnv
+from stable_baselines3.common.vec_env.base_vec_env import VecEnvStepReturn, VecEnvObs
+from stable_baselines3.common.vec_env.subproc_vec_env import _flatten_obs
+
+from src.env.logger import InfoCollector
+from src.env.rfunc import RewardFunc
+from src.smb.asyncsimlt import AsycSimltPool
+from src.smb.level import lvlhcat
+from src.gan.gans import SAGenerator
+from src.gan.gankits import *
+from src.smb.proxy import MarioProxy, MarioJavaAgents
+from src.utils.datastruct import RingQueue
+
+
+def get_padded_obs(vecs, histlen, add_batch_dim=False):
+    if len(vecs) < histlen:
+        lack = histlen - len(vecs)
+        pad = [np.zeros([nz], np.float32) for _ in range(lack)]
+        obs = np.concatenate(pad + vecs)
+    else:
+        obs = np.concatenate(vecs)
+    if add_batch_dim:
+        obs = np.reshape(obs, [1, -1])
+    return obs
+
+
+class SingleProcessOLGenEnv(gym.Env):
+    def __init__(self, rfunc, decoder: SAGenerator, eplen: int=50, device='cuda:0'):
+        self.rfunc = rfunc
+        self.initvec_set = np.load(getpath('smb/init_latvecs.npy'))
+        self.decoder = decoder
+        self.decoder.to(device)
+        self.hist_len = self.rfunc.get_n()
+        self.eplen = eplen
+        self.device = device
+        self.action_space = gym.spaces.Box(-1, 1, (nz,))
+        self.observation_space = gym.spaces.Box(-1, 1, (self.hist_len * nz,))
+        # self.obs_queue = RingQueue(self.hist_len)
+        self.lat_vecs = []
+        self.simulator = MarioProxy()
+        pass
+
+    def step(self, action):
+        self.lat_vecs.append(action)
+        done = len(self.lat_vecs) == (self.eplen + 1)
+        info = {}
+        if done:
+            rewss = self.__evalute()
+            info['rewss'] = rewss
+            rewsums = [sum(items) for items in zip(*rewss.values())]
+            info['transitions'] = self.__process_traj(rewsums[-self.eplen:])
+            self.reset()
+        return self.getobs(), 0, done, info
+
+    def __evalute(self):
+        z = torch.tensor(np.stack(self.lat_vecs).reshape([-1, nz, 1, 1]), device=self.device, dtype=torch.float)
+        # print(z.shape)
+        segs = process_onehot(self.decoder(z))
+        lvl = lvlhcat(segs)
+        simlt_res = MarioProxy.get_seg_infos(self.simulator.simulate_complete(lvl))
+        rewardss = self.rfunc.get_rewards(segs=segs, simlt_res=simlt_res)
+        return rewardss
+
+    def __process_traj(self, rewards):
+        obs = []
+        for i in range(1, len(self.lat_vecs) + 1):
+            ob = get_padded_obs(self.lat_vecs[max(0, i - self.hist_len): i], self.hist_len)
+            obs.append(ob)
+        traj = [(obs[i], self.lat_vecs[i+1], rewards[i], obs[i+1]) for i in range(len(self.lat_vecs) - 1)]
+        return traj
+
+    def reset(self):
+        self.lat_vecs.clear()
+        z0 = self.initvec_set[random.randrange(0, len(self.initvec_set))]
+        self.lat_vecs.append(z0)
+        return self.getobs()
+
+    def getobs(self):
+        s = max(0, len(self.lat_vecs) - self.hist_len)
+        return get_padded_obs(self.lat_vecs[s:], self.hist_len, True)
+
+    def render(self, mode="human"):
+        pass
+
+    def generate_levels(self, agent, n=1, max_parallel=None):
+        if max_parallel is None:
+            max_parallel = min(n, 512)
+        levels = []
+        latvecs = []
+        obs_queues = [RingQueue(self.hist_len) for _ in range(max_parallel)]
+        while len(levels) < n:
+            veclists = [[] for _ in range(min(max_parallel, n - len(levels)))]
+            for queue, veclist in zip(obs_queues, veclists):
+                queue.clear()
+                init_latvec = self.initvec_set[random.randrange(0, len(self.initvec_set))]
+                queue.push(init_latvec)
+                veclist.append(init_latvec)
+            for _ in range(self.eplen):
+                obs = np.stack([get_padded_obs(queue.to_list(), self.hist_len) for queue in obs_queues])
+                actions = agent.make_decision(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=self.device).view(-1, nz, 1, 1)
+                lvl = lvlhcat(process_onehot(self.decoder(z)))
+                levels.append(lvl)
+        return levels, latvecs
+
+
+class AsyncOlGenEnv:
+    def __init__(self, histlen, decoder: SAGenerator, eval_pool: AsycSimltPool, eplen: int=50, device='cuda:0'):
+        self.initvec_set = np.load(getpath('smb/init_latvecs.npy'))
+        self.decoder = decoder
+        self.decoder.to(device)
+        self.device = device
+        self.eval_pool = eval_pool
+        self.eplen = eplen
+        self.tid = 0
+        self.histlen = histlen
+
+        self.cur_vectraj = []
+        self.buffer = {}
+
+    def reset(self):
+        if len(self.cur_vectraj) > 0:
+            self.buffer[self.tid] = self.cur_vectraj
+            self.cur_vectraj = []
+            self.tid += 1
+        z0 = self.initvec_set[random.randrange(0, len(self.initvec_set))]
+        self.cur_vectraj.append(z0)
+        return self.getobs()
+
+    def step(self, action):
+        self.cur_vectraj.append(action)
+        done = len(self.cur_vectraj) == (self.eplen + 1)
+        if done:
+            self.__submit_eval_task()
+            self.reset()
+        return self.getobs(), done
+
+    def getobs(self):
+        s = max(0, len(self.cur_vectraj) - self.histlen)
+        return get_padded_obs(self.cur_vectraj[s:], self.histlen, True)
+
+    def __submit_eval_task(self):
+        z = torch.tensor(np.stack(self.cur_vectraj).reshape([-1, nz, 1, 1]), device=self.device)
+        segs = process_onehot(self.decoder(torch.clamp(z, -1, 1)))
+        lvl = lvlhcat(segs)
+        args = (self.tid, str(lvl))
+        self.eval_pool.put('evaluate', args)
+
+    def refresh(self):
+        if self.eval_pool is not None:
+            self.eval_pool.refresh()
+
+    def rollout(self, close=False, wait=False) -> Tuple[List[Tuple], List[Dict[str, List]]]:
+        transitions, rewss = [], []
+        if close:
+            eval_res = self.eval_pool.close()
+        else:
+            eval_res = self.eval_pool.get(wait)
+        for tid, rewards in eval_res:
+            rewss.append(rewards)
+            rewsums = [sum(items) for items in zip(*rewards.values())]
+            vectraj = self.buffer.pop(tid)
+            transitions += self.__process_traj(vectraj, rewsums[-self.eplen:])
+        return transitions, rewss
+
+    def __process_traj(self, vectraj, rewards):
+        obs = []
+        for i in range(1, len(vectraj) + 1):
+            ob = get_padded_obs(vectraj[max(0, i - self.histlen): i], self.histlen)
+            obs.append(ob)
+        traj = [(obs[i], vectraj[i+1], rewards[i], obs[i+1]) for i in range(len(vectraj) - 1)]
+        return traj
+
+    def close(self):
+        res = self.rollout(True)
+        self.eval_pool = None
+        return res
+
+    def generate_levels(self, agent, n=1, max_parallel=None):
+        if max_parallel is None:
+            max_parallel = min(n, 512)
+        levels = []
+        latvecs = []
+        obs_queues = [RingQueue(self.histlen) for _ in range(max_parallel)]
+        while len(levels) < n:
+            veclists = [[] for _ in range(min(max_parallel, n - len(levels)))]
+            for queue, veclist in zip(obs_queues, veclists):
+                queue.clear()
+                init_latvec = self.initvec_set[random.randrange(0, len(self.initvec_set))]
+                queue.push(init_latvec)
+                veclist.append(init_latvec)
+            for _ in range(self.eplen):
+                obs = np.stack([get_padded_obs(queue.to_list(), self.histlen) for queue in obs_queues])
+                actions = agent.make_decision(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=self.device).view(-1, nz, 1, 1)
+                lvl = lvlhcat(process_onehot(self.decoder(z)))
+                levels.append(lvl)
+        return levels, latvecs
+
+
+######### Adopt from https://github.com/SUSTechGameAI/MFEDRL #########
+class SyncOLGenWorkerEnv(gym.Env):
+    def __init__(self, rfunc=None, hist_len=5, eplen=25, return_lvl=False, init_one=False, play_style='Runner'):
+        self.rfunc = RewardFunc() if rfunc is None else rfunc
+        self.mario_proxy = MarioProxy() if self.rfunc.require_simlt else None
+        self.action_space = gym.spaces.Box(-1, 1, (nz,))
+        self.hist_len = hist_len
+        self.observation_space = gym.spaces.Box(-1, 1, (hist_len * nz,))
+        self.segs = []
+        self.latvec_archive = RingQueue(hist_len)
+        self.eplen = eplen
+        self.counter = 0
+        # self.repairer = DivideConquerRepairer()
+        self.init_one = init_one
+        self.backup_latvecs = None
+        self.backup_strsegs = None
+        self.return_lvl = return_lvl
+        self.jagent = MarioJavaAgents.__getitem__(play_style)
+        self.simlt_k = 80 if play_style == 'Runner' else 320
+
+    def receive(self, **kwargs):
+        for key in kwargs.keys():
+            setattr(self, key, kwargs[key])
+
+    def step(self, data):
+        action, strseg = data
+        seg = MarioLevel(strseg)
+        self.latvec_archive.push(action)
+
+        self.counter += 1
+        self.segs.append(seg)
+        done = self.counter >= self.eplen
+        if done:
+            full_level = lvlhcat(self.segs)
+            # full_level = self.repairer.repair(full_level)
+            w = MarioLevel.seg_width
+            segs = [full_level[:, s: s + w] for s in range(0, full_level.w, w)]
+            if self.mario_proxy:
+                raw_simlt_res = self.mario_proxy.simulate_complete(lvlhcat(segs), self.jagent, self.simlt_k)
+                simlt_res = MarioProxy.get_seg_infos(raw_simlt_res)
+            else:
+                simlt_res = None
+            rewards = self.rfunc.get_rewards(segs=segs, simlt_res=simlt_res)
+            info = {}
+            total_score = 0
+            if self.return_lvl:
+                info['LevelStr'] = str(full_level)
+            for key in rewards:
+                info[f'{key}_reward_list'] = rewards[key][-self.eplen:]
+                info[f'{key}'] = sum(rewards[key][-self.eplen:])
+                total_score += info[f'{key}']
+            info['TotalScore'] = total_score
+            info['EpLength'] = self.counter
+        else:
+            info = {}
+        return self.__get_obs(), 0, done, info
+
+    def reset(self):
+        self.segs.clear()
+        self.latvec_archive.clear()
+        for latvec, strseg in zip(self.backup_latvecs, self.backup_strsegs):
+            self.latvec_archive.push(latvec)
+            self.segs.append(MarioLevel(strseg))
+
+        self.backup_latvecs, self.backup_strsegs = None, None
+        self.counter = 0
+        return self.__get_obs()
+
+    def __get_obs(self):
+        lack = self.hist_len - len(self.latvec_archive)
+        pad = [np.zeros([nz], np.float32) for _ in range(lack)]
+        return np.concatenate([*pad, *self.latvec_archive.to_list()])
+
+    def render(self, mode='human'):
+        pass
+
+
+class VecOLGenEnv(SubprocVecEnv):
+    def __init__(
+        self, env_fns: List[Callable[[], gym.Env]], start_method: Optional[str] = None, hist_len=5, eplen=50,
+        init_one=True, log_path=None, log_itv=-1, log_targets=None, device='cuda:0'
+    ):
+        super(VecOLGenEnv, self).__init__(env_fns, start_method)
+        self.decoder = get_decoder(device=device)
+
+        if log_path:
+            self.logger = InfoCollector(log_path, log_itv, log_targets)
+        else:
+            self.logger = None
+        self.hist_len = hist_len
+        self.total_steps = 0
+        self.start_time = time.time()
+        self.eplen = eplen
+        self.device = device
+        self.init_one = init_one
+        self.latvec_set = np.load(getpath('smb/init_latvecs.npy'))
+
+    def step_async(self, actions: np.ndarray) -> None:
+        with torch.no_grad():
+            z = torch.tensor(actions.astype(np.float32), device=self.device).view(-1, nz, 1, 1)
+            segs = process_onehot(self.decoder(z))
+        for remote, action, seg in zip(self.remotes, actions, segs):
+            remote.send(("step", (action, str(seg))))
+        self.waiting = True
+
+    def step_wait(self) -> VecEnvStepReturn:
+        self.total_steps += self.num_envs
+        results = [remote.recv() for remote in self.remotes]
+        self.waiting = False
+        obs, rews, dones, infos = zip(*results)
+
+        envs_to_send = [i for i in range(self.num_envs) if dones[i]]
+        self.send_reset_data(envs_to_send)
+
+        if self.logger is not None:
+            for i in range(self.num_envs):
+                if infos[i]:
+                    infos[i]['TotalSteps'] = self.total_steps
+                    infos[i]['TimePassed'] = time.time() - self.start_time
+            self.logger.on_step(dones, infos)
+        return _flatten_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos
+
+    def reset(self) -> VecEnvObs:
+        self.send_reset_data()
+        for remote in self.remotes:
+            remote.send(("reset", None))
+        obs = [remote.recv() for remote in self.remotes]
+        self.send_reset_data()
+        return _flatten_obs(obs, self.observation_space)
+
+    def send_reset_data(self, env_ids=None):
+        if env_ids is None:
+            env_ids = [*range(self.num_envs)]
+        target_remotes = self._get_target_remotes(env_ids)
+
+        n_inits = 1 if self.init_one else self.hist_len
+        # latvecs = [sample_latvec(n_inits, tensor=False) for _ in range(len(env_ids))]
+
+        latvecs = [self.latvec_set[random.sample(range(len(self.latvec_set)), n_inits)] for _ in range(len(env_ids))]
+        with torch.no_grad():
+            segss = [[] for _ in range(len(env_ids))]
+            for i in range(len(env_ids)):
+                z = torch.tensor(latvecs[i]).view(-1, nz, 1, 1).to(self.device)
+                # print(self.decoder(z).shape)
+                segss[i] = [process_onehot(self.decoder(z))] if self.init_one else process_onehot(self.decoder(z))
+        for remote, latvec, segs in zip(target_remotes, latvecs, segss):
+            kwargs = {'backup_latvecs': latvec, 'backup_strsegs': [str(seg) for seg in segs]}
+            remote.send(("env_method", ('receive', [], kwargs)))
+        for remote in target_remotes:
+            remote.recv()
+
+    def close(self) -> None:
+        super().close()
+        if self.logger is not None:
+            self.logger.close()
+
+
+def make_vec_offrew_env(
+        num_envs, rfunc=None, log_path=None, eplen=25, log_itv=-1, hist_len=5, init_one=True,
+        play_style='Runner', device='cuda:0', log_targets=None, return_lvl=False
+    ):
+    return make_vec_env(
+        SyncOLGenWorkerEnv, n_envs=num_envs, vec_env_cls=VecOLGenEnv,
+        vec_env_kwargs={
+            'log_path': log_path,
+            'log_itv': log_itv,
+            'log_targets': log_targets,
+            'device': device,
+            'eplen': eplen,
+            'hist_len': hist_len,
+            'init_one': init_one
+        },
+        env_kwargs={
+            'rfunc': rfunc,
+            'eplen': eplen,
+            'return_lvl': return_lvl,
+            'play_style': play_style,
+            'hist_len': hist_len,
+            'init_one': init_one
+        }
+    )

src/env/logger.py

@@ -0,0 +1,182 @@
+import json
+import os
+import csv
+import numpy as np
+from itertools import product
+from src.smb.level import save_batch
+from src.utils.filesys import getpath
+
+############### Loggers for async environment ###############
+class AsyncCsvLogger:
+    def __init__(self, target, rfunc, buffer_size=50):
+        self.rterms = tuple(term.get_name() for term in rfunc.terms)
+        self.cols = ('steps', *self.rterms, 'reward_sum', 'time', 'trans', 'updates', '')
+        self.buffer = []
+        self.buffer_size = buffer_size
+        self.ftarget = open(getpath(target), 'w', newline='')
+        self.writer = csv.writer(self.ftarget)
+        self.writer.writerow(self.cols)
+
+    def on_episode(self, **kwargs):
+        for rews in kwargs['rewss']:
+            rews_list = [sum(rews[key]) for key in self.rterms]
+            self.buffer.append([
+                kwargs['steps'], *rews_list, sum(rews_list),
+                kwargs['time'], kwargs['trans'], kwargs['updates']
+            ])
+        self.__try_write()
+        if kwargs['close']:
+            self.close()
+
+    def __try_write(self):
+        if len(self.buffer) < self.buffer_size:
+            return
+        self.writer.writerows(self.buffer)
+        self.ftarget.flush()
+        self.buffer.clear()
+
+    def close(self):
+        self.writer.writerows(self.buffer)
+        self.ftarget.close()
+        pass
+
+
+class AsyncStdLogger:
+    def __init__(self, rfunc, itv=2000, path=''):
+        self.rterms = tuple(term.get_name() for term in rfunc.terms)
+        self.rews = {rterm: 0. for rterm in self.rterms}
+        self.n = 0
+        self.itv = itv
+        self.horizon = itv
+        if not len(path):
+            self.f = None
+        else:
+            self.f = open(getpath(path), 'w')
+        self.last_steps = 0
+        self.last_trans = 0
+        self.last_updates = 0
+        self.buffer = []
+        pass
+
+    def on_episode(self, **kwargs):
+        newrews = {k: self.rews[k] for k in self.rterms}
+        for rews, k in product(kwargs['rewss'], self.rterms):
+            newrews[k] = newrews[k] + sum(rews[k])
+        self.rews = newrews
+        self.n += len(kwargs['rewss'])
+        if kwargs['steps'] >= self.horizon or kwargs['close']:
+            self.__output(**kwargs)
+            self.horizon += self.itv
+            self.rews = {k: 0 for k in self.rews.keys()}
+            self.n = 0
+            self.last_steps = kwargs['steps']
+            self.last_trans = kwargs['trans']
+            self.last_updates = kwargs['updates']
+        if kwargs['close'] and self.f is not None:
+            self.f.close()
+
+    def __output(self, **kwargs):
+        steps = kwargs['steps']
+        if kwargs['close']:
+            head = '-' * 20 + 'Closing rollouts' + '-' * 20
+        else:
+            head = '-' * 20 + f'Rollout of {self.last_steps}-{steps} steps' + '-' * 20
+        self.buffer.append(head)
+        rsum = 0
+        for t in self.rterms:
+            v = 0 if self.n == 0 else self.rews[t] / self.n
+            rsum += v
+            self.buffer.append(f'{t}: {v:.2f}')
+        self.buffer.append(f'Reward sum: {rsum: .2f}')
+        self.buffer.append('Time elapsed: %.1fs' % kwargs['time'])
+        self.buffer.append('Transitions collected: %d (%d in total)' % (kwargs['trans'] - self.last_trans, kwargs['trans']))
+        self.buffer.append('Number of updates: %d (%d in total)' % (kwargs['updates']- self.last_updates, kwargs['updates']))
+        if self.f is None:
+            print('\n'.join(self.buffer) + '\n')
+        else:
+            self.f.write('\n'.join(self.buffer) + '\n')
+            self.f.flush()
+        self.buffer.clear()
+        pass
+
+    def __reset(self):
+        pass
+
+
+class GenResLogger:
+    def __init__(self, root_path, k, itv=5000):
+        self.k = k
+        self.itv = itv
+        self.horizon = 0
+        self.path = getpath(f'{root_path}/gen_log')
+        os.makedirs(self.path, exist_ok=True)
+
+    def on_episode(self, env, agent, steps):
+        if steps >= self.horizon:
+            lvls, vectraj = env.generate_levels(agent, self.k)
+            # np.save(f'{self.path}/step{steps}', vectraj)
+            if len(lvls):
+                save_batch(lvls, f'{self.path}/step{steps}')
+            self.horizon += self.itv
+        pass
+    pass
+
+
+############### Loggers for sync environment, from https://github.com/SUSTechGameAI/MFEDRL ###############
+class InfoCollector:
+    ignored_keys = {'episode', 'terminal_observation'}
+    save_itv = 1000
+
+    def __init__(self, path, log_itv=100, log_targets=None):
+        self.data = []
+        self.path = path
+        self.msg_itv = log_itv
+        self.time_before_save = InfoCollector.save_itv
+        self.msg_ptr = 0
+        self.log_targets = [] if log_targets is None else log_targets
+        if 'file' in log_targets:
+            with open(f'{self.path}/log.txt', 'w') as f:
+                f.write('')
+        self.recent_time = 0
+
+    def on_step(self, dones, infos):
+        for done, info in zip(dones, infos):
+            if done:
+                self.data.append({
+                    key: val for key, val in info.items()
+                    if key not in InfoCollector.ignored_keys and 'reward_list' not in key
+                })
+                self.time_before_save -= 1
+        if self.time_before_save <= 0:
+            with open(f'{self.path}/ep_infos.json', 'w') as f:
+                json.dump(self.data, f)
+            self.time_before_save += InfoCollector.save_itv
+
+        if self.log_targets and 0 < self.msg_itv <= (len(self.data) - self.msg_ptr):
+            keys = set(self.data[-1].keys()) - {'TotalSteps', 'TimePassed', 'TotalScore', 'EpLen'}
+
+            msg = '%sTotal steps: %d%s\n' % ('-' * 16, self.data[-1]['TotalSteps'], '-' * 16)
+            msg += 'Time passed: %ds\n' % self.data[-1]['TimePassed']
+            t = self.data[-1]['TimePassed'] - self.recent_time
+            self.recent_time = self.data[-1]['TimePassed']
+            f = sum(item['EpLength'] for item in self.data[self.msg_ptr:])
+            msg += 'fps: %.3g\n' % (f/t)
+            for key in keys:
+                values = [item[key] for item in self.data[self.msg_ptr:]]
+                values = np.array(values)
+                msg += '%s: %.2f +- %.2f\n' % (key, values.mean(), values.std())
+            values = [item['TotalScore'] for item in self.data[self.msg_ptr:]]
+            values = np.array(values)
+            msg += 'TotalScore: %.2f +- %.2f\n' % (values.mean(), values.std())
+
+            if 'file' in self.log_targets:
+                with open(f'{self.path}/log.txt', 'a') as f:
+                    f.write(msg + '\n')
+            if 'std' in self.log_targets:
+                print(msg)
+            self.msg_ptr = len(self.data)
+            pass
+
+    def close(self):
+        with open(self.path, 'w') as f:
+            json.dump(self.data, f)

src/env/rfunc.py

@@ -0,0 +1,219 @@
+import numpy as np
+from math import ceil
+from abc import abstractmethod
+from src.utils.mymath import a_clip
+from src.smb.level import *
+
+defaults = {'n': 5, 'gl': 0.14, 'gg': 0.30, 'wl': 2, 'wg': 10}
+
+
+class RewardFunc:
+    def __init__(self, *args):
+        self.terms = args
+        self.require_simlt = any(term.require_simlt for term in self.terms)
+
+    def get_rewards(self, **kwargs):
+        return {
+            term.get_name(): term.compute_rewards(**kwargs)
+            for term in self.terms
+        }
+
+    def get_n(self):
+        n = 1
+        for term in self.terms:
+            try:
+                n = max(n, term.n)
+            except AttributeError:
+                pass
+        return n
+
+    def __str__(self):
+        return 'Reward Function:\n' + ',\n'.join('\t' + str(term) for term in self.terms)
+
+
+class RewardTerm:
+    def __init__(self, require_simlt):
+        self.require_simlt = require_simlt
+
+    def get_name(self):
+        return self.__class__.__name__
+
+    @abstractmethod
+    def compute_rewards(self, **kwargs):
+        pass
+
+
+class Playability(RewardTerm):
+    def __init__(self, magnitude=1):
+        super(Playability, self).__init__(True)
+        self.magnitude=magnitude
+
+    def compute_rewards(self, **kwargs):
+        simlt_res = kwargs['simlt_res']
+        return [0 if item['playable'] else -self.magnitude for item in simlt_res[1:]]
+
+    def __str__(self):
+        return f'{self.magnitude} * Playability'
+
+
+class MeanDivergenceFun(RewardTerm):
+    def __init__(self, goal_div, n=defaults['n'], s=8):
+        super().__init__(False)
+        self.l = goal_div * 0.26 / 0.6
+        self.u = goal_div * 0.94 / 0.6
+        self.n = n
+        self.s = s
+
+    def compute_rewards(self, **kwargs):
+        segs = kwargs['segs']
+        rewards = []
+        for i in range(1, len(segs)):
+            seg = segs[i]
+            histroy = lvlhcat(segs[max(0, i - self.n): i])
+            k = 0
+            divergences = []
+            while k * self.s <= (min(self.n, i) - 1) * MarioLevel.seg_width:
+                cmp_seg = histroy[:, k * self.s: k * self.s + MarioLevel.seg_width]
+                # print(i, nd, cmp_seg.shape)
+                divergences.append(tile_pattern_js_div(seg, cmp_seg))
+                k += 1
+            mean_d = sum(divergences) / len(divergences)
+            if mean_d < self.l:
+                rewards.append(-(mean_d - self.l) ** 2)
+            elif mean_d > self.u:
+                rewards.append(-(mean_d - self.u) ** 2)
+            else:
+                rewards.append(0)
+        return rewards
+
+
+class SACNovelty(RewardTerm):
+    def __init__(self, magnitude, goal_div, require_simlt, n):
+        super().__init__(require_simlt)
+        self.g = goal_div
+        self.magnitude = magnitude
+        self.n = n
+
+    def compute_rewards(self, **kwargs):
+        n_segs = len(kwargs['segs'])
+        rewards = []
+        for i in range(1, n_segs):
+            reward = 0
+            r_sum = 0
+            for k in range(1, self.n + 1):
+                j = i - k
+                if j < 0:
+                    break
+                r = 1 - k / (self.n + 1)
+                r_sum += r
+                reward += a_clip(self.disim(i, j, **kwargs), self.g, r)
+            rewards.append(reward * self.magnitude / r_sum)
+        return rewards
+
+    @abstractmethod
+    def disim(self, i, j, **kwargs):
+        pass
+
+
+class LevelSACN(SACNovelty):
+    def __init__(self, magnitude=1, g=defaults['gl'], w=defaults['wl'], n=defaults['n']):
+        super(LevelSACN, self).__init__(magnitude, g, False, n)
+        self.w = w
+
+    def disim(self, i, j, **kwargs):
+        segs = kwargs['segs']
+        seg1, seg2 = segs[i], segs[j]
+        return tile_pattern_js_div(seg1, seg2, self.w)
+
+    def __str__(self):
+        s = f'{self.magnitude} * LevelSACN(g={self.g:.3g}, w={self.w}, n={self.n})'
+        return s
+
+
+class GameplaySACN(SACNovelty):
+    def __init__(self, magnitude=1, g=defaults['gg'], w=defaults['wg'], n=defaults['n']):
+        super(GameplaySACN, self).__init__(magnitude, g, True, n)
+        self.w = w
+
+    def disim(self, i, j, **kwargs):
+        simlt_res = kwargs['simlt_res']
+        trace1, trace2 = simlt_res[i]['trace'], simlt_res[j]['trace']
+        return trace_div(trace1, trace2, self.w)
+
+    def __str__(self):
+        s = f'{self.magnitude} * GameplaySACN(g={self.g:.3g}, w={self.w}, n={self.n})'
+        return s
+
+
+class Fun(RewardTerm):
+    def __init__(self, magnitude=1., num_windows=3, lb=0.26, ub=0.94, stride=8):
+        super().__init__(False)
+        self.lb, self.ub = lb, ub
+        self.magnitude = magnitude
+        self.stride = stride
+        self.num_windows = num_windows
+        self.n = ceil(num_windows * stride / MarioLevel.seg_width - 1e-8)
+
+    def compute_rewards(self, **kwargs):
+        n_segs = len(kwargs['segs'])
+        lvl = lvlhcat(kwargs['segs'])
+        W = MarioLevel.seg_width
+        rewards = []
+        for i in range(1, n_segs):
+            seg = lvl[:, W*i: W*(i+1)]
+            divs = []
+            for k in range(0, self.num_windows + 1):
+                s = W * i - k * self.stride
+                if s < 0:
+                    break
+                cmp_seg = lvl[:, s:s+W]
+                divs.append(tile_pattern_kl_div(seg, cmp_seg))
+            mean_div = np.mean(divs)
+            rew = 0
+            if mean_div > self.ub:
+                rew = -(self.ub - mean_div) ** 2
+            if mean_div < self.lb:
+                rew = -(self.lb - mean_div) ** 2
+            rewards.append(rew * self.magnitude)
+        return rewards
+
+    def __str__(self):
+        s = f'{self.magnitude} * Fun(lb={self.lb:.2f}, ub={self.ub:.2f}, n={self.num_windows}, stride={self.stride})'
+        return s
+
+
+class HistoricalDeviation(RewardTerm):
+    def __init__(self, magnitude=1., m=3, n=10):
+        super().__init__(False)
+        self.magnitude = magnitude
+        self.m = m
+        self.n = n
+
+    def compute_rewards(self, **kwargs):
+        segs = kwargs['segs']
+        n_segs = len(kwargs['segs'])
+        rewards = []
+        for i in range(1, n_segs):
+            divs = []
+            for k in range(1, self.n+1):
+                j = i - k
+                if j < 0:
+                    break
+                divs.append(tile_pattern_kl_div(segs[i], segs[j]))
+            divs.sort()
+            m = min(i, self.m)
+            rew = np.mean(divs[:m])
+            rewards.append(rew * self.magnitude)
+        return rewards
+
+    def __str__(self):
+        return f'{self.magnitude} * HistoricalDeviation(m={self.m}, n={self.n})'
+
+
+if __name__ == '__main__':
+    # print(type(ceil(0.2)))
+    # arr = [1., 3., 2.]
+    # arr.sort()
+    # print(arr)
+    rfunc = HistoricalDeviation()
+

src/env/rfuncs.py

@@ -0,0 +1,9 @@
+from src.env.rfunc import *
+
+default = lambda : RewardFunc(LevelSACN(), GameplaySACN(), Playability())
+fp = lambda : RewardFunc(Fun(200), Playability(10))
+hp = lambda : RewardFunc(HistoricalDeviation(1), Playability(1))
+fhp = lambda : RewardFunc(Fun(30, num_windows=21), HistoricalDeviation(3), Playability(3))
+lp = lambda : RewardFunc(LevelSACN(2.), Playability(2.))
+gp = lambda : RewardFunc(GameplaySACN(2.), Playability(2.))
+lgp = lambda : RewardFunc(LevelSACN(), GameplaySACN(), Playability())