mirror of
https://github.com/AtsushiSakai/PythonRobotics.git
synced 2026-01-14 09:08:01 -05:00
c92aaf36d8
* feat: add ElasticBands * feat: Elastic Bands update * feat: ElasticBands update * feat: ElasticBands add test * feat: ElasticBands reduce occupation * fix: ElasticBands test * feat: ElasticBands remove tangential component * feat: Elastic Bands update * feat: Elastic Bands doc * feat: Elastic Bands update * feat: ElasticBands update
301 lines
10 KiB
Python
301 lines
10 KiB
Python
"""
|
||
Elastic Bands
|
||
|
||
author: Wang Zheng (@Aglargil)
|
||
|
||
Ref:
|
||
|
||
- [Elastic Bands: Connecting Path Planning and Control]
|
||
(http://www8.cs.umu.se/research/ifor/dl/Control/elastic%20bands.pdf)
|
||
"""
|
||
|
||
import numpy as np
|
||
import sys
|
||
import pathlib
|
||
import matplotlib.pyplot as plt
|
||
from matplotlib.patches import Circle
|
||
|
||
sys.path.append(str(pathlib.Path(__file__).parent.parent.parent))
|
||
|
||
from Mapping.DistanceMap.distance_map import compute_sdf_scipy
|
||
|
||
# Elastic Bands Params
|
||
MAX_BUBBLE_RADIUS = 100
|
||
MIN_BUBBLE_RADIUS = 10
|
||
RHO0 = 20.0 # Maximum distance for applying repulsive force
|
||
KC = 0.05 # Contraction force gain
|
||
KR = -0.1 # Repulsive force gain
|
||
LAMBDA = 0.7 # Overlap constraint factor
|
||
STEP_SIZE = 3.0 # Step size for calculating gradient
|
||
|
||
# Visualization Params
|
||
ENABLE_PLOT = True
|
||
# ENABLE_INTERACTIVE is True allows user to add obstacles by left clicking
|
||
# and add path points by right clicking and start planning by middle clicking
|
||
ENABLE_INTERACTIVE = False
|
||
# ENABLE_SAVE_DATA is True allows saving the path and obstacles which added
|
||
# by user in interactive mode to file
|
||
ENABLE_SAVE_DATA = False
|
||
MAX_ITER = 50
|
||
|
||
|
||
class Bubble:
|
||
def __init__(self, position, radius):
|
||
self.pos = np.array(position) # Bubble center coordinates [x, y]
|
||
self.radius = radius # Safety distance radius ρ(b)
|
||
if self.radius > MAX_BUBBLE_RADIUS:
|
||
self.radius = MAX_BUBBLE_RADIUS
|
||
if self.radius < MIN_BUBBLE_RADIUS:
|
||
self.radius = MIN_BUBBLE_RADIUS
|
||
|
||
|
||
class ElasticBands:
|
||
def __init__(
|
||
self,
|
||
initial_path,
|
||
obstacles,
|
||
rho0=RHO0,
|
||
kc=KC,
|
||
kr=KR,
|
||
lambda_=LAMBDA,
|
||
step_size=STEP_SIZE,
|
||
):
|
||
self.distance_map = compute_sdf_scipy(obstacles)
|
||
self.bubbles = [
|
||
Bubble(p, self.compute_rho(p)) for p in initial_path
|
||
] # Initialize bubble chain
|
||
self.kc = kc # Contraction force gain
|
||
self.kr = kr # Repulsive force gain
|
||
self.rho0 = rho0 # Maximum distance for applying repulsive force
|
||
self.lambda_ = lambda_ # Overlap constraint factor
|
||
self.step_size = step_size # Step size for calculating gradient
|
||
self._maintain_overlap()
|
||
|
||
def compute_rho(self, position):
|
||
"""Compute the distance field value at the position"""
|
||
return self.distance_map[int(position[0]), int(position[1])]
|
||
|
||
def contraction_force(self, i):
|
||
"""Calculate internal contraction force for the i-th bubble"""
|
||
if i == 0 or i == len(self.bubbles) - 1:
|
||
return np.zeros(2)
|
||
|
||
prev = self.bubbles[i - 1].pos
|
||
next_ = self.bubbles[i + 1].pos
|
||
current = self.bubbles[i].pos
|
||
|
||
# f_c = kc * ( (prev-current)/|prev-current| + (next-current)/|next-current| )
|
||
dir_prev = (prev - current) / (np.linalg.norm(prev - current) + 1e-6)
|
||
dir_next = (next_ - current) / (np.linalg.norm(next_ - current) + 1e-6)
|
||
return self.kc * (dir_prev + dir_next)
|
||
|
||
def repulsive_force(self, i):
|
||
"""Calculate external repulsive force for the i-th bubble"""
|
||
h = self.step_size # Step size
|
||
b = self.bubbles[i].pos
|
||
rho = self.bubbles[i].radius
|
||
|
||
if rho >= self.rho0:
|
||
return np.zeros(2)
|
||
|
||
# Finite difference approximation of the gradient ∂ρ/∂b
|
||
dx = np.array([h, 0])
|
||
dy = np.array([0, h])
|
||
grad_x = (self.compute_rho(b - dx) - self.compute_rho(b + dx)) / (2 * h)
|
||
grad_y = (self.compute_rho(b - dy) - self.compute_rho(b + dy)) / (2 * h)
|
||
grad = np.array([grad_x, grad_y])
|
||
|
||
return self.kr * (self.rho0 - rho) * grad
|
||
|
||
def update_bubbles(self):
|
||
"""Update bubble positions"""
|
||
new_bubbles = []
|
||
for i in range(len(self.bubbles)):
|
||
if i == 0 or i == len(self.bubbles) - 1:
|
||
new_bubbles.append(self.bubbles[i]) # Fixed start and end points
|
||
continue
|
||
|
||
f_total = self.contraction_force(i) + self.repulsive_force(i)
|
||
v = self.bubbles[i - 1].pos - self.bubbles[i + 1].pos
|
||
|
||
# Remove tangential component
|
||
f_star = f_total - f_total * v * v / (np.linalg.norm(v) ** 2 + 1e-6)
|
||
|
||
alpha = self.bubbles[i].radius # Adaptive step size
|
||
new_pos = self.bubbles[i].pos + alpha * f_star
|
||
new_pos = np.clip(new_pos, 0, 499)
|
||
new_radius = self.compute_rho(new_pos)
|
||
|
||
# Update bubble and maintain overlap constraint
|
||
new_bubble = Bubble(new_pos, new_radius)
|
||
new_bubbles.append(new_bubble)
|
||
|
||
self.bubbles = new_bubbles
|
||
self._maintain_overlap()
|
||
|
||
def _maintain_overlap(self):
|
||
"""Maintain bubble chain continuity (simplified insertion/deletion mechanism)"""
|
||
# Insert bubbles
|
||
i = 0
|
||
while i < len(self.bubbles) - 1:
|
||
bi, bj = self.bubbles[i], self.bubbles[i + 1]
|
||
dist = np.linalg.norm(bi.pos - bj.pos)
|
||
if dist > self.lambda_ * (bi.radius + bj.radius):
|
||
new_pos = (bi.pos + bj.pos) / 2
|
||
rho = self.compute_rho(
|
||
new_pos
|
||
) # Calculate new radius using environment model
|
||
self.bubbles.insert(i + 1, Bubble(new_pos, rho))
|
||
i += 2 # Skip the processed region
|
||
else:
|
||
i += 1
|
||
|
||
# Delete redundant bubbles
|
||
i = 1
|
||
while i < len(self.bubbles) - 1:
|
||
prev = self.bubbles[i - 1]
|
||
next_ = self.bubbles[i + 1]
|
||
dist = np.linalg.norm(prev.pos - next_.pos)
|
||
if dist <= self.lambda_ * (prev.radius + next_.radius):
|
||
del self.bubbles[i] # Delete if redundant
|
||
else:
|
||
i += 1
|
||
|
||
|
||
class ElasticBandsVisualizer:
|
||
def __init__(self):
|
||
self.obstacles = np.zeros((500, 500))
|
||
self.obstacles_points = []
|
||
self.path_points = []
|
||
self.elastic_band = None
|
||
self.running = True
|
||
|
||
if ENABLE_PLOT:
|
||
self.fig, self.ax = plt.subplots(figsize=(8, 8))
|
||
self.fig.canvas.mpl_connect("close_event", self.on_close)
|
||
self.ax.set_xlim(0, 500)
|
||
self.ax.set_ylim(0, 500)
|
||
|
||
if ENABLE_INTERACTIVE:
|
||
self.path_points = [] # Add a list to store path points
|
||
# Connect mouse events
|
||
self.fig.canvas.mpl_connect("button_press_event", self.on_click)
|
||
else:
|
||
self.path_points = np.load(pathlib.Path(__file__).parent / "path.npy")
|
||
self.obstacles_points = np.load(
|
||
pathlib.Path(__file__).parent / "obstacles.npy"
|
||
)
|
||
for x, y in self.obstacles_points:
|
||
self.add_obstacle(x, y)
|
||
self.plan_path()
|
||
|
||
self.plot_background()
|
||
|
||
def on_close(self, event):
|
||
"""Handle window close event"""
|
||
self.running = False
|
||
plt.close("all") # Close all figure windows
|
||
|
||
def plot_background(self):
|
||
"""Plot the background grid"""
|
||
if not ENABLE_PLOT or not self.running:
|
||
return
|
||
|
||
self.ax.cla()
|
||
self.ax.set_xlim(0, 500)
|
||
self.ax.set_ylim(0, 500)
|
||
self.ax.grid(True)
|
||
|
||
if ENABLE_INTERACTIVE:
|
||
self.ax.set_title(
|
||
"Elastic Bands Path Planning\n"
|
||
"Left click: Add obstacles\n"
|
||
"Right click: Add path points\n"
|
||
"Middle click: Start planning",
|
||
pad=20,
|
||
)
|
||
else:
|
||
self.ax.set_title("Elastic Bands Path Planning", pad=20)
|
||
|
||
if self.path_points:
|
||
self.ax.plot(
|
||
[p[0] for p in self.path_points],
|
||
[p[1] for p in self.path_points],
|
||
"yo",
|
||
markersize=8,
|
||
)
|
||
|
||
self.ax.imshow(self.obstacles.T, origin="lower", cmap="binary", alpha=0.8)
|
||
self.ax.plot([], [], color="black", label="obstacles")
|
||
if self.elastic_band is not None:
|
||
path = [b.pos.tolist() for b in self.elastic_band.bubbles]
|
||
path = np.array(path)
|
||
self.ax.plot(path[:, 0], path[:, 1], "b-", linewidth=2, label="path")
|
||
|
||
for bubble in self.elastic_band.bubbles:
|
||
circle = Circle(
|
||
bubble.pos, bubble.radius, fill=False, color="g", alpha=0.3
|
||
)
|
||
self.ax.add_patch(circle)
|
||
self.ax.plot(bubble.pos[0], bubble.pos[1], "bo", markersize=10)
|
||
self.ax.plot([], [], color="green", label="bubbles")
|
||
|
||
self.ax.legend(loc="upper right")
|
||
plt.draw()
|
||
plt.pause(0.01)
|
||
|
||
def add_obstacle(self, x, y):
|
||
"""Add an obstacle at the given coordinates"""
|
||
size = 30 # Side length of the square
|
||
half_size = size // 2
|
||
x_start = max(0, x - half_size)
|
||
x_end = min(self.obstacles.shape[0], x + half_size)
|
||
y_start = max(0, y - half_size)
|
||
y_end = min(self.obstacles.shape[1], y + half_size)
|
||
self.obstacles[x_start:x_end, y_start:y_end] = 1
|
||
|
||
def on_click(self, event):
|
||
"""Handle mouse click events"""
|
||
if event.inaxes != self.ax:
|
||
return
|
||
|
||
x, y = int(event.xdata), int(event.ydata)
|
||
|
||
if event.button == 1: # Left click to add obstacles
|
||
self.add_obstacle(x, y)
|
||
self.obstacles_points.append([x, y])
|
||
|
||
elif event.button == 3: # Right click to add path points
|
||
self.path_points.append([x, y])
|
||
|
||
elif event.button == 2: # Middle click to end path input and start planning
|
||
if len(self.path_points) >= 2:
|
||
if ENABLE_SAVE_DATA:
|
||
np.save(
|
||
pathlib.Path(__file__).parent / "path.npy", self.path_points
|
||
)
|
||
np.save(
|
||
pathlib.Path(__file__).parent / "obstacles.npy",
|
||
self.obstacles_points,
|
||
)
|
||
self.plan_path()
|
||
|
||
self.plot_background()
|
||
|
||
def plan_path(self):
|
||
"""Plan the path"""
|
||
|
||
initial_path = self.path_points
|
||
# Create an elastic band object and optimize
|
||
self.elastic_band = ElasticBands(initial_path, self.obstacles)
|
||
for _ in range(MAX_ITER):
|
||
self.elastic_band.update_bubbles()
|
||
self.path_points = [b.pos for b in self.elastic_band.bubbles]
|
||
self.plot_background()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
_ = ElasticBandsVisualizer()
|
||
if ENABLE_PLOT:
|
||
plt.show(block=True)
|