diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 78e2d1536..4f166ba3e 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -81,7 +81,7 @@ jobs: pip install -r requirements-build.txt --upgrade --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple pip install -r requirements-test.txt --upgrade --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple pip uninstall DI-engine -y - pip install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple git+https://github.com/opendilab/DI-engine.git@main#egg=DI-engine[common_env] --upgrade + pip install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple "DI-engine[common_env] @ git+https://github.com/opendilab/DI-engine.git@main" --upgrade pip install --no-cache-dir transformers --upgrade -i https://pypi.tuna.tsinghua.edu.cn/simple pip install opencv-python --upgrade --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple - name: Verify Python.h availability diff --git a/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.cpp b/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.cpp index 1044b1acf..ce6cd0568 100644 --- a/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.cpp +++ b/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.cpp @@ -285,52 +285,52 @@ namespace tree else { // discrete action space for sampled algo.. + std::vector candidate_actions; + if (this->legal_actions.empty()) + { + for (int i = 0; i < this->action_space_size; ++i) + { + candidate_actions.push_back(i); + } + } + else + { + for (int i = 0; i < this->legal_actions.size(); ++i) + { + candidate_actions.push_back(int(this->legal_actions[i].value[0])); + } + } - //======================================================== - // python code - //======================================================== - // if self.legal_actions is not None: - // # fisrt 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 = [] - // then empty the self.legal_actions - // prob = torch.softmax(torch.tensor(policy_logits), dim=-1) - // sampled_actions = torch.multinomial(prob, self.num_of_sampled_actions, replacement=False) - - //======================================================== - // TODO(pu): legal actions - //======================================================== - // std::vector policy_tmp; - // for (int i = 0; i < this->action_space_size; ++i) - // { - // policy_tmp.push_back(0.); - // } - // for (int i = 0; i < this->legal_actions.size(); ++i) - // { - // policy_tmp[this->legal_actions[i].value] = policy_logits[i]; - // } - // for (int i = 0; i < this->action_space_size; ++i) - // { - // policy_logits[i] = policy_tmp[i]; - // } - // std::cout << "position 3" << std::endl; + assert(candidate_actions.size() > 0); + assert(policy_logits.size() == this->action_space_size || policy_logits.size() == candidate_actions.size()); - // python code: legal_actions = [] - std::vector legal_actions; + std::vector candidate_logits; + for (int i = 0; i < candidate_actions.size(); ++i) + { + int action = candidate_actions[i]; + if (policy_logits.size() == this->action_space_size) + { + candidate_logits.push_back(policy_logits[action]); + } + else + { + candidate_logits.push_back(policy_logits[i]); + } + } - // python code: probs = softmax(policy_logits) float logits_exp_sum = 0; - for (int i = 0; i < policy_logits.size(); ++i) + float max_logit = candidate_logits[0]; + for (int i = 1; i < candidate_logits.size(); ++i) + { + max_logit = std::max(max_logit, candidate_logits[i]); + } + for (int i = 0; i < candidate_logits.size(); ++i) { - logits_exp_sum += exp(policy_logits[i]); + logits_exp_sum += exp(candidate_logits[i] - max_logit); } - for (int i = 0; i < policy_logits.size(); ++i) + for (int i = 0; i < candidate_logits.size(); ++i) { - probs.push_back(exp(policy_logits[i]) / (logits_exp_sum + 1e-6)); + probs.push_back(exp(candidate_logits[i] - max_logit) / (logits_exp_sum + 1e-6)); } unsigned seed = std::chrono::system_clock::now().time_since_epoch().count(); @@ -392,10 +392,10 @@ namespace tree std::sort(disc_action_with_probs.begin(), disc_action_with_probs.end(), cmp); - // take the fist ``num_of_sampled_actions`` actions - for (int k = 0; k < num_of_sampled_actions; ++k) + int num_sampled = std::min(this->num_of_sampled_actions, int(candidate_actions.size())); + for (int k = 0; k < num_sampled; ++k) { - sampled_actions.push_back(disc_action_with_probs[k].first); + sampled_actions.push_back(candidate_actions[disc_action_with_probs[k].first]); // disc_action_with_probs[k].second is disturbed_probs // sampled_actions_probs.push_back(disc_action_with_probs[k].second); sampled_actions_probs.push_back(probs[disc_action_with_probs[k].first]); @@ -419,18 +419,23 @@ namespace tree disc_action_with_probs.clear(); // Empty the collection to prepare for the next sampling. } - float prior; - for (int i = 0; i < this->num_of_sampled_actions; ++i) - { + this->legal_actions.clear(); + this->sampled_actions.clear(); - if (this->continuous_action_space == true) + if (this->continuous_action_space == true) + { + for (int i = 0; i < this->num_of_sampled_actions; ++i) { CAction action = CAction(sampled_actions_after_tanh[i], 0); std::vector legal_actions; this->children[action.get_combined_hash()] = CNode(sampled_actions_log_probs_after_tanh[i], legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space); // only for muzero/efficient zero, not support alphazero this->legal_actions.push_back(action); + this->sampled_actions.push_back(action); } - else + } + else + { + for (int i = 0; i < sampled_actions.size(); ++i) { std::vector sampled_action_tmp; for (size_t iter = 0; iter < 1; iter++) @@ -441,6 +446,11 @@ namespace tree std::vector legal_actions; this->children[action.get_combined_hash()] = CNode(sampled_actions_probs[i], legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space); // only for muzero/efficient zero, not support alphazero this->legal_actions.push_back(action); + this->sampled_actions.push_back(action); + } + while (!this->sampled_actions.empty() && this->sampled_actions.size() < this->num_of_sampled_actions) + { + this->sampled_actions.push_back(this->sampled_actions.back()); } } @@ -461,7 +471,8 @@ namespace tree - noises: the vector of noises added to each child node. */ float noise, prior; - for (int i = 0; i < this->num_of_sampled_actions; ++i) + int noise_num = std::min(int(this->legal_actions.size()), int(noises.size())); + for (int i = 0; i < noise_num; ++i) { noise = noises[i]; @@ -596,6 +607,10 @@ namespace tree CNode *child = this->get_child(a); distribution.push_back(child->visit_count); } + while (distribution.size() < this->sampled_actions.size()) + { + distribution.push_back(0); + } } return distribution; } @@ -644,13 +659,13 @@ namespace tree for (int i = 0; i < this->root_num; ++i) { - if (this->continuous_action_space == true and this->legal_actions_list[0][0] == -1) + if (this->continuous_action_space == true) { // continous action space std::vector legal_actions; this->roots.push_back(CNode(0, legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space)); } - else if (this->continuous_action_space == false or this->legal_actions_list[0][0] == -1) + else if (this->legal_actions_list.size() <= i || this->legal_actions_list[i].empty() || this->legal_actions_list[i][0] == -1) { // sampled // discrete action space without action mask @@ -662,9 +677,11 @@ namespace tree { // TODO(pu): discrete action space std::vector c_legal_actions; - for (int i = 0; i < this->legal_actions_list.size(); ++i) + for (int j = 0; j < this->legal_actions_list[i].size(); ++j) { - CAction c_legal_action = CAction(legal_actions_list[i], 0); + std::vector legal_action; + legal_action.push_back(this->legal_actions_list[i][j]); + CAction c_legal_action = CAction(legal_action, 0); c_legal_actions.push_back(c_legal_action); } this->roots.push_back(CNode(0, c_legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space)); @@ -773,10 +790,10 @@ namespace tree for (int i = 0; i < this->root_num; ++i) { std::vector sampled_action; - sampled_action = this->roots[i].legal_actions; + sampled_action = this->roots[i].sampled_actions; std::vector > python_sampled_action; - for (int j = 0; j < this->roots[i].legal_actions.size(); ++j) + for (int j = 0; j < sampled_action.size(); ++j) { python_sampled_action.push_back(sampled_action[j].value); } diff --git a/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.h b/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.h index 533335591..999adde25 100644 --- a/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.h +++ b/lzero/mcts/ctree/ctree_sampled_efficientzero/lib/cnode.h @@ -48,6 +48,7 @@ namespace tree std::map children; std::vector legal_actions; + std::vector sampled_actions; CNode(); // sampled related core code @@ -120,4 +121,4 @@ namespace tree void cbatch_traverse(CRoots *roots, int pb_c_base, float pb_c_init, float discount_factor, tools::CMinMaxStatsList *min_max_stats_lst, CSearchResults &results, std::vector &virtual_to_play_batch, bool continuous_action_space); } -#endif \ No newline at end of file +#endif diff --git a/lzero/mcts/ctree/ctree_sampled_muzero/lib/cnode.cpp b/lzero/mcts/ctree/ctree_sampled_muzero/lib/cnode.cpp index 012dc1a48..2c2f6be43 100644 --- a/lzero/mcts/ctree/ctree_sampled_muzero/lib/cnode.cpp +++ b/lzero/mcts/ctree/ctree_sampled_muzero/lib/cnode.cpp @@ -362,52 +362,53 @@ namespace tree else { // discrete action space for sampled algo.. + std::vector candidate_actions; + if (this->legal_actions.empty()) + { + for (int i = 0; i < this->action_space_size; ++i) + { + candidate_actions.push_back(i); + } + } + else + { + for (int i = 0; i < this->legal_actions.size(); ++i) + { + candidate_actions.push_back(int(this->legal_actions[i].value[0])); + } + } - //======================================================== - // python code - //======================================================== - // if self.legal_actions is not None: - // # fisrt 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 = [] - // then empty the self.legal_actions - // prob = torch.softmax(torch.tensor(policy_logits), dim=-1) - // sampled_actions = torch.multinomial(prob, self.num_of_sampled_actions, replacement=False) - - //======================================================== - // TODO(pu): legal actions - //======================================================== - // std::vector policy_tmp; - // for (int i = 0; i < this->action_space_size; ++i) - // { - // policy_tmp.push_back(0.); - // } - // for (int i = 0; i < this->legal_actions.size(); ++i) - // { - // policy_tmp[this->legal_actions[i].value] = policy_logits[i]; - // } - // for (int i = 0; i < this->action_space_size; ++i) - // { - // policy_logits[i] = policy_tmp[i]; - // } - // std::cout << "position 3" << std::endl; + assert(candidate_actions.size() > 0); + assert(policy_logits.size() == this->action_space_size || policy_logits.size() == candidate_actions.size()); - // python code: legal_actions = [] - std::vector legal_actions; + std::vector candidate_logits; + for (int i = 0; i < candidate_actions.size(); ++i) + { + int action = candidate_actions[i]; + if (policy_logits.size() == this->action_space_size) + { + candidate_logits.push_back(policy_logits[action]); + } + else + { + candidate_logits.push_back(policy_logits[i]); + } + } - // python code: probs = softmax(policy_logits) + // python code: probs = softmax(candidate_logits) float logits_exp_sum = 0; - for (int i = 0; i < policy_logits.size(); ++i) + float max_logit = candidate_logits[0]; + for (int i = 1; i < candidate_logits.size(); ++i) + { + max_logit = std::max(max_logit, candidate_logits[i]); + } + for (int i = 0; i < candidate_logits.size(); ++i) { - logits_exp_sum += exp(policy_logits[i]); + logits_exp_sum += exp(candidate_logits[i] - max_logit); } - for (int i = 0; i < policy_logits.size(); ++i) + for (int i = 0; i < candidate_logits.size(); ++i) { - probs.push_back(exp(policy_logits[i]) / (logits_exp_sum + 1e-6)); + probs.push_back(exp(candidate_logits[i] - max_logit) / (logits_exp_sum + 1e-6)); } unsigned seed = std::chrono::system_clock::now().time_since_epoch().count(); @@ -444,10 +445,10 @@ namespace tree std::sort(disc_action_with_probs.begin(), disc_action_with_probs.end(), cmp); - // take the fist ``num_of_sampled_actions`` actions - for (int k = 0; k < num_of_sampled_actions; ++k) + int num_sampled = std::min(this->num_of_sampled_actions, int(candidate_actions.size())); + for (int k = 0; k < num_sampled; ++k) { - sampled_actions.push_back(disc_action_with_probs[k].first); + sampled_actions.push_back(candidate_actions[disc_action_with_probs[k].first]); // disc_action_with_probs[k].second is disturbed_probs // sampled_actions_probs.push_back(disc_action_with_probs[k].second); sampled_actions_probs.push_back(probs[disc_action_with_probs[k].first]); @@ -471,18 +472,20 @@ namespace tree disc_action_with_probs.clear(); // Empty the collection to prepare for the next sampling. } - float prior; - for (int i = 0; i < this->num_of_sampled_actions; ++i) + this->legal_actions.clear(); + if (this->continuous_action_space == true) { - - if (this->continuous_action_space == true) + for (int i = 0; i < this->num_of_sampled_actions; ++i) { CAction action = CAction(sampled_actions_after_tanh[i], 0); std::vector legal_actions; this->children[action.get_combined_hash()] = CNode(sampled_actions_log_probs_after_tanh[i], legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space); // only for muzero/efficient zero, not support alphazero this->legal_actions.push_back(action); } - else + } + else + { + for (int i = 0; i < sampled_actions.size(); ++i) { std::vector sampled_action_tmp; for (size_t iter = 0; iter < 1; iter++) @@ -513,7 +516,8 @@ namespace tree - noises: the vector of noises added to each child node. */ float noise, prior; - for (int i = 0; i < this->num_of_sampled_actions; ++i) + int noise_num = std::min(int(this->legal_actions.size()), int(noises.size())); + for (int i = 0; i < noise_num; ++i) { noise = noises[i]; CNode *child = this->get_child(this->legal_actions[i]); @@ -698,17 +702,17 @@ namespace tree for (int i = 0; i < this->root_num; ++i) { - if (this->continuous_action_space == true and this->legal_actions_list[0][0] == -1) + if (this->continuous_action_space == true) { // continous action space std::vector legal_actions; this->roots.push_back(CNode(0, legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space)); } - else if (this->continuous_action_space == false or this->legal_actions_list[0][0] == -1) + else if (this->legal_actions_list.size() <= i || this->legal_actions_list[i].empty() || this->legal_actions_list[i][0] == -1) { // sampled // discrete action space without action mask - std::vector legal_actions; + std::vector legal_actions; this->roots.push_back(CNode(0, legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space)); } @@ -716,9 +720,11 @@ namespace tree { // TODO(pu): discrete action space std::vector c_legal_actions; - for (int i = 0; i < this->legal_actions_list.size(); ++i) + for (int j = 0; j < this->legal_actions_list[i].size(); ++j) { - CAction c_legal_action = CAction(legal_actions_list[i], 0); + std::vector legal_action; + legal_action.push_back(this->legal_actions_list[i][j]); + CAction c_legal_action = CAction(legal_action, 0); c_legal_actions.push_back(c_legal_action); } this->roots.push_back(CNode(0, c_legal_actions, this->action_space_size, this->num_of_sampled_actions, this->continuous_action_space)); diff --git a/lzero/mcts/ptree/ptree_sez.py b/lzero/mcts/ptree/ptree_sez.py index 63d1dde2e..99112f70e 100644 --- a/lzero/mcts/ptree/ptree_sez.py +++ b/lzero/mcts/ptree/ptree_sez.py @@ -55,6 +55,7 @@ def __init__( self.value_prefix = 0.0 self.children = {} self.children_index = [] + self.sampled_actions = [] self.simulation_index = 0 self.batch_index = 0 @@ -113,35 +114,56 @@ def expand( log_prob = dist.log_prob(sampled_actions_before_tanh).unsqueeze(-1) log_prob = log_prob - torch.log(y).sum(-1, keepdim=True) self.legal_actions = [] + self.sampled_actions = [] for action_index in range(self.num_of_sampled_actions): - self.children[Action(sampled_actions[action_index].detach().cpu().numpy())] = Node( + action = Action(sampled_actions[action_index].detach().cpu().numpy()) + self.children[action] = Node( log_prob[action_index], action_space_size=self.action_space_size, num_of_sampled_actions=self.num_of_sampled_actions, continuous_action_space=self.continuous_action_space ) - self.legal_actions.append(Action(sampled_actions[action_index].detach().cpu().numpy())) + self.legal_actions.append(action) + self.sampled_actions.append(action) else: - 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 + candidate_actions = list(range(self.action_space_size)) if self.legal_actions is None else list(self.legal_actions) + if len(candidate_actions) == 0: + raise ValueError("legal_actions must not be empty when expanding a discrete Sampled EfficientZero node") + + policy_logits = torch.as_tensor(policy_logits, dtype=torch.float32) + candidate_actions_tensor = torch.as_tensor(candidate_actions, dtype=torch.long) + if policy_logits.numel() == self.action_space_size: + candidate_logits = policy_logits[candidate_actions_tensor] + elif policy_logits.numel() == len(candidate_actions): + candidate_logits = policy_logits + else: + raise ValueError( + f"policy_logits length {policy_logits.numel()} does not match action_space_size " + f"{self.action_space_size} or legal action count {len(candidate_actions)}" + ) + + prob = torch.softmax(candidate_logits, dim=-1) + num_sampled = min(self.num_of_sampled_actions, len(candidate_actions)) + sampled_action_indices = torch.multinomial(prob, num_sampled, replacement=False) + sampled_actions = candidate_actions_tensor[sampled_action_indices] + self.legal_actions = [] - prob = torch.softmax(torch.tensor(policy_logits), dim=-1) - sampled_actions = torch.multinomial(prob, self.num_of_sampled_actions, replacement=False) + self.sampled_actions = [] - for action_index in range(self.num_of_sampled_actions): - self.children[Action(sampled_actions[action_index].detach().cpu().numpy())] = Node( - prob[sampled_actions[action_index]], # + for action_index in range(num_sampled): + action = Action(sampled_actions[action_index].detach().cpu().numpy()) + self.children[action] = Node( + prob[sampled_action_indices[action_index]], action_space_size=self.action_space_size, num_of_sampled_actions=self.num_of_sampled_actions, continuous_action_space=self.continuous_action_space ) - self.legal_actions.append(Action(sampled_actions[action_index].detach().cpu().numpy())) + self.legal_actions.append(action) + self.sampled_actions.append(action) + + if self.sampled_actions: + self.sampled_actions.extend([self.sampled_actions[-1]] * (self.num_of_sampled_actions - num_sampled)) def add_exploration_noise_to_sample_distribution( self, exploration_fraction: float, noises: List[float], policy_logits: List[float] @@ -255,6 +277,7 @@ def get_children_distribution(self) -> List[Union[int, float]]: distribution[a] = child.visit_count # only take the visit counts distribution = [v for k, v in distribution.items()] + distribution.extend([0] * (len(self.sampled_actions) - len(distribution))) return distribution def get_child(self, action: Union[int, float]) -> "Node": @@ -459,7 +482,7 @@ def get_sampled_actions(self) -> List[List[Union[int, float]]]: # TODO(pu): root_sampled_actions bug in discere action space? sampled_actions = [] for i in range(self.root_num): - sampled_actions.append(self.roots[i].legal_actions) + sampled_actions.append(self.roots[i].sampled_actions) return sampled_actions diff --git a/lzero/mcts/tests/test_ptree_sez_discrete_sampling.py b/lzero/mcts/tests/test_ptree_sez_discrete_sampling.py new file mode 100644 index 000000000..84c6fbeb1 --- /dev/null +++ b/lzero/mcts/tests/test_ptree_sez_discrete_sampling.py @@ -0,0 +1,97 @@ +import numpy as np +import pytest +import torch + +from lzero.mcts.ptree.ptree_sez import Action, Node + + +def _action_value(action): + return int(np.asarray(action.value).item()) + + +@pytest.mark.unittest +def test_ptree_sez_samples_only_legal_actions_and_uses_action_indexed_logits(): + torch.manual_seed(0) + node = Node(0, legal_actions=[1, 4], action_space_size=5, num_of_sampled_actions=2) + + node.expand(-1, 0, 0, 0.0, [100.0, -10.0, 90.0, 80.0, 10.0]) + + sampled_actions = [_action_value(action) for action in node.legal_actions] + assert set(sampled_actions) == {1, 4} + assert len(node.children) == 2 + + action_1_prior = float(node.get_child(Action(np.array(1))).prior) + action_4_prior = float(node.get_child(Action(np.array(4))).prior) + assert action_4_prior > action_1_prior + + +@pytest.mark.unittest +def test_ptree_sez_pads_sampled_actions_and_distributions_when_legal_count_is_smaller_than_k(): + torch.manual_seed(0) + node = Node(0, legal_actions=[2], action_space_size=5, num_of_sampled_actions=4) + + node.expand(-1, 0, 0, 0.0, [0.0, 1.0, 2.0, 3.0, 4.0]) + + assert [_action_value(action) for action in node.legal_actions] == [2] + assert [_action_value(action) for action in node.sampled_actions] == [2, 2, 2, 2] + + node.add_exploration_noise(0.25, [0.25, 0.25, 0.25, 0.25]) + node.get_child(node.legal_actions[0]).visit_count = 3 + assert node.get_children_distribution() == [3, 0, 0, 0] + + +@pytest.mark.unittest +def test_ctree_sampled_efficientzero_discrete_sampling_respects_legal_actions(): + ezs_tree = pytest.importorskip("lzero.mcts.ctree.ctree_sampled_efficientzero.ezs_tree") + + roots = ezs_tree.Roots(2, [[1, 4], [2]], 5, 4, False) + roots.prepare( + 0.25, + [[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], + [0.0, 0.0], + [ + [100.0, -10.0, 90.0, 80.0, 10.0], + [0.0, 1.0, 2.0, 3.0, 4.0], + ], + [-1, -1], + ) + + sampled_actions = roots.get_sampled_actions() + distributions = roots.get_distributions() + + assert len(sampled_actions[0]) == 4 + assert len(distributions[0]) == 4 + assert set(int(action[0]) for action in sampled_actions[0]) == {1, 4} + assert len(set(int(action[0]) for action in sampled_actions[0][:2])) == 2 + + assert len(sampled_actions[1]) == 4 + assert len(distributions[1]) == 4 + assert [int(action[0]) for action in sampled_actions[1]] == [2, 2, 2, 2] + + +@pytest.mark.unittest +def test_ctree_sampled_muzero_discrete_sampling_respects_legal_actions(): + smz_tree = pytest.importorskip("lzero.mcts.ctree.ctree_sampled_muzero.smz_tree") + + roots = smz_tree.Roots(2, [[1, 4], [2]], 5, 4, False) + roots.prepare( + 0.25, + [[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], + [0.0, 0.0], + [ + [100.0, -10.0, 90.0, 80.0, 10.0], + [0.0, 1.0, 2.0, 3.0, 4.0], + ], + [-1, -1], + ) + + sampled_actions = roots.get_sampled_actions() + distributions = roots.get_distributions() + + assert len(sampled_actions[0]) == 2 + assert len(distributions[0]) == 2 + assert set(int(action[0]) for action in sampled_actions[0]) == {1, 4} + + assert len(sampled_actions[1]) == 1 + assert len(distributions[1]) == 1 + assert [int(action[0]) for action in sampled_actions[1]] == [2]