(rewards)
| 73 | cbar.set_label('reward') |
| 74 | |
| 75 | def prepare_reward(rewards): |
| 76 | # rewards.shape = [num_experiment, num_generation, num_population] |
| 77 | # merge each experiment into one for each generation |
| 78 | if isinstance(rewards, list): |
| 79 | rewards = torch.stack(rewards) |
| 80 | |
| 81 | rewards = rewards.permute(1, 0, 2).reshape(rewards.shape[1], -1) |
| 82 | return rewards |
| 83 | |
| 84 | def range_plot(x, color=None, label=None): |
| 85 | print(f'{len(x)} experiments, (num_generation, num_population)={x[0].shape}') |