def update(self, new_node: Node) -> None:
dist_to_end = new_node.distance_to(self._cfg.end_node)
if dist_to_end < self.global_closest_dist:
self.global_closest_dist = dist_to_end
self.global_closest = new_node
print(f'[{self.step_counter}] (new node / closest global) = '
f'({dist_to_end:.2f} / global {self.global_closest_dist:.2f})')
if new_node.distance_to(self._cfg.end_node) < self._cfg.eps:
self.running = False
print('Target found or steps done!')
plot_tree(self.ax,
self.rr_tree,
self._cfg.grid_size,
self._cfg.start_node,
self._cfg.end_node,
obstacles=self.obstacles,
reached_target=True,
last_node=new_node,
fast_plot=self._cfg.fast_plot)
input()
elif self.step_counter == self._cfg.max_steps:
print('Max steps reached - target not found.')
self.running = False
input()
else:
self.step_counter += 1
start_time = time.time()
plot_tree(self.ax,
self.rr_tree,
self._cfg.grid_size,
self._cfg.start_node,
self._cfg.end_node,
obstacles=self.obstacles,
fast_plot=self._cfg.fast_plot)
print(f'\tplotting took {time.time() - start_time:.2f}')
class RrtConfig:
grid_size: int = 100
start_node: Node = Node(Pose(10, 10))
end_node: Node = Node(Pose(90, 65))
eps: float = 3
max_steps: int = 1000
clamp_dist: float = 3
fast_plot: bool = True
def plot_tree(rr_tree: Node, end_node: Node, reached_target: bool = False) -> None:
global app
plot_start = time.time()
nodes = rr_tree.to_array()
connections = rr_tree.adjacency_nodes()
# cmap = pg.ColorMap(np.array([0, 100]), np.array([pg.mkColor('r'), pg.mkColor('b')]))
# distances = [end_node.distance_to(node) for node in rr_tree.traverse()]
# print(distances)
# brushes = [pg.mkBrush(col) for col in cmap.map(distances)]
graph_item.setData(pos=nodes, adj=np.array(connections))
app.processEvents()
print(f'\tPlotting took {time.time() - plot_start:.2f}s')
def insert_new_node(nearest_node: Node, new_node: Node,
clamp_distance: float) -> Node:
"""If a nearest node has been found, insert the new node into the tree appropriately and also return the Node."""
distance_to_new = new_node.distance_to(nearest_node)
if distance_to_new > clamp_distance:
dist_vec = Pose.normalize(
Pose.subtract(nearest_node.pose, new_node.pose))
new_pose = Pose.add(
nearest_node.pose,
Pose(dist_vec.x * clamp_distance, dist_vec.y * clamp_distance))
new_node = Node(new_pose)
nearest_node.children.append(new_node)
new_node.parent = nearest_node
return new_node
def plot_tree(ax, rr_tree: Node, grid_size: int, start: Node, end: Node,
reached_target: bool = False, last_node: Node = None, obstacles: list[Obstacle] = None,
fast_plot: bool = False) -> None:
# Add obstacles to the Axes
if obstacles is not None:
for obs in obstacles:
rect = patches.Rectangle((obs.x, obs.y), obs.width, obs.height,
linewidth=1, edgecolor='none', facecolor='gray')
ax.add_patch(rect)
ax.set_xlim(-grid_size * 0.2, grid_size * 1.2)
ax.set_ylim(-grid_size * 0.2, grid_size * 1.2)
ax.set_aspect('equal')
ax.grid()
nodes = []
connections = []
for node in rr_tree.traverse():
nodes.append(node)
for child in node.children:
x1, y1 = node.pose.x, node.pose.y
x2, y2 = child.pose.x, child.pose.y
connections.append([[x1, x2], [y1, y2]])
# Draw connections to children
for connection in connections:
ax.plot(connection[0], connection[1], 'k--', linewidth=0.5, zorder=8)
# Draw the nodes themselves
x_nodes = [node.pose.x for node in nodes]
y_nodes = [node.pose.y for node in nodes]
if fast_plot: # use plt.plot instead of plt.scatter which is super slow
ax.plot(x_nodes, y_nodes, 'b.', zorder=10)
else:
colors = [end.distance_to(node) for node in nodes]
ax.scatter(x_nodes, y_nodes, c=colors, s=15, vmin=0, vmax=90, zorder=10)
ax.scatter([start.pose.x], [start.pose.y], c='red', s=60, zorder=11)
ax.scatter([end.pose.x], [end.pose.y], c='green', s=100, zorder=11)
if reached_target:
# Draw the actual path from goal to start
node = last_node
while node.parent is not None:
ax.plot([node.pose.x, node.parent.pose.x], [node.pose.y, node.parent.pose.y], 'r-')
plt.draw()
plt.pause(1e-17)
node = node.parent
time.sleep(0.05)
plt.draw()
plt.pause(1e-17)
ax.clear()
def sample_new_node(obstacles: list[Obstacle], grid_size: int) -> Node:
while True:
new_node = Node(
Pose(x=random.uniform(0, grid_size),
y=random.uniform(0, grid_size)))
if not any([obs.contains(new_node.pose) for obs in obstacles]):
return new_node
def rrt_path_finder(cfg: RrtConfig):
rr_tree = cfg.start_node
counter = 1
found_target = False
global_closest = 1e6
while not found_target:
last_time = time.time()
not_valid_sample = True
while not_valid_sample:
new_node = Node(Pose(x=random.uniform(0, cfg.grid_size),
y=random.uniform(0, cfg.grid_size)))
if not any([obs.contains(new_node.pose) for obs in obstacles]):
not_valid_sample = False
closest_node = None
closest_distance = 1e6
for node in rr_tree.traverse():
distance = new_node.distance_to(node)
if distance < closest_distance:
closest_node = node
closest_distance = distance
all_object_lines = []
for obs in obstacles:
all_object_lines.extend(obs.segments)
seg_to_closest = [(new_node.pose.x, new_node.pose.y), (closest_node.pose.x, closest_node.pose.y)]
if any([intersects(seg, seg_to_closest) for seg in all_object_lines]):
continue
distance_to_new = new_node.distance_to(closest_node)
if distance_to_new > cfg.clamp_dist:
dist_vec = Pose.normalize(Pose.subtract(closest_node.pose, new_node.pose))
new_pose = Pose.add(closest_node.pose, Pose(dist_vec.x * cfg.clamp_dist, dist_vec.y * cfg.clamp_dist))
else:
new_pose = new_node.pose
new_node = Node(new_pose)
closest_node.children.append(new_node)
new_node.parent = closest_node
dist_to_end = new_node.distance_to(cfg.end_node)
global_closest = dist_to_end if dist_to_end < global_closest else global_closest
print(f'[{counter}] (new node / closest global) = ({closest_distance:.2f} / global {global_closest:.2f})')
if new_node.distance_to(cfg.end_node) < cfg.eps or counter == cfg.max_steps:
found_target = True
print('Target found or steps done!')
plot_tree(rr_tree, end_node=cfg.end_node, reached_target=True)
input()
sys.exit(0)
counter += 1
plot_tree(rr_tree, end_node=cfg.end_node, reached_target=False)
print(f'\ttook {time.time() - last_time:.2f}s')
def test_traversal():
root = Node(Pose(1, 1),
children=[
Node(Pose(2, 2)),
Node(Pose(3, 3),
children=[Node(Pose(5, 5)),
Node(Pose(4, 4))]),
Node(Pose(6, 6))
])
for node in root.traverse():
print(node)
def find_closest_node(new_node: Node, rr_tree: Node,
obstacles: list[Obstacle]) -> Optional[Node]:
"""For a new node, find the nearest node in the tree. Return this nearest node if there is an obstacle-free
connection between the new and nearest node, else return None."""
closest_node = None
closest_distance = 1e6
for node in rr_tree.traverse():
distance = new_node.distance_to(node)
if distance < closest_distance:
closest_node = node
closest_distance = distance
all_object_lines = []
for obs in obstacles:
all_object_lines.extend(obs.segments)
seg_to_closest = [(new_node.pose.x, new_node.pose.y),
(closest_node.pose.x, closest_node.pose.y)]
if any([intersects(seg, seg_to_closest) for seg in all_object_lines]):
return None
else:
return closest_node