repose_motion_planning/
stomp.rs

1//! STOMP: stochastic trajectory optimization, anytime and seeded.
2//!
3//! Like [`CuJointPathOptimizer`](crate::CuJointPathOptimizer) the
4//! path is a fixed number of waypoints seeded on the straight joint line, but
5//! updates need no gradients. Each iteration draws `rollouts` smooth noisy
6//! copies of the path, scores every waypoint of every rollout by its clearance
7//! penalty plus a local acceleration (smoothness) cost, weights the noise per
8//! waypoint by `exp(-sensitivity · normalized cost)`, and adds the smoothed
9//! weighted noise to the path. That copes with costs that are flat or
10//! discontinuous where a gradient would vanish. The shortest collision-free
11//! path seen so far is the result. All randomness comes from the job seed.
12
13use crate::CuJointPlanPath;
14use crate::CuJointPlanRequest;
15use crate::CuJointPlanner;
16use crate::CuJointPlannerTask;
17use crate::CuPlanStatus;
18use crate::CuPlannerKind;
19use crate::MAX_WAYPOINTS;
20use crate::common::check_endpoints;
21use crate::common::length_quality;
22use crate::common::path_length;
23use crate::common::write_header;
24use crate::validity::CuSceneValidity;
25use crate::validity::joint_distance;
26use bincode::Decode;
27use bincode::Encode;
28use cu_rng::prelude::*;
29use cu29::cutask_anytime::Quality;
30use cu29::prelude::*;
31use cu29::units::si::f64::Length;
32use cu29::units::si::length::meter;
33use repose_collision::CuRobotCollisionModel;
34use repose_kinematics::CuJointMetric;
35use repose_kinematics::CuJointPosition;
36use repose_kinematics::CuJointPositionsPayload;
37use repose_kinematics::CuJointResolution;
38use repose_kinematics::CuKinematicChain;
39use serde::Deserialize;
40use serde::Serialize;
41
42/// Rollout capacity per iteration.
43pub const STOMP_MAX_ROLLOUTS: usize = 16;
44
45/// Parameters of [`CuJointStomp`].
46#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Encode, Decode, Reflect)]
47#[serde(default, deny_unknown_fields)]
48pub struct CuJointStompParams {
49    /// Path waypoints including start and goal (`3..=MAX_WAYPOINTS`).
50    pub waypoints: u32,
51    /// Noisy rollouts per iteration (`2..=STOMP_MAX_ROLLOUTS`).
52    pub rollouts: u32,
53    /// Standard deviation of the exploration noise (dimensionless,
54    /// metric-scaled).
55    pub noise: f64,
56    /// Exponent scale turning normalized rollout costs into weights
57    /// (dimensionless).
58    pub sensitivity: f64,
59    /// Clearance below which the obstacle cost applies.
60    pub clearance: Length,
61    /// Weight of the obstacle cost (dimensionless, `>= 0`).
62    pub obstacle_weight: f64,
63    /// Weight of the smoothness cost (dimensionless, `>= 0`).
64    pub smoothness_weight: f64,
65    /// Iterations after which the job ends.
66    pub max_iterations: u32,
67    /// Resolution of the segment collision checks.
68    pub collision_check_step: CuJointResolution,
69    /// Scaling of dimensionless joint-space distances and costs.
70    pub metric: CuJointMetric,
71}
72
73impl Default for CuJointStompParams {
74    fn default() -> Self {
75        Self {
76            waypoints: 12,
77            rollouts: 8,
78            noise: 0.15,
79            sensitivity: 10.0,
80            clearance: Length::new::<meter>(0.05),
81            obstacle_weight: 40.0,
82            smoothness_weight: 1.0,
83            max_iterations: 200,
84            collision_check_step: CuJointResolution::default(),
85            metric: CuJointMetric::default(),
86        }
87    }
88}
89
90impl CuJointStompParams {
91    /// Checks every parameter, the metric, and the collision-check resolution.
92    pub fn validate(&self) -> CuResult<()> {
93        self.metric.validate()?;
94        self.collision_check_step.metric().validate()?;
95        let valid = {
96            let positive = |v: f64| v.is_finite() && v > 0.0;
97            (3..=MAX_WAYPOINTS as u32).contains(&self.waypoints)
98                && (2..=STOMP_MAX_ROLLOUTS as u32).contains(&self.rollouts)
99                && positive(self.noise)
100                && positive(self.sensitivity)
101                && positive(self.clearance.get::<meter>())
102                && self.obstacle_weight.is_finite()
103                && self.obstacle_weight >= 0.0
104                && self.smoothness_weight.is_finite()
105                && self.smoothness_weight >= 0.0
106                && self.max_iterations > 0
107        };
108        if valid {
109            Ok(())
110        } else {
111            Err(CuError::from(
112                "repose-motion-planning: invalid STOMP parameters",
113            ))
114        }
115    }
116}
117
118/// Stochastic trajectory optimizer.
119///
120/// Built with [`CuJointPlanner::new`] from robot models and [`CuJointStompParams`];
121/// runs in [`CuJointStompTask`].
122pub struct CuJointStomp<const N: usize, const M: usize = N> {
123    validity: CuSceneValidity<N, M>,
124    params: CuJointStompParams,
125    path: [[f64; N]; MAX_WAYPOINTS],
126    kinds: [CuJointPosition; N],
127    best: [[CuJointPosition; N]; MAX_WAYPOINTS],
128    count: usize,
129    best_cost: f64,
130    straight_cost: f64,
131    rng: CuRng,
132    iterations: u32,
133    status: CuPlanStatus,
134}
135
136/// Approximately standard normal sample (Irwin–Hall with four uniforms).
137fn normal(rng: &mut CuRng) -> f64 {
138    let sum: f64 = (0..4).map(|_| rng.random::<f64>()).sum();
139    (sum - 2.0) * libm::sqrt(3.0)
140}
141
142/// In-place `[1/4, 1/2, 1/4]` smoothing along the path, endpoints fixed at zero.
143fn smooth<const N: usize>(values: &mut [[f64; N]], count: usize) {
144    let copy: [[f64; N]; MAX_WAYPOINTS] =
145        core::array::from_fn(|i| if i < count { values[i] } else { [0.0; N] });
146    for i in 1..count - 1 {
147        for j in 0..N {
148            values[i][j] = 0.25 * copy[i - 1][j] + 0.5 * copy[i][j] + 0.25 * copy[i + 1][j];
149        }
150    }
151    values[0] = [0.0; N];
152    values[count - 1] = [0.0; N];
153}
154
155/// [`CuJointPlannerTask`] running [`CuJointStomp`].
156pub type CuJointStompTask<const N: usize, const M: usize = N> =
157    CuJointPlannerTask<CuJointStomp<N, M>, N, M>;
158
159impl<const N: usize, const M: usize> CuJointStomp<N, M> {
160    /// The validated planner parameters.
161    pub const fn params(&self) -> &CuJointStompParams {
162        &self.params
163    }
164
165    fn build(
166        chain: CuKinematicChain<N>,
167        model: CuRobotCollisionModel<N, M>,
168        params: CuJointStompParams,
169    ) -> Self {
170        Self {
171            validity: CuSceneValidity::new(chain, model, params.collision_check_step),
172            params,
173            path: [[0.0; N]; MAX_WAYPOINTS],
174            kinds: [CuJointPosition::default(); N],
175            best: [[CuJointPosition::default(); N]; MAX_WAYPOINTS],
176            count: 0,
177            best_cost: f64::INFINITY,
178            straight_cost: 0.0,
179            rng: CuRng::from_seed(0),
180            iterations: 0,
181            status: CuPlanStatus::InvalidRequest,
182        }
183    }
184
185    /// Iterations run on the current job.
186    pub const fn iterations(&self) -> u32 {
187        self.iterations
188    }
189
190    fn scale(&self, j: usize) -> f64 {
191        self.params.metric.scale(self.kinds[j])
192    }
193
194    fn configuration(&self, x: &[f64; N]) -> [CuJointPosition; N] {
195        core::array::from_fn(|j| self.kinds[j].with_si_value(x[j] * self.scale(j)))
196    }
197
198    fn clamp(&self, x: &mut [f64; N]) {
199        for (j, value) in x.iter_mut().enumerate() {
200            if matches!(self.kinds[j], CuJointPosition::Fixed) {
201                *value = 0.0;
202                continue;
203            }
204            let limits = self.validity.chain().joints()[j].limits;
205            let lo = limits.lower().si_value() / self.scale(j);
206            let hi = limits.upper().si_value() / self.scale(j);
207            *value = value.clamp(lo, hi);
208        }
209    }
210
211    fn obstacle_cost(&self, x: &[f64; N]) -> f64 {
212        let eps = self.params.clearance.get::<meter>();
213        match self.validity.clearance(&self.configuration(x)) {
214            Some(c) if c >= eps => 0.0,
215            Some(c) => self.params.obstacle_weight * (eps - c) * (eps - c) / eps,
216            None => self.params.obstacle_weight * eps,
217        }
218    }
219
220    fn record_if_better(&mut self) {
221        let mut q = [[CuJointPosition::default(); N]; MAX_WAYPOINTS];
222        for (slot, x) in q.iter_mut().zip(&self.path[..self.count]) {
223            *slot = self.configuration(x);
224        }
225        let free = (1..self.count - 1).all(|i| self.validity.is_configuration_free(&q[i]))
226            && (0..self.count - 1).all(|i| self.validity.is_segment_free(&q[i], &q[i + 1]));
227        if !free {
228            return;
229        }
230        let cost = path_length(&q[..self.count], self.params.metric);
231        if cost < self.best_cost {
232            self.best_cost = cost;
233            self.best = q;
234        }
235    }
236
237    // Rollout, waypoint and joint indices address several parallel arrays;
238    // explicit indices read clearer than zipped iterators here.
239    #[allow(clippy::needless_range_loop)]
240    fn iterate(&mut self) {
241        let k = self.params.rollouts as usize;
242        let count = self.count;
243        let mut noise = [[[0.0_f64; N]; MAX_WAYPOINTS]; STOMP_MAX_ROLLOUTS];
244        let mut costs = [[0.0_f64; MAX_WAYPOINTS]; STOMP_MAX_ROLLOUTS];
245        for r in 0..k {
246            for i in 1..count - 1 {
247                for j in 0..N {
248                    noise[r][i][j] = self.params.noise * normal(&mut self.rng);
249                }
250            }
251            smooth(&mut noise[r], count);
252            for i in 1..count - 1 {
253                let mut x = self.path[i];
254                for j in 0..N {
255                    x[j] += noise[r][i][j];
256                }
257                self.clamp(&mut x);
258                let mut acceleration = 0.0;
259                for j in 0..N {
260                    let a = (self.path[i - 1][j] + noise[r][i - 1][j]) - 2.0 * x[j]
261                        + (self.path[i + 1][j] + noise[r][i + 1][j]);
262                    acceleration += a * a;
263                }
264                costs[r][i] = self.obstacle_cost(&x) + self.params.smoothness_weight * acceleration;
265            }
266        }
267        let mut update = [[0.0_f64; N]; MAX_WAYPOINTS];
268        for i in 1..count - 1 {
269            let (lo, hi) = (0..k).fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), r| {
270                (lo.min(costs[r][i]), hi.max(costs[r][i]))
271            });
272            let span = (hi - lo).max(1e-12);
273            let mut weights = [0.0_f64; STOMP_MAX_ROLLOUTS];
274            let mut total = 0.0;
275            for r in 0..k {
276                weights[r] = libm::exp(-self.params.sensitivity * (costs[r][i] - lo) / span);
277                total += weights[r];
278            }
279            for r in 0..k {
280                for j in 0..N {
281                    update[i][j] += weights[r] / total * noise[r][i][j];
282                }
283            }
284        }
285        smooth(&mut update, count);
286        for i in 1..count - 1 {
287            let mut x = self.path[i];
288            for j in 0..N {
289                x[j] += update[i][j];
290            }
291            self.clamp(&mut x);
292            self.path[i] = x;
293        }
294    }
295}
296
297impl<const N: usize, const M: usize> CuJointPlanner<N, M> for CuJointStomp<N, M> {
298    const KIND: CuPlannerKind = CuPlannerKind::Optimization;
299
300    type Params = CuJointStompParams;
301
302    fn new(
303        chain: CuKinematicChain<N>,
304        collision_model: CuRobotCollisionModel<N, M>,
305        params: CuJointStompParams,
306    ) -> CuResult<Self> {
307        params.validate()?;
308        Ok(Self::build(chain, collision_model, params))
309    }
310
311    fn begin(&mut self, request: &CuJointPlanRequest<N>, seed: u64) -> CuResult<CuPlanStatus> {
312        self.rng.reseed(seed);
313        self.iterations = 0;
314        self.best_cost = f64::INFINITY;
315        self.count = self.params.waypoints as usize;
316        self.kinds = request.start;
317        self.straight_cost = joint_distance(&request.start, &request.goal, self.params.metric);
318        self.status = check_endpoints(&mut self.validity, request, self.params.metric);
319        if self.status != CuPlanStatus::Searching {
320            return Ok(self.status);
321        }
322        let last = (self.count - 1) as f64;
323        for i in 0..self.count {
324            let t = i as f64 / last;
325            for j in 0..N {
326                let a = self.params.metric.normalized(request.start[j]);
327                let b = self.params.metric.normalized(request.goal[j]);
328                self.path[i][j] = a + t * (b - a);
329            }
330        }
331        // A free straight line is already the shortest path.
332        self.record_if_better();
333        if self.best_cost.is_finite() {
334            self.status = CuPlanStatus::Converged;
335        }
336        Ok(self.status)
337    }
338
339    fn step(&mut self, budget: u32) -> CuResult<CuPlanStatus> {
340        for _ in 0..budget {
341            if self.status.is_final() {
342                break;
343            }
344            self.iterate();
345            self.iterations += 1;
346            self.record_if_better();
347            let done = self.iterations >= self.params.max_iterations;
348            self.status = match (self.best_cost.is_finite(), done) {
349                (true, true) => CuPlanStatus::Converged,
350                (true, false) => CuPlanStatus::Improving,
351                (false, true) => CuPlanStatus::Exhausted,
352                (false, false) => CuPlanStatus::Searching,
353            };
354        }
355        Ok(self.status)
356    }
357
358    fn write_best(&self, out: &mut CuJointPlanPath<N>) -> CuPlanStatus {
359        write_header(out, &self.validity, self.status, self.iterations);
360        if self.status.has_path() {
361            for (slot, q) in out.waypoints.iter_mut().zip(&self.best[..self.count]) {
362                *slot = CuJointPositionsPayload { positions: *q };
363            }
364            out.len = self.count as u32;
365            out.cost = self.best_cost;
366        }
367        self.status
368    }
369
370    fn quality(&self) -> Quality {
371        length_quality(self.straight_cost, self.best_cost)
372    }
373}
374
375crate::impl_planner_type_path!(CuJointStomp, "repose_motion_planning::CuJointStomp");
repose_motion_planning/stomp.rs

repose_motion_planning/
stomp.rs
