(observations, generations, rewards, ax=None, box_size=1.5, box_size_y=2, dt=25, color_bar=False)
| 37 | plt.plot([x0[i]-cart_size,x0[i]+cart_size], [y0[i], y0[i]], color=color, alpha=alpha * max_alpha) |
| 38 | |
| 39 | def plot_cartpole(observations, generations, rewards, ax=None, box_size=1.5, box_size_y=2, dt=25, color_bar=False): |
| 40 | if ax is None: |
| 41 | ax = plt.gca() |
| 42 | |
| 43 | for g, t in enumerate(generations): |
| 44 | ax.axhline(g * 2, color='black', alpha=0.05) |
| 45 | for i in range(len(observations[0])): |
| 46 | c = custom_cmap(rewards[t][i] / 500) |
| 47 | ang = observations[t][i][::dt, 2] |
| 48 | pos = observations[t][i][::dt, 0] / 4.8 |
| 49 | cartpole_plot(ang, pos, max_alpha=0.5, decay=2, y0_shift=g * box_size_y, box_size=box_size, color=c) |
| 50 | |
| 51 | x = np.arange(0, 501, 50) |
| 52 | x_corr = x / dt * box_size |
| 53 | ax.set_xticks(x_corr, x) |
| 54 | ax.set_yticks(np.arange(0, len(generations) * box_size_y, box_size_y), generations+1) |
| 55 | ax.set_ylim(-0.5, None) |
| 56 | ax.set_xlabel('time steps') |
| 57 | ax.set_ylabel('generation') |
| 58 | |
| 59 | if color_bar: |
| 60 | # Adding the horizontal color bar inside the plot using inset_axes |
| 61 | cbar_ax = inset_axes(ax, width="20%", height="3%", loc='lower right', |
| 62 | bbox_to_anchor=(0.05, 0.15, 0.9, 0.95), |
| 63 | bbox_transform=ax.transAxes, borderpad=0) |
| 64 | |
| 65 | # Correctly referencing the figure associated with ax |
| 66 | cbar = ax.figure.colorbar(plt.cm.ScalarMappable(cmap=custom_cmap), cax=cbar_ax, orientation='horizontal') |
| 67 | |
| 68 | # Remove color bar ticks |
| 69 | cbar.ax.set_xticks([]) |
| 70 | cbar.ax.set_yticks([]) |
| 71 | |
| 72 | # Set color bar label |
| 73 | cbar.set_label('reward') |
| 74 | |
| 75 | def prepare_reward(rewards): |
| 76 | # rewards.shape = [num_experiment, num_generation, num_population] |
no test coverage detected