update

maik97 · Sep 29, 2021 · c984bf5 · c984bf5
1 parent fb285d8
commit c984bf5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/wacky_rl/agents/a2c.py b/wacky_rl/agents/a2c.py
@@ -14,13 +14,19 @@
 from wacky_rl.losses import ActorLoss
 from wacky_rl.transform import ExpectedReturnsCalculator
 from wacky_rl.trainer import Trainer
+from wacky_rl.logger import StatusPrinter
 
 
 class A2C(AgentCore):
 
-    def __init__(self, env, approximate_contin=False):
+    def __init__(self, env, approximate_contin=False, logger=None):
         super(A2C, self).__init__()
 
+        if logger is None:
+            self.logger = StatusPrinter('a2c')
+        else:
+            self.logger = logger
+
         self.approximate_contin = approximate_contin
         self.memory = BufferMemory()
         self.calc_returns = ExpectedReturnsCalculator()

diff --git a/wacky_rl/agents/ppo.py b/wacky_rl/agents/ppo.py
@@ -25,7 +25,7 @@ def __init__(self, env, approximate_contin=True, logger=None):
         super(PPO, self).__init__()
 
         if logger is None:
-            self.logger = StatusPrinter('test')
+            self.logger = StatusPrinter('ppo')
         else:
             self.logger = logger