%matplotlib inline import pylab as plt from IPython.display import clear_output import numpy as np import asyncio class TSP(object): ''' 用 Q-Learning 求解 TSP 问题 作者 Surfer Zen @ https://www.zhihu.com/people/surfer-zen ''' def __init__(self, num_cities=15, map_size=(800.0, 600.0), alpha=2, beta=1, learning_rate=0.001, eps=0.1, ): ''' Args: num_cities (int): 城市数目 map_size (int, int): 地图尺寸（宽，高） alpha (float): 一个超参，值越大，越优先探索最近的点 beta (float): 一个超参，值越大，越优先探索可能导向总距离最优的点 learning_rate (float)： 学习率 eps (float): 探索率，值越大，探索性越强，但越难收敛 ''' self.num_cities =num_cities self.map_size = map_size self.alpha = alpha self.beta = beta self.eps = eps self.learning_rate = learning_rate self.cities = self.generate_cities() self.distances = self.get_dist_matrix() self.mean_distance = self.distances.mean() self.qualities = np.zeros([num_cities, num_cities]) self.normalizers = np.zeros(num_cities) self.best_path = None self.best_path_length = np.inf def generate_cities(self): ''' 随机生成城市（坐标） Returns: cities: [[x1, x2, x3...], [y1, y2, y3...]] 城市坐标 ''' max_width, max_height = self.map_size cities = np.random.random([2, self.num_cities]) \ * np.array([max_width, max_height]).reshape(2, -1) return cities def get_dist_matrix(self): ''' 根据城市坐标，计算距离矩阵 ''' dist_matrix = np.zeros([self.num_cities, self.num_cities]) for i in range(self.num_cities): for j in range(self.num_cities): if i == j: continue xi, xj = self.cities[0, i], self.cities[0, j] yi, yj = self.cities[1, i], self.cities[1, j] dist_matrix[i, j] = np.sqrt((xi-xj)**2 + (yi-yj)**2) return dist_matrix def rollout(self, start_city_id=None): ''' 从 start_city 出发，根据策略，在城市间游走，直到所有城市都走了一遍 ''' cities_visited = [] action_probs = [] if start_city_id is None: start_city_id = np.random.randint(self.num_cities) current_city_id = start_city_id cities_visited.append(current_city_id) while len(cities_visited) < self.num_cities: current_city_id, action_prob = self.choose_next_city(cities_visited) cities_visited.append(current_city_id) action_probs.append(action_prob) cities_visited.append(cities_visited[0]) action_probs.append(1.0) path_length = self.calc_path_length(cities_visited) if path_length < self.best_path_length: self.best_path = cities_visited self.best_path_length = path_length rewards = self.calc_path_rewards(cities_visited, path_length) return cities_visited, action_probs, rewards def choose_next_city(self, cities_visited): ''' 根据策略选择下一个城市 ''' current_city_id = cities_visited[-1] # 对 quality 取指数，计算 softmax 概率用 probabilities = np.exp(self.qualities[current_city_id]) # 将已经走过的城市概率设置为零 for city_visited in cities_visited: probabilities[city_visited] = 0 # 计算 softmax 概率 probabilities = probabilities/probabilities.sum() if np.random.random() < self.eps: # 以 eps 概率按softmax概率密度进行随机采样 next_city_id = np.random.choice(range(len(probabilities)), p=probabilities) else: # 以 (1 - eps) 概率选择当前最优策略 next_city_id = probabilities.argmax() # 计算当前决策/action 的概率 if probabilities.argmax() == next_city_id: action_prob = probabilities[next_city_id]*self.eps + (1-self.eps) else: action_prob = probabilities[next_city_id]*self.eps return next_city_id, action_prob def calc_path_rewards(self, path, path_length): ''' 计算给定路径的奖励/rewards Args: path (list[int]): 路径，每个元素代表城市的 id path_length (float): 路径长路 Returns: rewards: 每一步的奖励，总距离以及当前这一步的距离越大，奖励越小 ''' rewards = [] for fr, to in zip(path[:-1], path[1:]): dist = self.distances[fr, to] reward = (self.mean_distance/path_length)**self.beta reward = reward*(self.mean_distance/dist)**self.alpha rewards.append(np.log(reward)) return rewards def calc_path_length(self, path): ''' 计算路径长度 ''' path_length = 0 for fr, to in zip(path[:-1], path[1:]): path_length += self.distances[fr, to] return path_length def calc_updates_for_one_rollout(self, path, action_probs, rewards): ''' 对于给定的一次 rollout 的结果，计算其对应的 qualities 和 normalizers ''' qualities = [] normalizers = [] for fr, to, reward, action_prob in zip(path[:-1], path[1:], rewards, action_probs): log_action_probability = np.log(action_prob) qualities.append(- reward*log_action_probability) normalizers.append(- (reward + 1)*log_action_probability) return qualities, normalizers def update(self, path, new_qualities, new_normalizers): ''' 用渐近平均的思想，对 qualities 和 normalizers 进行更新 ''' lr = self.learning_rate for fr, to, new_quality, new_normalizer in zip( path[:-1], path[1:], new_qualities, new_normalizers): self.normalizers[fr] = (1-lr)*self.normalizers[fr] + lr*new_normalizer self.qualities[fr, to] = (1-lr)*self.qualities[fr, to] + lr*new_quality async def train_for_one_rollout(self, start_city_id): ''' 对一次 rollout 的结果进行训练的流程 ''' path, action_probs, rewards = self.rollout(start_city_id=start_city_id) new_qualities, new_normalizers = self.calc_updates_for_one_rollout(path, action_probs, rewards) self.update(path, new_qualities, new_normalizers) async def train_for_one_epoch(self): ''' 对一个 epoch 的结果进行训练的流程， 一个 epoch 对应于从每个 city 出发进行一次 rollout ''' tasks = [self.train_for_one_rollout(start_city_id) for start_city_id in range(self.num_cities)] await asyncio.gather(*tasks) async def train(self, num_epochs=1000, display=True): ''' 总训练流程 ''' for epoch in range(num_epochs): await self.train_for_one_epoch() if display: self.draw(epoch) def draw(self, epoch): ''' 绘图 ''' _ = plt.scatter(*self.cities) for fr, to in zip(self.best_path[:-1], self.best_path[1:]): x1, y1 = self.cities[:, fr] x2, y2 = self.cities[:, to] dx, dy = x2-x1, y2-y1 plt.arrow(x1, y1, dx, dy, width=0.01*min(self.map_size), edgecolor='orange', facecolor='white', animated=True, length_includes_head=True) nrs = np.exp(self.qualities) for i in range(self.num_cities): nrs[i, i] = 0 gap = np.abs(np.exp(self.normalizers) - nrs.sum(-1)).mean() plt.title(f'epoch {epoch}: path length = {self.best_path_length:.2f}, normalizer error = {gap:.3f}') plt.savefig('tsp.png') plt.show() clear_output(wait=True)