google-deepmind / mctx

Monte Carlo tree search in JAX
Apache License 2.0
2.31k stars 188 forks source link

best practices for speed #48

Closed wcarvalho closed 1 year ago

wcarvalho commented 1 year ago

Hey, thanks for this amazing library.

Question: are there any best practices for ensuring that MCTX is fast? I find it quite slow and was curious if there's common gotchas I might be missing.

fidlej commented 1 year ago

Hi Wilka. Do you have an example of your code? I assume that you use jit. How many simulations do you use? With Gumbel MuZero you can train even with 2 simulations.

wcarvalho commented 1 year ago

Thank you, that's helpful.