Spaces:
Sleeping
Sleeping
| import torch | |
| from easydict import EasyDict | |
| from lzero.policy import inverse_scalar_transform, select_action | |
| import numpy as np | |
| import random | |
| from lzero.mcts.tree_search.mcts_ptree import EfficientZeroMCTSPtree as MCTSPtree | |
| from lzero.mcts.tree_search.mcts_ctree import EfficientZeroMCTSCtree as MCTSCtree | |
| import time | |
| class MuZeroModelFake(torch.nn.Module): | |
| """ | |
| Overview: | |
| Fake MuZero model just for test EfficientZeroMCTSPtree. | |
| Interfaces: | |
| __init__, initial_inference, recurrent_inference | |
| """ | |
| def __init__(self, action_num): | |
| super().__init__() | |
| self.action_num = action_num | |
| def initial_inference(self, observation): | |
| encoded_state = observation | |
| batch_size = encoded_state.shape[0] | |
| value = torch.zeros(size=(batch_size, 601)) | |
| value_prefix = [0. for _ in range(batch_size)] | |
| policy_logits = torch.zeros(size=(batch_size, self.action_num)) | |
| latent_state = torch.zeros(size=(batch_size, 12, 3, 3)) | |
| reward_hidden_state_state = (torch.zeros(size=(1, batch_size, 16)), torch.zeros(size=(1, batch_size, 16))) | |
| output = { | |
| 'searched_value': value, | |
| 'value_prefix': value_prefix, | |
| 'policy_logits': policy_logits, | |
| 'latent_state': latent_state, | |
| 'reward_hidden_state': reward_hidden_state_state | |
| } | |
| return EasyDict(output) | |
| def recurrent_inference(self, hidden_states, reward_hidden_states, actions): | |
| batch_size = hidden_states.shape[0] | |
| latent_state = torch.zeros(size=(batch_size, 12, 3, 3)) | |
| reward_hidden_state_state = (torch.zeros(size=(1, batch_size, 16)), torch.zeros(size=(1, batch_size, 16))) | |
| value = torch.zeros(size=(batch_size, 601)) | |
| value_prefix = torch.zeros(size=(batch_size, 601)) | |
| policy_logits = torch.zeros(size=(batch_size, self.action_num)) | |
| output = { | |
| 'searched_value': value, | |
| 'value_prefix': value_prefix, | |
| 'policy_logits': policy_logits, | |
| 'latent_state': latent_state, | |
| 'reward_hidden_state': reward_hidden_state_state | |
| } | |
| return EasyDict(output) | |
| def ptree_func(policy_config, num_simulations): | |
| """ | |
| Overview: | |
| Search on the tree of the Python implementation and record the time spent at different stages. | |
| Arguments: | |
| - policy_config: config of game. | |
| - num_simulations: Number of simulations. | |
| Returns: | |
| - build_time: Type builds take time. | |
| - prepare_time: time for prepare. | |
| - search_time. | |
| - total_time. | |
| """ | |
| batch_size = env_nums = policy_config.batch_size | |
| action_space_size = policy_config.action_space_size | |
| build_time = [] | |
| prepare_time = [] | |
| search_time = [] | |
| total_time = [] | |
| for n_s in num_simulations: | |
| t0 = time.time() | |
| model = MuZeroModelFake(action_num=action_space_size) | |
| stack_obs = torch.zeros( | |
| size=( | |
| batch_size, | |
| n_s, | |
| ), dtype=torch.float | |
| ) | |
| policy_config.num_simulations = n_s | |
| network_output = model.initial_inference(stack_obs.float()) | |
| latent_state_roots = network_output['latent_state'] | |
| reward_hidden_state_state = network_output['reward_hidden_state'] | |
| pred_values_pool = network_output['value'] | |
| value_prefix_pool = network_output['value_prefix'] | |
| policy_logits_pool = network_output['policy_logits'] | |
| # network output process | |
| pred_values_pool = inverse_scalar_transform(pred_values_pool, | |
| policy_config.model.support_scale).detach().cpu().numpy() | |
| latent_state_roots = latent_state_roots.detach().cpu().numpy() | |
| reward_hidden_state_state = ( | |
| reward_hidden_state_state[0].detach().cpu().numpy(), reward_hidden_state_state[1].detach().cpu().numpy() | |
| ) | |
| policy_logits_pool = policy_logits_pool.detach().cpu().numpy().tolist() | |
| action_mask = [[random.randint(0, 1) for _ in range(action_space_size)] for _ in range(env_nums)] | |
| assert len(action_mask) == batch_size | |
| assert len(action_mask[0]) == action_space_size | |
| action_num = [int(np.array(action_mask[i]).sum()) for i in range(env_nums)] | |
| legal_actions_list = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(env_nums)] | |
| to_play = [np.random.randint(1, 3) for i in range(env_nums)] | |
| assert len(to_play) == batch_size | |
| # ============================================ptree=====================================# | |
| for i in range(env_nums): | |
| assert action_num[i] == len(legal_actions_list[i]) | |
| t1 = time.time() | |
| roots = MCTSPtree.roots(env_nums, legal_actions_list) | |
| build_time.append(time.time() - t1) | |
| noises = [ | |
| np.random.dirichlet([policy_config.root_dirichlet_alpha] * int(sum(action_mask[j])) | |
| ).astype(np.float32).tolist() for j in range(env_nums) | |
| ] | |
| t1 = time.time() | |
| roots.prepare(policy_config.root_noise_weight, noises, value_prefix_pool, policy_logits_pool, to_play) | |
| prepare_time.append(time.time() - t1) | |
| t1 = time.time() | |
| MCTSPtree(policy_config).search(roots, model, latent_state_roots, reward_hidden_state_state, to_play) | |
| search_time.append(time.time() - t1) | |
| total_time.append(time.time() - t0) | |
| roots_distributions = roots.get_distributions() | |
| roots_values = roots.get_values() | |
| assert len(roots_values) == env_nums | |
| assert len(roots_values) == env_nums | |
| for i in range(env_nums): | |
| assert len(roots_distributions[i]) == action_num[i] | |
| temperature = [1 for _ in range(env_nums)] | |
| for i in range(env_nums): | |
| distributions = roots_distributions[i] | |
| action_index, visit_count_distribution_entropy = select_action( | |
| distributions, temperature=temperature[i], deterministic=False | |
| ) | |
| action = np.where(np.array(action_mask[i]) == 1.0)[0][action_index] | |
| assert action_index < action_num[i] | |
| assert action == legal_actions_list[i][action_index] | |
| print('\n action_index={}, legal_action={}, action={}'.format(action_index, legal_actions_list[i], action)) | |
| return build_time, prepare_time, search_time, total_time | |
| def ctree_func(policy_config, num_simulations): | |
| """ | |
| Overview: | |
| Search on the tree of the C++ implementation and record the time spent at different stages. | |
| Arguments: | |
| - policy_config: config of game. | |
| - num_simulations: Number of simulations. | |
| Returns: | |
| - build_time: Type builds take time. | |
| - prepare_time: time for prepare. | |
| - search_time. | |
| - total_time. | |
| """ | |
| batch_size = env_nums = policy_config.batch_size | |
| action_space_size = policy_config.action_space_size | |
| build_time = [] | |
| prepare_time = [] | |
| search_time = [] | |
| total_time = [] | |
| for n_s in num_simulations: | |
| t0 = time.time() | |
| model = MuZeroModelFake(action_num=action_space_size) | |
| stack_obs = torch.zeros( | |
| size=( | |
| batch_size, | |
| n_s, | |
| ), dtype=torch.float | |
| ) | |
| policy_config.num_simulations = n_s | |
| network_output = model.initial_inference(stack_obs.float()) | |
| latent_state_roots = network_output['latent_state'] | |
| reward_hidden_state_state = network_output['reward_hidden_state'] | |
| pred_values_pool = network_output['value'] | |
| value_prefix_pool = network_output['value_prefix'] | |
| policy_logits_pool = network_output['policy_logits'] | |
| # network output process | |
| pred_values_pool = inverse_scalar_transform(pred_values_pool, | |
| policy_config.model.support_scale).detach().cpu().numpy() | |
| latent_state_roots = latent_state_roots.detach().cpu().numpy() | |
| reward_hidden_state_state = ( | |
| reward_hidden_state_state[0].detach().cpu().numpy(), reward_hidden_state_state[1].detach().cpu().numpy() | |
| ) | |
| policy_logits_pool = policy_logits_pool.detach().cpu().numpy().tolist() | |
| action_mask = [[random.randint(0, 1) for _ in range(action_space_size)] for _ in range(env_nums)] | |
| assert len(action_mask) == batch_size | |
| assert len(action_mask[0]) == action_space_size | |
| action_num = [int(np.array(action_mask[i]).sum()) for i in range(env_nums)] | |
| legal_actions_list = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(env_nums)] | |
| to_play = [np.random.randint(1, 3) for i in range(env_nums)] | |
| assert len(to_play) == batch_size | |
| # ============================================ctree=====================================# | |
| for i in range(env_nums): | |
| assert action_num[i] == len(legal_actions_list[i]) | |
| t1 = time.time() | |
| roots = MCTSCtree.roots(env_nums, legal_actions_list) | |
| build_time.append(time.time() - t1) | |
| noises = [ | |
| np.random.dirichlet([policy_config.root_dirichlet_alpha] * int(sum(action_mask[j])) | |
| ).astype(np.float32).tolist() for j in range(env_nums) | |
| ] | |
| t1 = time.time() | |
| roots.prepare(policy_config.root_noise_weight, noises, value_prefix_pool, policy_logits_pool, to_play) | |
| prepare_time.append(time.time() - t1) | |
| t1 = time.time() | |
| MCTSCtree(policy_config).search(roots, model, latent_state_roots, reward_hidden_state_state, to_play) | |
| search_time.append(time.time() - t1) | |
| total_time.append(time.time() - t0) | |
| roots_distributions = roots.get_distributions() | |
| roots_values = roots.get_values() | |
| assert len(roots_values) == env_nums | |
| assert len(roots_values) == env_nums | |
| for i in range(env_nums): | |
| assert len(roots_distributions[i]) == action_num[i] | |
| temperature = [1 for _ in range(env_nums)] | |
| for i in range(env_nums): | |
| distributions = roots_distributions[i] | |
| action_index, visit_count_distribution_entropy = select_action( | |
| distributions, temperature=temperature[i], deterministic=False | |
| ) | |
| action = np.where(np.array(action_mask[i]) == 1.0)[0][action_index] | |
| assert action_index < action_num[i] | |
| assert action == legal_actions_list[i][action_index] | |
| print('\n action_index={}, legal_action={}, action={}'.format(action_index, legal_actions_list[i], action)) | |
| return build_time, prepare_time, search_time, total_time | |
| def plot(ctree_time, ptree_time, iters, label): | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from matplotlib import pyplot | |
| plt.style.use('seaborn-whitegrid') | |
| palette = pyplot.get_cmap('Set1') | |
| font1 = { | |
| 'family': 'Times New Roman', | |
| 'weight': 'normal', | |
| 'size': 18, | |
| } | |
| plt.figure(figsize=(20, 10)) | |
| # ctree | |
| color = palette(0) | |
| avg = np.mean(ctree_time, axis=0) | |
| std = np.std(ctree_time, axis=0) | |
| r1 = list(map(lambda x: x[0] - x[1], zip(avg, std))) | |
| r2 = list(map(lambda x: x[0] + x[1], zip(avg, std))) | |
| plt.plot(iters, avg, color=color, label="ctree", linewidth=3.0) | |
| plt.fill_between(iters, r1, r2, color=color, alpha=0.2) | |
| # ptree | |
| ptree_time = np.array(ptree_time) | |
| color = palette(1) | |
| avg = np.mean(ptree_time, axis=0) | |
| std = np.std(ptree_time, axis=0) | |
| r1 = list(map(lambda x: x[0] - x[1], zip(avg, std))) | |
| r2 = list(map(lambda x: x[0] + x[1], zip(avg, std))) | |
| plt.plot(iters, avg, color=color, label="ptree", linewidth=3.0) | |
| plt.fill_between(iters, r1, r2, color=color, alpha=0.2) | |
| plt.legend(loc='lower right', prop=font1) | |
| plt.title('{}'.format(label)) | |
| plt.xlabel('simulations', fontsize=22) | |
| plt.ylabel('time', fontsize=22) | |
| plt.savefig('{}-time.png'.format(label)) | |
| if __name__ == "__main__": | |
| # cProfile.run("ctree_func()", filename="ctree_result.out", sort="cumulative") | |
| # cProfile.run("ptree_func()", filename="ptree_result.out", sort="cumulative") | |
| policy_config = EasyDict( | |
| dict( | |
| lstm_horizon_len=5, | |
| model=dict( | |
| support_scale=300, | |
| categorical_distribution=True, | |
| ), | |
| action_space_size=100, | |
| num_simulations=100, | |
| batch_size=512, | |
| pb_c_base=1, | |
| pb_c_init=1, | |
| discount_factor=0.9, | |
| root_dirichlet_alpha=0.3, | |
| root_noise_weight=0.2, | |
| dirichlet_alpha=0.3, | |
| exploration_fraction=1, | |
| device='cpu', | |
| value_delta_max=0.01, | |
| ) | |
| ) | |
| ACTION_SPCAE_SIZE = [16, 50] | |
| BATCH_SIZE = [8, 64, 512] | |
| NUM_SIMULATIONS = [i for i in range(20, 200, 20)] | |
| # ACTION_SPCAE_SIZE = [50] | |
| # BATCH_SIZE = [512] | |
| # NUM_SIMULATIONS = [i for i in range(10, 50, 10)] | |
| for action_space_size in ACTION_SPCAE_SIZE: | |
| for batch_size in BATCH_SIZE: | |
| policy_config.batch_size = batch_size | |
| policy_config.action_space_size = action_space_size | |
| ctree_build_time = [] | |
| ctree_prepare_time = [] | |
| ctree_search_time = [] | |
| ptree_build_time = [] | |
| ptree_prepare_time = [] | |
| ptree_search_time = [] | |
| ctree_total_time = [] | |
| ptree_total_time = [] | |
| num_simulations = NUM_SIMULATIONS | |
| for i in range(3): | |
| build_time, prepare_time, search_time, total_time = ctree_func( | |
| policy_config, num_simulations=num_simulations | |
| ) | |
| ctree_build_time.append(build_time) | |
| ctree_prepare_time.append(prepare_time) | |
| ctree_search_time.append(search_time) | |
| ctree_total_time.append(total_time) | |
| for i in range(3): | |
| build_time, prepare_time, search_time, total_time = ptree_func( | |
| policy_config, num_simulations=num_simulations | |
| ) | |
| ptree_build_time.append(build_time) | |
| ptree_prepare_time.append(prepare_time) | |
| ptree_search_time.append(search_time) | |
| ptree_total_time.append(total_time) | |
| label = 'action_space_size_{}_batch_size_{}'.format(action_space_size, batch_size) | |
| plot(ctree_build_time, ptree_build_time, iters=num_simulations, label=label + '_bulid_time') | |
| plot(ctree_prepare_time, ptree_prepare_time, iters=num_simulations, label=label + '_prepare_time') | |
| plot(ctree_search_time, ptree_search_time, iters=num_simulations, label=label + '_search_time') | |
| plot(ctree_total_time, ptree_total_time, iters=num_simulations, label=label + '_total_time') | |