Fix on store training

yura-hb · Mar 7, 2024 · 273e04c · 273e04c
1 parent c3fe8ef
commit 273e04c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/diploma_thesis/agents/base/marl_agent.py b/diploma_thesis/agents/base/marl_agent.py
@@ -66,7 +66,7 @@ def train_step(self):
 
     @filter(lambda self, *args: self.phase == TrainingPhase())
     def store(self, key: Key, sample: TrainingSample):
-        self.trainer[key].store(sample)
+        self.trainer[key].store(sample, self.model.policy)
 
     def loss_record(self):
         if self.keys is None:

diff --git a/diploma_thesis/agents/base/rl_agent.py b/diploma_thesis/agents/base/rl_agent.py
@@ -36,7 +36,7 @@ def train_step(self):
 
     @filter(lambda self, *args: self.phase != EvaluationPhase())
     def store(self, key: Key, sample: TrainingSample):
-        self.trainer.store(sample)
+        self.trainer.store(sample, self.model.policy)
 
     def loss_record(self):
         return self.trainer.loss_record()