matplotlib 绘制轮廓图和 3D 图

contour([X, Y,] Z, [levels], **kwargs)

import numpy as np
import matplotlib.pyplot as plt


policy = np.loadtxt('p.txt')
policy = policy[21:42,:]

X = np.arange(0, 21)
Y = np.arange(0, 21)
X, Y = np.meshgrid(X, Y)

levels = range(-5, 6, 1)
CS = plt.contourf(X, Y, policy, levels)
cbar = plt.colorbar(CS)
cbar.ax.set_ylabel('action')
plt.title('Policy')
plt.xlabel("Cars at Location 2")
plt.ylabel("Cars at Location 1")
plt.show()

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

V = np.loadtxt('value.txt')
fig = plt.figure(figsize=(10, 8))
ax_surf = fig.gca(projection='3d')
ax_surf.set_position([0.1,0.15,0.7,0.7])
X, Y = np.meshgrid(np.arange(0, 21), np.arange(0, 21))
surf = ax_surf.plot_surface(X, Y, V, cmap=cm.coolwarm,
                                linewidth=0, antialiased=False)
ax_surf.set_xticks(np.arange(0,21,4).astype(int))
ax_surf.set_yticks(np.arange(0,21,4).astype(int))
ax_surf.set_xlabel('Number of cars at location 2')
ax_surf.set_ylabel('Number of cars at location 1')
ax_surf.set_title('Value function of optimal policy')
ax_color = fig.add_axes([0.85,0.25,0.03,0.5])
cbar = fig.colorbar(surf, cax=ax_color, 
                        orientation='vertical')