| import time | |
| import torch | |
| from hpc_rll.origin.gae import gae, gae_data | |
| from hpc_rll.rl_utils.gae import GAE | |
| from testbase import mean_relative_error, times | |
| assert torch.cuda.is_available() | |
| use_cuda = True | |
| T = 1024 | |
| B = 64 | |
| def gae_val(): | |
| value = torch.randn(T + 1, B) | |
| reward = torch.randn(T, B) | |
| hpc_gae = GAE(T, B) | |
| if use_cuda: | |
| value = value.cuda() | |
| reward = reward.cuda() | |
| hpc_gae = hpc_gae.cuda() | |
| ori_adv = gae(gae_data(value, reward)) | |
| hpc_adv = hpc_gae(value, reward) | |
| if use_cuda: | |
| torch.cuda.synchronize() | |
| mre = mean_relative_error( | |
| torch.flatten(ori_adv).cpu().detach().numpy(), | |
| torch.flatten(hpc_adv).cpu().detach().numpy() | |
| ) | |
| print("gae mean_relative_error: " + str(mre)) | |
| def gae_perf(): | |
| value = torch.randn(T + 1, B) | |
| reward = torch.randn(T, B) | |
| hpc_gae = GAE(T, B) | |
| if use_cuda: | |
| value = value.cuda() | |
| reward = reward.cuda() | |
| hpc_gae = hpc_gae.cuda() | |
| for i in range(times): | |
| t = time.time() | |
| adv = gae(gae_data(value, reward)) | |
| if use_cuda: | |
| torch.cuda.synchronize() | |
| print('epoch: {}, original gae cost time: {}'.format(i, time.time() - t)) | |
| for i in range(times): | |
| t = time.time() | |
| hpc_adv = hpc_gae(value, reward) | |
| if use_cuda: | |
| torch.cuda.synchronize() | |
| print('epoch: {}, hpc gae cost time: {}'.format(i, time.time() - t)) | |
| if __name__ == '__main__': | |
| print("target problem: T = {}, B = {}".format(T, B)) | |
| print("================run gae validation test================") | |
| gae_val() | |
| print("================run gae performance test================") | |
| gae_perf() | |