fix mcts update_root arg order

lowrollr · Feb 2, 2024 · 4e55c72 · 4e55c72
1 parent 18fdf52
commit 4e55c72
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/core/evaluators/alphazero.py b/core/evaluators/alphazero.py
@@ -27,7 +27,7 @@ def get_config(self) -> Dict:
             **super().get_config()
         }
 
-    def update_root(self, tree: MCTSTree, root_embedding: chex.ArrayTree, root_metadata: StepMetadata, params: chex.ArrayTree) -> MCTSTree:
+    def update_root(self, tree: MCTSTree, root_embedding: chex.ArrayTree, params: chex.ArrayTree, root_metadata: StepMetadata) -> MCTSTree:
         key, tree = get_rng(tree)
         root_policy_logits, root_value = self.eval_fn(root_embedding, params, key)
         root_policy = jax.nn.softmax(root_policy_logits)

diff --git a/core/evaluators/mcts/mcts.py b/core/evaluators/mcts/mcts.py
@@ -78,7 +78,7 @@ def evaluate(self,
         Returns:
         - `MCTSOutput`: contains new tree state, selected action, root value, and policy weights
         """
-        eval_state = self.update_root(eval_state, env_state, root_metadata, params)
+        eval_state = self.update_root(eval_state, env_state, params, root_metadata=root_metadata)
         iterate = partial(self.iterate, 
             params=params,
             env_step_fn=env_step_fn