ptree_sez bugs for discrete action spaces

[esparano]

Hi, I am trying Sampled EfficientZero for a large, discrete action space, and I ran into ~3 bugs in ptree_sez.py. I believe the following code has some issues:

            if self.legal_actions is not None:
                # first use the self.legal_actions to exclude the illegal actions
                policy_tmp = [0. for _ in range(self.action_space_size)]
                for index, legal_action in enumerate(self.legal_actions):
                    policy_tmp[legal_action] = policy_logits[index]
                policy_logits = policy_tmp
            # then empty the self.legal_actions
            self.legal_actions = []
            prob = torch.softmax(torch.tensor(policy_logits), dim=-1)
            sampled_actions = torch.multinomial(prob, self.num_of_sampled_actions, replacement=False)

            for action_index in range(self.num_of_sampled_actions):
                self.children[Action(sampled_actions[action_index].detach().cpu().numpy())] = Node( ...

1. By setting policy_tmp to 0, then taking the softmax, illegal moves still have non-zero probability. I just set the default to -1e9 as a hack, and it seems to work now, but it's probably more efficient to extract the list of legal moves before taking the softmax.

2. The line policy_tmp[legal_action] = policy_logits[index] is copying the wrong policy logits. I think it should be policy_tmp[legal_action] = policy_logits[legal_action].

3. I think instead of for action_index in range(self.num_of_sampled_actions):, it should be something like for action_index in range(len(sampled_actions)): This is because there could be fewer legal moves than sampled actions (K), and torch.multinomial is used with replacement=False so sampled_actions[action_index] will have an index out of bounds issue.

Sorry for not just submitting PRs for these bugs - I am not confident in the best way to fix it. I'm still working my way through my own fix so I may run into more related bugs later. I haven't verified that the rest of MCTS works, and I haven't checked the "ctree" equivalent, either.

[esparano]

Note: Solution 3 is causing other issues with other code (like game buffer code) which expects the sampled actions always to have the same size. Maybe it's possible to just pad with 0s, or similar. I'll let you all figure out the rest!

[puyuan1996]

Hi,

Thank you very much for your detailed report in #489! Your analysis of the discrete-action sampling path was spot-on and helped us catch the bug quickly.

We have fixed the issue in PR #491. The patch now:

• Gathers logits only for legal candidate actions before applying softmax
• Uses policy_logits[legal_action] for the full action-space logits
• Caps sampling by the actual number of legal actions to avoid out-of-bounds access when legal_count < K

The same legal-action masking and K-bound fix was also applied to the sampled MuZero C++ tree, and we added regression tests for ptree, sampled EfficientZero ctree, and sampled MuZero ctree.

Could you please pull the latest code and verify the fix locally? It would be great to confirm everything works correctly in your environment. (The fix was co-authored with Codex, which helped trace the C++ tree behavior.)

Feel free to open new issues if you run into anything else or have further suggestions. Thanks again for the contribution!