Skip to main content

rust_robotics_planning/
chomp.rs

1//! CHOMP (Covariant Hamiltonian Optimization for Motion Planning).
2//!
3//! This module provides a practical 2D CHOMP-style optimizer that balances
4//! smoothness and obstacle avoidance costs over a fixed waypoint trajectory.
5
6use 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/// Configuration for CHOMP trajectory optimization.
18#[derive(Debug, Clone, Copy)]
19pub struct ChompConfig {
20    /// Number of waypoints including start and goal.
21    pub n_waypoints: usize,
22    /// Time step \[s\] used for smoothness scaling.
23    pub dt: f64,
24    /// Maximum optimization iterations.
25    pub max_iterations: usize,
26    /// Gradient descent step size.
27    pub learning_rate: f64,
28    /// Weight for obstacle potential.
29    pub obstacle_cost_weight: f64,
30    /// Weight for smoothness penalty.
31    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/// Optimization result for CHOMP.
48#[derive(Debug, Clone)]
49pub struct ChompResult {
50    pub path: Path2D,
51    pub cost: f64,
52    pub iterations: usize,
53}
54
55/// CHOMP planner.
56pub 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                // Keep boundary conditions fixed.
103                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                // Potential increases near/inside obstacles; gradient points toward
170                // obstacle so gradient descent moves away.
171                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}