| 22 | |
| 23 | |
| 24 | def cartpole_plot(angles, positions, box_size=1, y0_shift=0, pole_length=1, cart_size=0.1, ang_scale=2.4, max_alpha=1, decay=10, color='black'): |
| 25 | total_time = len(angles) |
| 26 | x0 = torch.linspace(0, total_time * box_size, total_time) |
| 27 | x0 = x0 + positions |
| 28 | y0 = x0 * 0 + y0_shift |
| 29 | |
| 30 | x1 = x0 + torch.sin(angles * ang_scale) * pole_length |
| 31 | y1 = y0 + torch.cos(angles * ang_scale) * pole_length |
| 32 | |
| 33 | alpha = 1 |
| 34 | i = len(x0) - 1 |
| 35 | plt.arrow(x0[i], y0[i], x1[i] - x0[i], y1[i] - y0[i], head_width=0.0, head_length=0.0, alpha=alpha * max_alpha, color=color) |
| 36 | # add a line to represent the cart |
| 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: |