From e44db33a91bfb56a2b1db34621772ba6863fe73e Mon Sep 17 00:00:00 2001
From: Dominik Roth <dominik.roth.dev@gmail.com>
Date: Wed, 3 May 2023 23:17:45 +0200
Subject: [PATCH] Better tests alias sit="desk sit"

---
 test.py | 83 ++++++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 67 insertions(+), 16 deletions(-)

diff --git a/test.py b/test.py
index e8c0655..b9b2360 100644
--- a/test.py
+++ b/test.py
@@ -3,6 +3,7 @@ from time import sleep, time
 import numpy as np
 import pygame
 import yaml
+import random
 
 from columbus import env
 from columbus.observables import Observable, CnnObservable
@@ -13,11 +14,14 @@ from priorConditionedAnnealing import pca
 
 
 def main():
-    agent_func = choosePlayType()
-    env = chooseEnv()
+    alg_name, agent_func = choosePlayType()
+    env = chooseEnv(alg_name)
     while True:
         playEnv(env, agent_func)
-        input('<again?>')
+        try:
+            agent_func.reset()
+        except:
+            pass
     env.close()
 
 
@@ -56,7 +60,7 @@ def loadConfigDefinedEnv(EnvClass):
     return EnvClass(fps=30, **cur)
 
 
-def chooseEnv():
+def chooseEnv(alg_name):
     envs = list(getAvaibleEnvs())
     for i, Env in enumerate(envs):
         print('['+str(i)+'] '+Env.__name__)
@@ -72,7 +76,7 @@ def chooseEnv():
         if envs[i] in [env.ColumbusConfigDefined]:
             return loadConfigDefinedEnv(envs[i])
         Env = envs[i]
-        return Env(fps=30, agent_draw_path=True, path_decay=1/256)
+        return Env(fps=30, agent_draw_path=True, path_decay=1/1024, title_appendix=' ['+alg_name+']', max_steps=30*10, clear_path_on_reset=False)
 
 
 def value_func(obs):
@@ -91,27 +95,33 @@ def human_input(obs, env):
 
 class Colored_Noise():
     def __init__(self, beta=1, dim_a=2, samples=2**18):
+        self.beta = beta
+        self.dim_a = dim_a
+        self.samples = samples
         self.index = 0
-        self.samples = cn.powerlaw_psd_gaussian(beta, (dim_a, samples))
+        self.reset()
 
     def __call__(self, obs, env):
         sample = self.samples[:, self.index]
         self.index += 1
         return sample
 
+    def reset(self):
+        self.samples = cn.powerlaw_psd_gaussian(
+            self.beta, (self.dim_a, self.samples), random_state=rand_seed())
+
 
 class PCA_Noise():
-    def __init__(self, dim_a=2, kernel_func='SE_1.41_1', window=8, ssf=-1):
-        self.dist = pca.PCA_Distribution(
-            action_dim=dim_a, par_strength='CONT_DIAG', kernel_func=kernel_func, window=window)
-        self.dist.proba_distribution(th.Tensor([[0]*2]), th.Tensor([[1]*2]))
-        self.traj = [[0]*dim_a]
-        self.ssf = 32
+    def __init__(self, dim_a=2, kernel_func='SE_1.41_1', window=64, ssf=-1):
+        self.dim_a = dim_a
+        self.kernel_func = kernel_func
+        self.window = window
+        self.ssf = ssf
         self.index = 0
+        self.reset()
 
     def __call__(self, obs, env):
-        if self.index == self.ssf:
-            self.index = 0
+        if self.ssf != -1 and self.index % self.ssf == 0:
             self.traj = [[0]*len(self.traj[0])]
         traj = th.Tensor(self.traj).unsqueeze(0)
         sample = self.dist.sample(traj).squeeze(0)
@@ -119,10 +129,50 @@ class PCA_Noise():
         self.index += 1
         return sample
 
+    def reset(self):
+        self.dist = pca.PCA_Distribution(
+            action_dim=self.dim_a, par_strength='CONT_DIAG', kernel_func=self.kernel_func, window=self.window)
+        self.dist.proba_distribution(th.Tensor([[0]*2]), th.Tensor([[1]*2]))
+        self.traj = [[0]*self.dim_a]
+
+
+class Colored_PCA_Noise():
+    def __init__(self, beta=1, dim_a=2, samples=2**18, kernel_func='SE_1.41_1', window=64, ssf=-1):
+        self.beta = beta
+        self.dim_a = dim_a
+        self.samples = samples
+        self.kernel_func = kernel_func
+        self.window = window
+        self.ssf = ssf
+        self.index = 0
+        self.reset()
+
+    def __call__(self, obs, env):
+        epsilon = self.samples[:, self.index]
+        if self.ssf != -1 and self.index % self.ssf == 0:
+            self.traj = [[0]*len(self.traj[0])]
+        traj = th.Tensor(self.traj).unsqueeze(0)
+        sample = self.dist.sample(traj).squeeze(0)
+        self.traj.append(sample)
+        self.index += 1
+        return sample
+
+    def reset(self):
+        self.dist = pca.PCA_Distribution(
+            action_dim=self.dim_a, par_strength='CONT_DIAG', kernel_func=self.kernel_func, window=self.window)
+        self.dist.proba_distribution(th.Tensor([[0]*2]), th.Tensor([[1]*2]))
+        self.traj = [[0]*self.dim_a]
+        self.samples = cn.powerlaw_psd_gaussian(
+            self.beta, (self.dim_a, self.samples), random_state=rand_seed())
+
+
+def rand_seed():
+    return int(99999999 * random.random())
+
 
 def choosePlayType():
     options = {'human': human_input, 'PCA': PCA_Noise(),
-               'REX': Colored_Noise(beta=0), 'PINK': Colored_Noise(beta=1), 'BROWN': Colored_Noise(beta=2), 'BETA.5': Colored_Noise(beta=.5)}
+               'REX': Colored_Noise(beta=0), 'PINK': Colored_Noise(beta=1), 'BROWN': Colored_Noise(beta=2), 'BETA.5': Colored_Noise(beta=.5), 'PINK_PCA': Colored_PCA_Noise(beta=1)}
     for i, name in enumerate(options):
         print('['+str(i)+'] '+name)
     while True:
@@ -136,7 +186,8 @@ def choosePlayType():
             print(
                 '[!] That is a number, but not one that makes sense in this context...')
         else:
-            return options[list(options.keys())[i]]
+            name = list(options.keys())[i]
+            return name, options[name]
 
 
 def playEnv(env, agent_func):