|
| 1 | +#!/usr/bin/env python |
| 2 | +import time |
| 3 | +import numpy as np |
| 4 | +from scipy import linalg |
| 5 | +from matplotlib import pyplot as plt |
| 6 | + |
| 7 | +class PrioritySet: |
| 8 | + def __init__(self): |
| 9 | + self._item_dict = {} |
| 10 | + |
| 11 | + def push(self, item, key): |
| 12 | + assert(item not in self._item_dict) |
| 13 | + self._item_dict[item] = key |
| 14 | + |
| 15 | + def pop(self): |
| 16 | + min_item, min_key = self.min_item() |
| 17 | + del self._item_dict[min_item] |
| 18 | + return min_item |
| 19 | + |
| 20 | + def min_item(self): |
| 21 | + min_key = None |
| 22 | + min_item = None |
| 23 | + for item in self._item_dict: |
| 24 | + key = self._item_dict[item] |
| 25 | + if min_key is None or key < min_key: |
| 26 | + min_key = key |
| 27 | + min_item = item |
| 28 | + return min_item, min_key |
| 29 | + |
| 30 | + def remove(self, item): |
| 31 | + del self._item_dict[item] |
| 32 | + |
| 33 | + def contains(self, item): |
| 34 | + return (item in self._item_dict) |
| 35 | + |
| 36 | +class SimpleProblem: |
| 37 | + def __init__(self, states, actions_list): |
| 38 | + # array of states. indices are the state ids |
| 39 | + self._states = states |
| 40 | + |
| 41 | + self._actions = {s:{} for s in range(len(self._states))} |
| 42 | + self._actions_rev = {s:{} for s in range(len(self._states))} |
| 43 | + for s_from, s_to, cost in actions_list: |
| 44 | + self._actions[s_from][s_to] = cost |
| 45 | + self._actions_rev[s_to][s_from] = cost |
| 46 | + |
| 47 | + # set of actions whose costs have changed since problem creation |
| 48 | + self._disabled_actions = {} |
| 49 | + self._disabled_actions_rev = {} |
| 50 | + self._changed_states = set() |
| 51 | + |
| 52 | + def num_states(self): |
| 53 | + return len(self._states) |
| 54 | + |
| 55 | + def succs(self, s): |
| 56 | + ''' |
| 57 | + Return dictionary whose keys are sucessor states and |
| 58 | + values are costs. |
| 59 | + ''' |
| 60 | + return self._actions[s] |
| 61 | + |
| 62 | + def preds(self, s): |
| 63 | + ''' |
| 64 | + Return dictionary whose keys are predeccessor states and |
| 65 | + values are costs. |
| 66 | + ''' |
| 67 | + return self._actions_rev[s] |
| 68 | + |
| 69 | + def h(self, s1, s2): |
| 70 | + return linalg.norm(self._states[s1] - self._states[s2]) |
| 71 | + |
| 72 | + def disable_state(self, s): |
| 73 | + # mark this state and all neighbors as changed |
| 74 | + self._changed_states.add(s) |
| 75 | + self._changed_states.update(self._actions[s].keys()) |
| 76 | + |
| 77 | + # save the actions so we can restore them later |
| 78 | + self._disabled_actions[s] = self._actions[s].copy() |
| 79 | + self._disabled_actions_rev[s] = self._actions[s].copy() |
| 80 | + |
| 81 | + # delete all actions to/from this state |
| 82 | + self._actions[s] = {} |
| 83 | + self._actions_rev[s] = {} |
| 84 | + |
| 85 | + def get_changed_states(self): |
| 86 | + return self._changed_states |
| 87 | + |
| 88 | +class DStarLite: |
| 89 | + ''' |
| 90 | + States are represented as python integers. |
| 91 | + ''' |
| 92 | + def __init__(self, problem, s_start, s_goal): |
| 93 | + self.problem = problem |
| 94 | + self.s_start = s_start |
| 95 | + self.s_goal = s_goal |
| 96 | + |
| 97 | + self.g = np.array([np.inf] * self.problem.num_states()) |
| 98 | + self.rhs = np.array([np.inf] * self.problem.num_states()) |
| 99 | + self.rhs[s_goal] = 0 |
| 100 | + |
| 101 | + self.open = PrioritySet() |
| 102 | + self.open.push(s_goal, self.key(s_goal)) |
| 103 | + |
| 104 | + # keeps track of how many vertices we've updated |
| 105 | + self.num_updates = 0 |
| 106 | + |
| 107 | + def key(self, s): |
| 108 | + ''' Calculate key for a state ''' |
| 109 | + return (min(self.g[s], self.rhs[s]) + self.problem.h(self.s_start, s), |
| 110 | + min(self.g[s], self.rhs[s])) |
| 111 | + |
| 112 | + def update_vertex(self, s): |
| 113 | + self.num_updates += 1 |
| 114 | + |
| 115 | + if s != self.s_goal: |
| 116 | + succs = self.problem.succs(s) |
| 117 | + self.rhs[s] = min([self.g[s_next] + succs[s_next] for s_next in succs]) |
| 118 | + |
| 119 | + if self.open.contains(s): |
| 120 | + self.open.remove(s) |
| 121 | + |
| 122 | + if self.g[s] != self.rhs[s]: |
| 123 | + self.open.push(s, self.key(s)) |
| 124 | + |
| 125 | + def compute_shortest_path(self): |
| 126 | + while True: |
| 127 | + min_item, min_key = self.open.min_item() |
| 128 | + #print '' |
| 129 | + #print 'Top: %d, %s' % (min_item, str(min_key)) |
| 130 | + #print ' g:', self.g |
| 131 | + #print 'rhs:', self.rhs |
| 132 | + #print self.open._item_dict |
| 133 | + |
| 134 | + if (min_item is None) or (min_key >= self.key(self.s_start)) and (self.rhs[self.s_start] == self.g[self.s_start]): |
| 135 | + break |
| 136 | + |
| 137 | + k_old = min_key |
| 138 | + u = self.open.pop() |
| 139 | + # print 'Pop %d' % u |
| 140 | + |
| 141 | + if k_old < self.key(u): |
| 142 | + # print ' Push u with updated key' |
| 143 | + self.open.push(u, self.key(u)) |
| 144 | + elif self.g[u] > self.rhs[u]: |
| 145 | + # print 'Update all preds' |
| 146 | + self.g[u] = self.rhs[u] |
| 147 | + for s_pred in self.problem.preds(u): |
| 148 | + self.update_vertex(s_pred) |
| 149 | + else: |
| 150 | + # print 'Update u and all preds' |
| 151 | + self.g[u] = np.inf |
| 152 | + for s_pred in self.problem.preds(u): |
| 153 | + self.update_vertex(s_pred) |
| 154 | + self.update_vertex(u) |
| 155 | + |
| 156 | + def update_start_state(self, s_start): |
| 157 | + self.s_start = s_start |
| 158 | + |
| 159 | + # update the keys for the open set (the hueristic |
| 160 | + # has changed since the start state has changed) |
| 161 | + for s in self.open._item_dict: |
| 162 | + self.open._item_dict[s] = self.key(s) |
| 163 | + |
| 164 | + #self.update_vertex(s_start) |
| 165 | + #self.update_vertex(s_goal) |
| 166 | + #self.g[s_goal] = np.inf |
| 167 | + |
| 168 | + # should call update_vertex here on and vertices with changed edge costs |
| 169 | + |
| 170 | +def plot_g(points, g): |
| 171 | + ''' Plot distance from goal/start using colors''' |
| 172 | + expanded_states = ~(g == np.inf) |
| 173 | + |
| 174 | + # plot the expanded states |
| 175 | + plt.scatter(points[expanded_states][:,0], points[expanded_states][:,1], s=200, marker='o', cmap=plt.cm.jet, |
| 176 | + c=g[expanded_states]) |
| 177 | + |
| 178 | + # plot the non-expanded states in black |
| 179 | + plt.scatter(points[~expanded_states][:,0], points[~expanded_states][:,1], s=200, marker='o', cmap=plt.cm.jet, |
| 180 | + c='black') |
| 181 | + |
| 182 | + # write the state on each point |
| 183 | + for s, x in enumerate(points): |
| 184 | + plt.annotate('%d' % s, x, color='m') |
| 185 | + |
| 186 | +def random_graph_2d(n, conn_dist): |
| 187 | + points = np.random.random((n, 2)) |
| 188 | + actions_list = [] |
| 189 | + for ii in range(len(points)): |
| 190 | + for jj in range(len(points)): |
| 191 | + d = linalg.norm(points[ii] - points[jj]) |
| 192 | + if d < conn_dist: |
| 193 | + actions_list.append((ii, jj, d)) |
| 194 | + return points, actions_list |
| 195 | + |
| 196 | +def regular_graph_2d(n, dx, connectivity=4): |
| 197 | + ''' Creates an n x n grid. ''' |
| 198 | + if connectivity == 4: |
| 199 | + conn_dist = dx*1.01 |
| 200 | + elif connectivity == 8: |
| 201 | + conn_dist = dx*1.5 |
| 202 | + else: |
| 203 | + raise ValueError('Invalid value for connectivity parameter') |
| 204 | + |
| 205 | + points = [] |
| 206 | + for ii in range(n): |
| 207 | + for jj in range(n): |
| 208 | + points.append((ii*dx, jj*dx)) |
| 209 | + points = np.array(points) |
| 210 | + |
| 211 | + actions_list = [] |
| 212 | + for ii in range(len(points)): |
| 213 | + for jj in range(len(points)): |
| 214 | + if ii == jj: |
| 215 | + continue |
| 216 | + d = linalg.norm(points[ii] - points[jj]) |
| 217 | + if d < dx * conn_dist: |
| 218 | + actions_list.append((ii, jj, d)) |
| 219 | + return points, actions_list |
| 220 | + |
| 221 | + |
| 222 | +n = 200 |
| 223 | +conn_dist = 0.1 |
| 224 | + |
| 225 | +#points, actions_list = random_graph_2d(n, conn_dist) |
| 226 | +points, actions_list = regular_graph_2d(int(n**0.5), 1.0, connectivity=8) |
| 227 | + |
| 228 | +s_start = 0 |
| 229 | +s_goal = 195 |
| 230 | + |
| 231 | +problem = SimpleProblem(points, actions_list) |
| 232 | +planner = DStarLite(problem, s_start, s_goal) |
| 233 | +while True: |
| 234 | + print '' |
| 235 | + print 'Planning' |
| 236 | + |
| 237 | + t0 = time.time() |
| 238 | + planner.num_updates = 0 |
| 239 | + planner.compute_shortest_path() |
| 240 | + t1 = time.time() |
| 241 | + print 'Updated %d/%d vertices in %fs' % (planner.num_updates, len(points), t1-t0) |
| 242 | + |
| 243 | + # plot connections in the graph |
| 244 | + for s_from, s_to, cost in actions_list: |
| 245 | + endpoints = np.array([points[s_from], points[s_to]]) |
| 246 | + plt.plot(endpoints[:,0], endpoints[:,1], c='b') |
| 247 | + plt.plot([points[s_start][0]], [points[s_start][1]], 'rx', markersize=60) |
| 248 | + plt.plot([points[s_goal][0]], [points[s_goal][1]], 'gx', markersize=60) |
| 249 | + plot_g(points, planner.g) |
| 250 | + plt.show() |
| 251 | + |
| 252 | + |
| 253 | + # ask the user where to move |
| 254 | + input_str = raw_input('Next state:') |
| 255 | + |
| 256 | + toks = input_str.split() |
| 257 | + if len(toks) == 0: |
| 258 | + print 'Bad input' |
| 259 | + elif toks[0] == 'q': |
| 260 | + # quit |
| 261 | + break |
| 262 | + elif toks[0] == 'm': |
| 263 | + # move to specified state |
| 264 | + s_best = int(toks[1]) |
| 265 | + print 'Moving to %d' % s_best |
| 266 | + elif toks[0] == 't': |
| 267 | + print 'Toggling', toks[1:] |
| 268 | + # Todo: toggle states (disable/enable actions |
| 269 | + |
| 270 | + plt.cla() |
| 271 | + |
| 272 | + |
| 273 | + if 0: |
| 274 | + # move to the best neighbor vertex |
| 275 | + s_best, cost_best = None, np.inf |
| 276 | + succs = problem.succs(s_start) |
| 277 | + for s_next in succs: |
| 278 | + cost_next = succs[s_next] + planner.g[s_next] |
| 279 | + if cost_next < cost_best: |
| 280 | + cost_best = cost_next |
| 281 | + s_best = s_next |
| 282 | + if s_best is None: |
| 283 | + raise ValueError('No valid move') |
| 284 | + |
| 285 | + s_start = s_best |
| 286 | + planner.update_start_state(s_start) |
| 287 | + |
| 288 | + |
0 commit comments