Skip to main content

rust_robotics_control/
behavior_tree.rs

1//! Behavior Tree implementation for mission planning.
2//!
3//! This module provides a compact behavior tree runtime suitable for
4//! robotics decision making with a shared blackboard.
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum BehaviorStatus {
8    Success,
9    Failure,
10    Running,
11}
12
13#[derive(Debug, Clone, Copy)]
14pub struct Blackboard {
15    pub battery_level: f64,
16    pub obstacle_detected: bool,
17    pub path_ready: bool,
18    pub goal_visible: bool,
19}
20
21impl Default for Blackboard {
22    fn default() -> Self {
23        Self {
24            battery_level: 100.0,
25            obstacle_detected: false,
26            path_ready: true,
27            goal_visible: true,
28        }
29    }
30}
31
32pub trait BehaviorNode {
33    fn tick(&mut self, blackboard: &Blackboard) -> BehaviorStatus;
34    fn reset(&mut self) {}
35    fn name(&self) -> &str;
36}
37
38pub struct BehaviorTree {
39    root: Box<dyn BehaviorNode>,
40}
41
42impl BehaviorTree {
43    pub fn new(root: Box<dyn BehaviorNode>) -> Self {
44        Self { root }
45    }
46
47    pub fn tick(&mut self, blackboard: &Blackboard) -> BehaviorStatus {
48        self.root.tick(blackboard)
49    }
50
51    pub fn reset(&mut self) {
52        self.root.reset();
53    }
54
55    pub fn root_name(&self) -> &str {
56        self.root.name()
57    }
58}
59
60pub struct ConditionNode {
61    name: String,
62    evaluator: Box<dyn Fn(&Blackboard) -> bool>,
63}
64
65impl ConditionNode {
66    pub fn new(name: impl Into<String>, evaluator: impl Fn(&Blackboard) -> bool + 'static) -> Self {
67        Self {
68            name: name.into(),
69            evaluator: Box::new(evaluator),
70        }
71    }
72}
73
74impl BehaviorNode for ConditionNode {
75    fn tick(&mut self, blackboard: &Blackboard) -> BehaviorStatus {
76        if (self.evaluator)(blackboard) {
77            BehaviorStatus::Success
78        } else {
79            BehaviorStatus::Failure
80        }
81    }
82
83    fn name(&self) -> &str {
84        &self.name
85    }
86}
87
88pub struct ActionNode {
89    name: String,
90    action: Box<dyn FnMut(&Blackboard) -> BehaviorStatus>,
91}
92
93impl ActionNode {
94    pub fn new(
95        name: impl Into<String>,
96        action: impl FnMut(&Blackboard) -> BehaviorStatus + 'static,
97    ) -> Self {
98        Self {
99            name: name.into(),
100            action: Box::new(action),
101        }
102    }
103}
104
105impl BehaviorNode for ActionNode {
106    fn tick(&mut self, blackboard: &Blackboard) -> BehaviorStatus {
107        (self.action)(blackboard)
108    }
109
110    fn name(&self) -> &str {
111        &self.name
112    }
113}
114
115pub struct SequenceNode {
116    name: String,
117    children: Vec<Box<dyn BehaviorNode>>,
118    current_index: usize,
119}
120
121impl SequenceNode {
122    pub fn new(name: impl Into<String>, children: Vec<Box<dyn BehaviorNode>>) -> Self {
123        Self {
124            name: name.into(),
125            children,
126            current_index: 0,
127        }
128    }
129}
130
131impl BehaviorNode for SequenceNode {
132    fn tick(&mut self, blackboard: &Blackboard) -> BehaviorStatus {
133        while self.current_index < self.children.len() {
134            match self.children[self.current_index].tick(blackboard) {
135                BehaviorStatus::Success => {
136                    self.children[self.current_index].reset();
137                    self.current_index += 1;
138                }
139                BehaviorStatus::Failure => {
140                    self.reset();
141                    return BehaviorStatus::Failure;
142                }
143                BehaviorStatus::Running => return BehaviorStatus::Running,
144            }
145        }
146
147        self.reset();
148        BehaviorStatus::Success
149    }
150
151    fn reset(&mut self) {
152        self.current_index = 0;
153        for child in &mut self.children {
154            child.reset();
155        }
156    }
157
158    fn name(&self) -> &str {
159        &self.name
160    }
161}
162
163pub struct SelectorNode {
164    name: String,
165    children: Vec<Box<dyn BehaviorNode>>,
166    current_index: usize,
167}
168
169impl SelectorNode {
170    pub fn new(name: impl Into<String>, children: Vec<Box<dyn BehaviorNode>>) -> Self {
171        Self {
172            name: name.into(),
173            children,
174            current_index: 0,
175        }
176    }
177}
178
179impl BehaviorNode for SelectorNode {
180    fn tick(&mut self, blackboard: &Blackboard) -> BehaviorStatus {
181        while self.current_index < self.children.len() {
182            match self.children[self.current_index].tick(blackboard) {
183                BehaviorStatus::Success => {
184                    self.reset();
185                    return BehaviorStatus::Success;
186                }
187                BehaviorStatus::Failure => {
188                    self.children[self.current_index].reset();
189                    self.current_index += 1;
190                }
191                BehaviorStatus::Running => return BehaviorStatus::Running,
192            }
193        }
194
195        self.reset();
196        BehaviorStatus::Failure
197    }
198
199    fn reset(&mut self) {
200        self.current_index = 0;
201        for child in &mut self.children {
202            child.reset();
203        }
204    }
205
206    fn name(&self) -> &str {
207        &self.name
208    }
209}
210
211pub fn condition(
212    name: impl Into<String>,
213    evaluator: impl Fn(&Blackboard) -> bool + 'static,
214) -> Box<dyn BehaviorNode> {
215    Box::new(ConditionNode::new(name, evaluator))
216}
217
218pub fn action(
219    name: impl Into<String>,
220    executor: impl FnMut(&Blackboard) -> BehaviorStatus + 'static,
221) -> Box<dyn BehaviorNode> {
222    Box::new(ActionNode::new(name, executor))
223}
224
225pub fn sequence(
226    name: impl Into<String>,
227    children: Vec<Box<dyn BehaviorNode>>,
228) -> Box<dyn BehaviorNode> {
229    Box::new(SequenceNode::new(name, children))
230}
231
232pub fn selector(
233    name: impl Into<String>,
234    children: Vec<Box<dyn BehaviorNode>>,
235) -> Box<dyn BehaviorNode> {
236    Box::new(SelectorNode::new(name, children))
237}
238
239pub fn create_demo_behavior_tree() -> BehaviorTree {
240    BehaviorTree::new(selector(
241        "mission_root",
242        vec![
243            sequence(
244                "navigate",
245                vec![
246                    condition("battery_ok", |bb| bb.battery_level > 20.0),
247                    condition("path_ready", |bb| bb.path_ready),
248                    condition("goal_visible", |bb| bb.goal_visible),
249                    action("move_to_goal", |_| BehaviorStatus::Success),
250                ],
251            ),
252            sequence(
253                "avoid",
254                vec![
255                    condition("obstacle_detected", |bb| bb.obstacle_detected),
256                    action("hover_and_replan", |_| BehaviorStatus::Running),
257                ],
258            ),
259            action("hold_position", |_| BehaviorStatus::Running),
260        ],
261    ))
262}
263
264pub fn demo_behavior_tree() {
265    let scenarios = [
266        (
267            "Nominal navigation",
268            Blackboard {
269                battery_level: 90.0,
270                obstacle_detected: false,
271                path_ready: true,
272                goal_visible: true,
273            },
274        ),
275        (
276            "Obstacle avoidance",
277            Blackboard {
278                battery_level: 90.0,
279                obstacle_detected: true,
280                path_ready: false,
281                goal_visible: false,
282            },
283        ),
284        (
285            "Low battery hold",
286            Blackboard {
287                battery_level: 5.0,
288                obstacle_detected: false,
289                path_ready: false,
290                goal_visible: false,
291            },
292        ),
293    ];
294
295    let mut tree = create_demo_behavior_tree();
296    println!("Behavior tree root: {}", tree.root_name());
297
298    for (label, blackboard) in scenarios {
299        tree.reset();
300        let status = tree.tick(&blackboard);
301        println!(
302            "{} => {:?} (battery={:.1}, obstacle={}, path_ready={}, goal_visible={})",
303            label,
304            status,
305            blackboard.battery_level,
306            blackboard.obstacle_detected,
307            blackboard.path_ready,
308            blackboard.goal_visible
309        );
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    #[test]
318    fn test_sequence_succeeds_when_all_children_succeed() {
319        let mut tree = BehaviorTree::new(sequence(
320            "sequence",
321            vec![
322                condition("battery_ok", |bb| bb.battery_level > 20.0),
323                action("move", |_| BehaviorStatus::Success),
324            ],
325        ));
326
327        let status = tree.tick(&Blackboard::default());
328
329        assert_eq!(status, BehaviorStatus::Success);
330    }
331
332    #[test]
333    fn test_selector_uses_fallback_branch() {
334        let mut tree = BehaviorTree::new(selector(
335            "selector",
336            vec![
337                condition("goal_visible", |bb| bb.goal_visible),
338                action("fallback", |_| BehaviorStatus::Success),
339            ],
340        ));
341
342        let status = tree.tick(&Blackboard {
343            goal_visible: false,
344            ..Default::default()
345        });
346
347        assert_eq!(status, BehaviorStatus::Success);
348    }
349
350    #[test]
351    fn test_running_action_resumes_until_success() {
352        let mut tree = BehaviorTree::new(sequence(
353            "warmup",
354            vec![
355                action("spin_up", {
356                    let mut count = 0;
357                    move |_| {
358                        count += 1;
359                        if count < 2 {
360                            BehaviorStatus::Running
361                        } else {
362                            BehaviorStatus::Success
363                        }
364                    }
365                }),
366                action("execute", |_| BehaviorStatus::Success),
367            ],
368        ));
369
370        assert_eq!(tree.tick(&Blackboard::default()), BehaviorStatus::Running);
371        assert_eq!(tree.tick(&Blackboard::default()), BehaviorStatus::Success);
372    }
373
374    #[test]
375    fn test_demo_tree_holds_position_on_low_battery() {
376        let mut tree = create_demo_behavior_tree();
377        let status = tree.tick(&Blackboard {
378            battery_level: 5.0,
379            obstacle_detected: false,
380            path_ready: false,
381            goal_visible: false,
382        });
383
384        assert_eq!(status, BehaviorStatus::Running);
385    }
386}