google-deepmind / mctx

Monte Carlo tree search in JAX
Apache License 2.0
2.33k stars 189 forks source link

Question: What is the purpose of extra_data in the class Tree? #30

Closed karroyan closed 1 year ago

karroyan commented 2 years ago

Hello, thanks for the nice project!

There's a parameter extra_data in function search, which is passed to class Tree. But it has not been used, so what's the purpose of this parameter?

fidlej commented 2 years ago

Thanks for asking.

The extra_data is used in gumbel_muzero_policy() to pass extra info about the root. https://github.com/deepmind/mctx/blob/8dac10f39786f13122785dd03b5ddcd3e393a655/mctx/_src/policies.py#L188

The gumbel_muzero_root_action_selection() uses that later: https://github.com/deepmind/mctx/blob/8dac10f39786f13122785dd03b5ddcd3e393a655/mctx/_src/action_selection.py#L149

You can use the gumbel_muzero_policy() and not worry about these internal details.