Skip to main content

mecha10_behavior_patterns/
ensemble.rs

1//! Ensemble pattern for multi-model fusion
2//!
3//! This module implements ensemble learning patterns where multiple AI models
4//! can be combined to make more robust decisions.
5
6use async_trait::async_trait;
7use mecha10_behavior_runtime::{BehaviorNode, BoxedBehavior, NodeStatus};
8use mecha10_core::Context;
9use serde::{Deserialize, Serialize};
10use tracing::{debug, info};
11
12/// A weighted model in an ensemble.
13#[derive(Debug)]
14pub struct WeightedModel {
15    /// The behavior/model
16    pub behavior: BoxedBehavior,
17    /// Weight for voting (only used in WeightedVote strategy)
18    pub weight: f32,
19    /// Optional name for debugging
20    pub name: Option<String>,
21}
22
23impl WeightedModel {
24    /// Create a new weighted model.
25    pub fn new(behavior: BoxedBehavior, weight: f32) -> Self {
26        Self {
27            behavior,
28            weight,
29            name: None,
30        }
31    }
32
33    mecha10_core::builder_setter!(
34        /// Set the name of this model.
35        into_some, with_name, name, impl Into<String>
36    );
37}
38
39/// Strategy for combining results from multiple models.
40#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
41#[serde(rename_all = "snake_case")]
42pub enum EnsembleStrategy {
43    /// Weighted voting: success if weighted sum of successes > threshold (default 0.5)
44    WeightedVote,
45
46    /// Conservative: all models must succeed
47    Conservative,
48
49    /// Optimistic: any model can succeed
50    Optimistic,
51
52    /// Majority: more than half must succeed
53    Majority,
54}
55
56/// Ensemble node for multi-model fusion.
57///
58/// This node runs multiple behaviors/models in parallel and combines their results
59/// using a configurable strategy. This is useful for:
60/// - Combining multiple AI models for robust predictions
61/// - Redundancy and fault tolerance
62/// - Multi-modal sensor fusion
63///
64/// # Strategies
65///
66/// - **WeightedVote**: Models vote with weights, success if total > threshold
67/// - **Conservative**: All models must succeed (high precision, low recall)
68/// - **Optimistic**: Any model can succeed (high recall, low precision)
69/// - **Majority**: More than half must succeed (balanced)
70///
71/// # Example
72///
73/// ```rust
74/// use mecha10_behavior_patterns::prelude::*;
75///
76/// # #[derive(Debug)]
77/// # struct YoloV8;
78/// # #[async_trait]
79/// # impl BehaviorNode for YoloV8 {
80/// #     async fn tick(&mut self, _ctx: &Context) -> anyhow::Result<NodeStatus> {
81/// #         Ok(NodeStatus::Success)
82/// #     }
83/// # }
84/// # #[derive(Debug)]
85/// # struct YoloV10;
86/// # #[async_trait]
87/// # impl BehaviorNode for YoloV10 {
88/// #     async fn tick(&mut self, _ctx: &Context) -> anyhow::Result<NodeStatus> {
89/// #         Ok(NodeStatus::Success)
90/// #     }
91/// # }
92/// # #[derive(Debug)]
93/// # struct CustomDetector;
94/// # #[async_trait]
95/// # impl BehaviorNode for CustomDetector {
96/// #     async fn tick(&mut self, _ctx: &Context) -> anyhow::Result<NodeStatus> {
97/// #         Ok(NodeStatus::Success)
98/// #     }
99/// # }
100///
101/// let ensemble = EnsembleNode::new(EnsembleStrategy::WeightedVote)
102///     .add_model(Box::new(YoloV8), 0.4)
103///     .add_model(Box::new(YoloV10), 0.3)
104///     .add_model(Box::new(CustomDetector), 0.3);
105/// ```
106#[derive(Debug)]
107pub struct EnsembleNode {
108    models: Vec<WeightedModel>,
109    strategy: EnsembleStrategy,
110    threshold: f32,
111    model_statuses: Vec<NodeStatus>,
112}
113
114impl EnsembleNode {
115    /// Create a new ensemble node with the given strategy.
116    pub fn new(strategy: EnsembleStrategy) -> Self {
117        Self {
118            models: Vec::new(),
119            strategy,
120            threshold: 0.5,
121            model_statuses: Vec::new(),
122        }
123    }
124
125    /// Add a model with a weight.
126    pub fn add_model(mut self, behavior: BoxedBehavior, weight: f32) -> Self {
127        self.models.push(WeightedModel::new(behavior, weight));
128        self.model_statuses.push(NodeStatus::Running);
129        self
130    }
131
132    /// Set the threshold for weighted voting (default 0.5).
133    pub fn with_threshold(mut self, threshold: f32) -> Self {
134        self.threshold = threshold;
135        self
136    }
137}
138
139#[async_trait]
140impl BehaviorNode for EnsembleNode {
141    async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
142        if self.models.is_empty() {
143            return Ok(NodeStatus::Failure);
144        }
145
146        // Execute all models in parallel
147        for (i, model) in self.models.iter_mut().enumerate() {
148            if self.model_statuses[i].is_running() {
149                let status = model.behavior.tick(ctx).await?;
150                self.model_statuses[i] = status;
151
152                let model_name = model.name.as_deref().unwrap_or("unnamed");
153                debug!("Ensemble: model '{}' status: {}", model_name, status);
154            }
155        }
156
157        // Combine results based on strategy
158        let overall_status = match self.strategy {
159            EnsembleStrategy::Conservative => self.conservative_combine(),
160            EnsembleStrategy::Optimistic => self.optimistic_combine(),
161            EnsembleStrategy::Majority => self.majority_combine(),
162            EnsembleStrategy::WeightedVote => self.weighted_vote_combine(),
163        };
164
165        info!("Ensemble ({:?}): overall status = {}", self.strategy, overall_status);
166        Ok(overall_status)
167    }
168
169    async fn reset(&mut self) -> anyhow::Result<()> {
170        for (i, model) in self.models.iter_mut().enumerate() {
171            model.behavior.reset().await?;
172            self.model_statuses[i] = NodeStatus::Running;
173        }
174        Ok(())
175    }
176
177    fn name(&self) -> &str {
178        "ensemble"
179    }
180
181    fn children_mut(&mut self) -> Vec<&mut dyn BehaviorNode> {
182        self.models
183            .iter_mut()
184            .map(|model| model.behavior.as_mut() as &mut dyn BehaviorNode)
185            .collect()
186    }
187}
188
189impl EnsembleNode {
190    fn conservative_combine(&self) -> NodeStatus {
191        // All must succeed
192        let failures = self.model_statuses.iter().filter(|s| s.is_failure()).count();
193        let running = self.model_statuses.iter().filter(|s| s.is_running()).count();
194
195        if failures > 0 {
196            NodeStatus::Failure
197        } else if running > 0 {
198            NodeStatus::Running
199        } else {
200            NodeStatus::Success
201        }
202    }
203
204    fn optimistic_combine(&self) -> NodeStatus {
205        // Any can succeed
206        let successes = self.model_statuses.iter().filter(|s| s.is_success()).count();
207        let running = self.model_statuses.iter().filter(|s| s.is_running()).count();
208
209        if successes > 0 {
210            NodeStatus::Success
211        } else if running > 0 {
212            NodeStatus::Running
213        } else {
214            NodeStatus::Failure
215        }
216    }
217
218    fn majority_combine(&self) -> NodeStatus {
219        // More than half must succeed
220        let successes = self.model_statuses.iter().filter(|s| s.is_success()).count();
221        let failures = self.model_statuses.iter().filter(|s| s.is_failure()).count();
222        let running = self.model_statuses.iter().filter(|s| s.is_running()).count();
223
224        let _total_decided = successes + failures;
225        let majority_threshold = self.models.len().div_ceil(2);
226
227        if successes >= majority_threshold {
228            NodeStatus::Success
229        } else if failures >= majority_threshold {
230            NodeStatus::Failure
231        } else if running > 0 {
232            NodeStatus::Running
233        } else {
234            // Tie - default to failure
235            NodeStatus::Failure
236        }
237    }
238
239    fn weighted_vote_combine(&self) -> NodeStatus {
240        // Weighted sum of successes must exceed threshold
241        let mut success_weight = 0.0;
242        let mut _failure_weight = 0.0;
243        let mut total_weight = 0.0;
244        let mut any_running = false;
245
246        for (model, status) in self.models.iter().zip(&self.model_statuses) {
247            total_weight += model.weight;
248            match status {
249                NodeStatus::Success => success_weight += model.weight,
250                NodeStatus::Failure => _failure_weight += model.weight,
251                NodeStatus::Running => any_running = true,
252            }
253        }
254
255        if any_running {
256            return NodeStatus::Running;
257        }
258
259        let success_ratio = if total_weight > 0.0 {
260            success_weight / total_weight
261        } else {
262            0.0
263        };
264
265        if success_ratio >= self.threshold {
266            NodeStatus::Success
267        } else {
268            NodeStatus::Failure
269        }
270    }
271}
272
273/// Configuration for an ensemble model (for JSON deserialization).
274#[allow(dead_code)] // Will be used when implementing JSON loading
275#[derive(Debug, Clone, Serialize, Deserialize)]
276pub struct EnsembleModelConfig {
277    /// Node type to instantiate
278    pub node: String,
279    /// Weight for this model
280    pub weight: f32,
281    /// Optional name for this model
282    #[serde(skip_serializing_if = "Option::is_none")]
283    pub name: Option<String>,
284    /// Configuration for the node
285    #[serde(default)]
286    pub config: serde_json::Value,
287}