diff --git a/MALKHUT/malkhut/training/cma_trainer.py b/MALKHUT/malkhut/training/cma_trainer.py index 8f695c8..e176cc0 100644 --- a/MALKHUT/malkhut/training/cma_trainer.py +++ b/MALKHUT/malkhut/training/cma_trainer.py @@ -1237,11 +1237,13 @@ class CMAESTrainer: evaluator: PolicyEvaluator, pool: SelfPlayPool, store: Optional[MalkhutCHStore] = None, + workers: int = 0, ) -> None: self.codec = codec self.evaluator = evaluator self.pool = pool self.store = store + self._workers = workers def train( self, @@ -1283,7 +1285,7 @@ class CMAESTrainer: candidate = self.codec.decode(x, version=f"cma_{time.time_ns()}_{evals}") score, results = self.evaluator.evaluate_candidate( params=candidate, scenarios=scenarios, rng_seed=seed + evals, - planner_type=planner_type, + planner_type=planner_type, workers=self._workers, ) perf_vec = PolicyEvaluator.performance_vector(results) snap = PolicySnapshot(