1use nalgebra::Vector2;
7use std::f64::consts::PI;
8
9use rust_robotics_core::{Path2D, PathPlanner, Point2D, RoboticsError};
10
11use crate::rrt::CircleObstacle;
12
13const DEFAULT_INFLUENCE_DISTANCE: f64 = 2.0;
14const DEFAULT_ROBOT_RADIUS: f64 = 0.8;
15const EPS: f64 = 1.0e-9;
16
17#[derive(Debug, Clone, Copy)]
19pub struct ChompConfig {
20 pub n_waypoints: usize,
22 pub dt: f64,
24 pub max_iterations: usize,
26 pub learning_rate: f64,
28 pub obstacle_cost_weight: f64,
30 pub smoothness_weight: f64,
32}
33
34impl Default for ChompConfig {
35 fn default() -> Self {
36 Self {
37 n_waypoints: 50,
38 dt: 0.1,
39 max_iterations: 100,
40 learning_rate: 0.01,
41 obstacle_cost_weight: 1.0,
42 smoothness_weight: 1.0,
43 }
44 }
45}
46
47#[derive(Debug, Clone)]
49pub struct ChompResult {
50 pub path: Path2D,
51 pub cost: f64,
52 pub iterations: usize,
53}
54
55pub struct ChompPlanner {
57 config: ChompConfig,
58 obstacles: Vec<CircleObstacle>,
59}
60
61impl ChompPlanner {
62 pub fn new(obstacles: Vec<CircleObstacle>, config: ChompConfig) -> Self {
63 Self { config, obstacles }
64 }
65
66 pub fn optimize(&self, start: Point2D, goal: Point2D) -> Result<ChompResult, RoboticsError> {
67 if self.config.n_waypoints < 2 {
68 return Err(RoboticsError::InvalidParameter(
69 "CHOMP requires at least 2 waypoints".to_string(),
70 ));
71 }
72 if self.config.dt <= 0.0 {
73 return Err(RoboticsError::InvalidParameter(
74 "CHOMP dt must be positive".to_string(),
75 ));
76 }
77 if self.config.learning_rate <= 0.0 {
78 return Err(RoboticsError::InvalidParameter(
79 "CHOMP learning_rate must be positive".to_string(),
80 ));
81 }
82
83 let mut waypoints = self.initialize_trajectory(start, goal);
84 let mut cost = self.trajectory_cost(&waypoints);
85 let mut iterations = 0usize;
86
87 for iter in 0..self.config.max_iterations {
88 iterations = iter + 1;
89 let mut step_size = self.config.learning_rate;
90 let mut accepted = None;
91
92 for _ in 0..8 {
93 let mut next = waypoints.clone();
94 for i in 1..(waypoints.len() - 1) {
95 let smooth = self.smoothness_gradient(&waypoints, i);
96 let obstacle = self.obstacle_gradient(&waypoints[i]);
97 let grad = self.config.smoothness_weight * smooth
98 + self.config.obstacle_cost_weight * obstacle;
99 next[i] -= step_size * grad;
100 }
101
102 next[0] = Vector2::new(start.x, start.y);
104 *next.last_mut().expect("non-empty trajectory") = Vector2::new(goal.x, goal.y);
105
106 let next_cost = self.trajectory_cost(&next);
107 if next_cost <= cost {
108 accepted = Some((next, next_cost));
109 break;
110 }
111 step_size *= 0.5;
112 }
113
114 let Some((next, next_cost)) = accepted else {
115 break;
116 };
117 if (cost - next_cost).abs() < 1.0e-9 {
118 waypoints = next;
119 cost = next_cost;
120 break;
121 }
122 waypoints = next;
123 cost = next_cost;
124 }
125
126 let points = waypoints
127 .into_iter()
128 .map(|p| Point2D::new(p.x, p.y))
129 .collect();
130
131 Ok(ChompResult {
132 path: Path2D::from_points(points),
133 cost,
134 iterations,
135 })
136 }
137
138 fn initialize_trajectory(&self, start: Point2D, goal: Point2D) -> Vec<Vector2<f64>> {
139 let mut trajectory = Vec::with_capacity(self.config.n_waypoints);
140 for i in 0..self.config.n_waypoints {
141 let t = if self.config.n_waypoints == 1 {
142 0.0
143 } else {
144 i as f64 / (self.config.n_waypoints - 1) as f64
145 };
146 trajectory.push(Vector2::new(
147 start.x + t * (goal.x - start.x),
148 start.y + t * (goal.y - start.y) + 1.0e-3 * (PI * t).sin(),
149 ));
150 }
151 trajectory
152 }
153
154 fn smoothness_gradient(&self, trajectory: &[Vector2<f64>], i: usize) -> Vector2<f64> {
155 let dt2 = self.config.dt * self.config.dt;
156 let second_diff = trajectory[i - 1] - 2.0 * trajectory[i] + trajectory[i + 1];
157 -2.0 * second_diff / dt2
158 }
159
160 fn obstacle_gradient(&self, point: &Vector2<f64>) -> Vector2<f64> {
161 let mut grad = Vector2::zeros();
162 for obs in &self.obstacles {
163 let center = Vector2::new(obs.x, obs.y);
164 let delta = *point - center;
165 let norm = delta.norm().max(EPS);
166 let signed_distance = norm - (obs.radius + DEFAULT_ROBOT_RADIUS);
167 if signed_distance < DEFAULT_INFLUENCE_DISTANCE {
168 let direction = delta / norm;
169 grad -= (DEFAULT_INFLUENCE_DISTANCE - signed_distance) * direction;
172 }
173 }
174 grad
175 }
176
177 fn smoothness_cost(&self, trajectory: &[Vector2<f64>]) -> f64 {
178 trajectory
179 .windows(3)
180 .map(|window| {
181 let second_diff = window[0] - 2.0 * window[1] + window[2];
182 second_diff.norm_squared() / (self.config.dt * self.config.dt)
183 })
184 .sum()
185 }
186
187 fn obstacle_cost(&self, trajectory: &[Vector2<f64>]) -> f64 {
188 let mut cost = 0.0;
189 for point in trajectory {
190 for obs in &self.obstacles {
191 let dx = point.x - obs.x;
192 let dy = point.y - obs.y;
193 let distance = (dx * dx + dy * dy).sqrt();
194 let signed_distance = distance - (obs.radius + DEFAULT_ROBOT_RADIUS);
195 if signed_distance < DEFAULT_INFLUENCE_DISTANCE {
196 let penetration = DEFAULT_INFLUENCE_DISTANCE - signed_distance;
197 cost += 0.5 * penetration * penetration;
198 }
199 }
200 }
201 cost
202 }
203
204 fn trajectory_cost(&self, trajectory: &[Vector2<f64>]) -> f64 {
205 self.config.smoothness_weight * self.smoothness_cost(trajectory)
206 + self.config.obstacle_cost_weight * self.obstacle_cost(trajectory)
207 }
208}
209
210impl PathPlanner for ChompPlanner {
211 fn plan(&self, start: Point2D, goal: Point2D) -> Result<Path2D, RoboticsError> {
212 self.optimize(start, goal).map(|result| result.path)
213 }
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219 use crate::rrt::{AreaBounds, RRTConfig, RRTNode, RRTPlanner};
220
221 fn smoothness_metric(path: &Path2D) -> f64 {
222 if path.points.len() < 3 {
223 return 0.0;
224 }
225 path.points
226 .windows(3)
227 .map(|w| {
228 let ax = w[0].x - 2.0 * w[1].x + w[2].x;
229 let ay = w[0].y - 2.0 * w[1].y + w[2].y;
230 ax * ax + ay * ay
231 })
232 .sum()
233 }
234
235 #[test]
236 fn test_chomp_converges_and_keeps_endpoints() {
237 let planner = ChompPlanner::new(
238 vec![CircleObstacle::new(5.0, 2.0, 1.0)],
239 ChompConfig {
240 obstacle_cost_weight: 3.0,
241 ..Default::default()
242 },
243 );
244 let start = Point2D::new(0.0, 0.0);
245 let goal = Point2D::new(10.0, 4.0);
246
247 let initial = planner.initialize_trajectory(start, goal);
248 let initial_cost = planner.trajectory_cost(&initial);
249 let result = planner
250 .optimize(start, goal)
251 .expect("optimization should succeed");
252
253 assert!(result.cost <= initial_cost);
254 assert_eq!(result.path.points.first().unwrap().x, start.x);
255 assert_eq!(result.path.points.first().unwrap().y, start.y);
256 assert_eq!(result.path.points.last().unwrap().x, goal.x);
257 assert_eq!(result.path.points.last().unwrap().y, goal.y);
258 }
259
260 #[test]
261 fn test_chomp_avoids_central_obstacle() {
262 let obstacles = vec![CircleObstacle::new(5.0, 0.0, 1.5)];
263 let planner = ChompPlanner::new(
264 obstacles.clone(),
265 ChompConfig {
266 max_iterations: 200,
267 obstacle_cost_weight: 5.0,
268 ..Default::default()
269 },
270 );
271 let path = planner
272 .plan(Point2D::new(0.0, 0.0), Point2D::new(10.0, 0.0))
273 .expect("planning should succeed");
274
275 let min_dist = path
276 .points
277 .iter()
278 .map(|p| ((p.x - obstacles[0].x).powi(2) + (p.y - obstacles[0].y).powi(2)).sqrt())
279 .fold(f64::INFINITY, f64::min);
280
281 assert!(
282 min_dist > obstacles[0].radius,
283 "optimized trajectory should bend around obstacle"
284 );
285 }
286
287 #[test]
288 fn test_chomp_path_is_smoother_than_rrt_path() {
289 let mut rrt = RRTPlanner::new(
290 vec![],
291 AreaBounds::new(-2.0, 15.0, -5.0, 5.0),
292 None,
293 RRTConfig::default(),
294 );
295 let samples = [
296 [2.0, 2.0],
297 [4.0, -2.0],
298 [6.0, 2.0],
299 [8.0, -2.0],
300 [10.0, 0.0],
301 [10.0, 0.0],
302 ];
303 let mut idx = 0usize;
304 let rrt_path = rrt
305 .plan_with_sampler(Point2D::new(0.0, 0.0), Point2D::new(10.0, 0.0), |_| {
306 let sample = samples[idx.min(samples.len() - 1)];
307 idx += 1;
308 RRTNode::new(sample[0], sample[1])
309 })
310 .expect("rrt should find a path");
311
312 let chomp = ChompPlanner::new(vec![], ChompConfig::default());
313 let chomp_path = chomp
314 .plan(Point2D::new(0.0, 0.0), Point2D::new(10.0, 0.0))
315 .expect("chomp should optimize a path");
316
317 assert!(smoothness_metric(&chomp_path) < smoothness_metric(&rrt_path));
318 }
319}