big_brain/
thinker.rs

1//! Thinkers are the "brain" of an entity. You attach Scorers to it, and the
2//! Thinker picks the right Action to run based on the resulting Scores.
3
4use std::{collections::VecDeque, sync::Arc};
5
6use bevy::{
7    prelude::*,
8    utils::{
9        tracing::{debug, field, span, Level, Span},
10        Duration, Instant,
11    },
12};
13
14#[cfg(feature = "trace")]
15use bevy::utils::tracing::trace;
16
17use crate::{
18    actions::{self, ActionBuilder, ActionBuilderWrapper, ActionState},
19    choices::{Choice, ChoiceBuilder},
20    pickers::Picker,
21    scorers::{Score, ScorerBuilder},
22};
23
24/// Wrapper for Actor entities. In terms of Scorers, Thinkers, and Actions,
25/// this is the [`Entity`] actually _performing_ the action, rather than the
26/// entity a Scorer/Thinker/Action is attached to. Generally, you will use
27/// this entity when writing Queries for Action and Scorer systems.
28#[derive(Debug, Clone, Component, Copy, Reflect)]
29pub struct Actor(pub Entity);
30
31#[derive(Debug, Clone, Copy, Reflect)]
32pub struct Action(pub Entity);
33
34impl Action {
35    pub fn entity(&self) -> Entity {
36        self.0
37    }
38}
39
40#[derive(Debug, Clone, Component)]
41pub struct ActionSpan {
42    pub(crate) span: Span,
43}
44
45impl ActionSpan {
46    pub(crate) fn new(action: Entity, label: Option<&str>) -> Self {
47        let span = span!(
48            Level::DEBUG,
49            "action",
50            ent = ?action,
51            label = field::Empty,
52        );
53        if let Some(label) = label {
54            span.record("label", label);
55        }
56        Self { span }
57    }
58
59    pub fn span(&self) -> &Span {
60        &self.span
61    }
62}
63
64#[derive(Debug, Clone, Copy, Reflect)]
65pub struct Scorer(pub Entity);
66
67#[derive(Debug, Clone, Component)]
68pub struct ScorerSpan {
69    pub(crate) span: Span,
70}
71
72impl ScorerSpan {
73    pub(crate) fn new(scorer: Entity, label: Option<&str>) -> Self {
74        let span = span!(
75            Level::DEBUG,
76            "scorer",
77            ent = ?scorer,
78            label = field::Empty,
79        );
80
81        if let Some(label) = label {
82            span.record("label", label);
83        }
84        Self { span }
85    }
86
87    pub fn span(&self) -> &Span {
88        &self.span
89    }
90}
91
92/// The "brains" behind this whole operation. A `Thinker` is what glues
93/// together `Actions` and `Scorers` and shapes larger, intelligent-seeming
94/// systems.
95///
96/// Note: Thinkers are also Actions, so anywhere you can pass in an Action (or
97/// [`ActionBuilder`]), you can pass in a Thinker (or [`ThinkerBuilder`]).
98///
99/// ### Example
100///
101/// ```
102/// # use bevy::prelude::*;
103/// # use big_brain::prelude::*;
104/// # #[derive(Component, Debug)]
105/// # struct Thirst(f32, f32);
106/// # #[derive(Component, Debug)]
107/// # struct Hunger(f32, f32);
108/// # #[derive(Clone, Component, Debug, ScorerBuilder)]
109/// # struct Thirsty;
110/// # #[derive(Clone, Component, Debug, ScorerBuilder)]
111/// # struct Hungry;
112/// # #[derive(Clone, Component, Debug, ActionBuilder)]
113/// # struct Drink;
114/// # #[derive(Clone, Component, Debug, ActionBuilder)]
115/// # struct Eat;
116/// # #[derive(Clone, Component, Debug, ActionBuilder)]
117/// # struct Meander;
118/// pub fn init_entities(mut cmd: Commands) {
119///     cmd.spawn((
120///         Thirst(70.0, 2.0),
121///         Hunger(50.0, 3.0),
122///         Thinker::build()
123///             .picker(FirstToScore::new(80.0))
124///             .when(Thirsty, Drink)
125///             .when(Hungry, Eat)
126///             .otherwise(Meander),
127///     ));
128/// }
129/// ```
130#[derive(Component, Debug, Reflect)]
131#[reflect(from_reflect = false)]
132pub struct Thinker {
133    #[reflect(ignore)]
134    picker: Arc<dyn Picker>,
135    #[reflect(ignore)]
136    otherwise: Option<ActionBuilderWrapper>,
137    #[reflect(ignore)]
138    choices: Vec<Choice>,
139    #[reflect(ignore)]
140    current_action: Option<(Action, ActionBuilderWrapper)>,
141    current_action_label: Option<Option<String>>,
142    #[reflect(ignore)]
143    span: Span,
144    #[reflect(ignore)]
145    scheduled_actions: VecDeque<ActionBuilderWrapper>,
146}
147
148impl Thinker {
149    /// Make a new [`ThinkerBuilder`]. This is what you'll actually use to
150    /// configure Thinker behavior.
151    pub fn build() -> ThinkerBuilder {
152        ThinkerBuilder::new()
153    }
154
155    pub fn schedule_action(&mut self, action: impl ActionBuilder + 'static) {
156        self.scheduled_actions
157            .push_back(ActionBuilderWrapper::new(Arc::new(action)));
158    }
159}
160
161/// This is what you actually use to configure Thinker behavior. It's a plain
162/// old [`ActionBuilder`], as well.
163#[derive(Component, Clone, Debug, Default)]
164pub struct ThinkerBuilder {
165    picker: Option<Arc<dyn Picker>>,
166    otherwise: Option<ActionBuilderWrapper>,
167    choices: Vec<ChoiceBuilder>,
168    label: Option<String>,
169}
170
171impl ThinkerBuilder {
172    pub(crate) fn new() -> Self {
173        Self {
174            picker: None,
175            otherwise: None,
176            choices: Vec::new(),
177            label: None,
178        }
179    }
180
181    /// Define a [`Picker`](crate::pickers::Picker) for this Thinker.
182    pub fn picker(mut self, picker: impl Picker + 'static) -> Self {
183        self.picker = Some(Arc::new(picker));
184        self
185    }
186
187    /// Define an [`ActionBuilder`](crate::actions::ActionBuilder) and
188    /// [`ScorerBuilder`](crate::scorers::ScorerBuilder) pair.
189    pub fn when(
190        mut self,
191        scorer: impl ScorerBuilder + 'static,
192        action: impl ActionBuilder + 'static,
193    ) -> Self {
194        self.choices
195            .push(ChoiceBuilder::new(Arc::new(scorer), Arc::new(action)));
196        self
197    }
198
199    /// Default `Action` to execute if the `Picker` did not pick any of the
200    /// given choices.
201    pub fn otherwise(mut self, otherwise: impl ActionBuilder + 'static) -> Self {
202        self.otherwise = Some(ActionBuilderWrapper::new(Arc::new(otherwise)));
203        self
204    }
205
206    /// * Configures a label to use for the thinker when logging.
207    pub fn label(mut self, label: impl AsRef<str>) -> Self {
208        self.label = Some(label.as_ref().to_string());
209        self
210    }
211}
212
213impl ActionBuilder for ThinkerBuilder {
214    fn build(&self, cmd: &mut Commands, action_ent: Entity, actor: Entity) {
215        let span = span!(
216            Level::DEBUG,
217            "thinker",
218            actor = ?actor,
219        );
220        let _guard = span.enter();
221        debug!("Spawning Thinker.");
222        let choices = self
223            .choices
224            .iter()
225            .map(|choice| choice.build(cmd, actor, action_ent))
226            .collect();
227        std::mem::drop(_guard);
228        cmd.entity(action_ent)
229            .insert(Thinker {
230                // TODO: reasonable default?...
231                picker: self
232                    .picker
233                    .clone()
234                    .expect("ThinkerBuilder must have a Picker"),
235                otherwise: self.otherwise.clone(),
236                choices,
237                current_action: None,
238                current_action_label: None,
239                span,
240                scheduled_actions: VecDeque::new(),
241            })
242            .insert(Name::new("Thinker"))
243            .insert(ActionState::Requested);
244    }
245
246    fn label(&self) -> Option<&str> {
247        self.label.as_deref()
248    }
249}
250
251pub fn thinker_component_attach_system(
252    mut cmd: Commands,
253    q: Query<(Entity, &ThinkerBuilder), Without<HasThinker>>,
254) {
255    for (entity, thinker_builder) in q.iter() {
256        let thinker = actions::spawn_action(thinker_builder, &mut cmd, entity);
257        cmd.entity(entity).insert(HasThinker(thinker));
258    }
259}
260
261pub fn thinker_component_detach_system(
262    mut cmd: Commands,
263    q: Query<(Entity, &HasThinker), Without<ThinkerBuilder>>,
264) {
265    for (actor, HasThinker(thinker)) in q.iter() {
266        if let Some(ent) = cmd.get_entity(*thinker) {
267            ent.despawn_recursive();
268        }
269        cmd.entity(actor).remove::<HasThinker>();
270    }
271}
272
273pub fn actor_gone_cleanup(
274    mut cmd: Commands,
275    actors: Query<&ThinkerBuilder>,
276    q: Query<(Entity, &Actor)>,
277) {
278    for (child, Actor(actor)) in q.iter() {
279        if actors.get(*actor).is_err() {
280            // Actor is gone. Let's clean up.
281            if let Some(ent) = cmd.get_entity(child) {
282                ent.despawn_recursive();
283            }
284        }
285    }
286}
287
288#[derive(Component, Debug, Reflect)]
289pub struct HasThinker(Entity);
290
291impl HasThinker {
292    pub fn entity(&self) -> Entity {
293        self.0
294    }
295}
296
297pub struct ThinkerIterations {
298    index: usize,
299    max_duration: Duration,
300}
301impl ThinkerIterations {
302    pub fn new(max_duration: Duration) -> Self {
303        Self {
304            index: 0,
305            max_duration,
306        }
307    }
308}
309impl Default for ThinkerIterations {
310    fn default() -> Self {
311        Self::new(Duration::from_millis(10))
312    }
313}
314
315pub fn thinker_system(
316    mut cmd: Commands,
317    mut iterations: Local<ThinkerIterations>,
318    mut thinker_q: Query<(Entity, &Actor, &mut Thinker)>,
319    scores: Query<&Score>,
320    mut action_states: Query<&mut actions::ActionState>,
321    action_spans: Query<&ActionSpan>,
322    scorer_spans: Query<&ScorerSpan>,
323) {
324    let start = Instant::now();
325    for (thinker_ent, Actor(actor), mut thinker) in thinker_q.iter_mut().skip(iterations.index) {
326        iterations.index += 1;
327
328        let thinker_state = action_states
329            .get_mut(thinker_ent)
330            .expect("Where is it?")
331            .clone();
332
333        let thinker_span = thinker.span.clone();
334        let _thinker_span_guard = thinker_span.enter();
335
336        match thinker_state {
337            ActionState::Init => {
338                let mut act_state = action_states.get_mut(thinker_ent).expect("???");
339                debug!("Initializing thinker.");
340                *act_state = ActionState::Requested;
341            }
342            ActionState::Requested => {
343                let mut act_state = action_states.get_mut(thinker_ent).expect("???");
344                debug!("Thinker requested. Starting execution.");
345                *act_state = ActionState::Executing;
346            }
347            ActionState::Success | ActionState::Failure => {}
348            ActionState::Cancelled => {
349                debug!("Thinker cancelled. Cleaning up.");
350                if let Some(current) = &mut thinker.current_action {
351                    let action_span = action_spans.get(current.0 .0).expect("Where is it?");
352                    debug!("Cancelling current action because thinker was cancelled.");
353                    let state = action_states.get_mut(current.0.0).expect("Couldn't find a component corresponding to the current action. This is definitely a bug.").clone();
354                    match state {
355                        ActionState::Success | ActionState::Failure => {
356                            debug!("Action already wrapped up on its own. Cleaning up action in Thinker.");
357                            if let Some(ent) = cmd.get_entity(current.0 .0) {
358                                ent.despawn_recursive();
359                            }
360                            thinker.current_action = None;
361                        }
362                        ActionState::Cancelled => {
363                            debug!("Current action already cancelled.");
364                        }
365                        _ => {
366                            let mut state = action_states.get_mut(current.0.0).expect("Couldn't find a component corresponding to the current action. This is definitely a bug.");
367                            debug!( "Action is still executing. Attempting to cancel it before wrapping up Thinker cancellation.");
368                            action_span.span.in_scope(|| {
369                                debug!("Parent thinker was cancelled. Cancelling action.");
370                            });
371                            *state = ActionState::Cancelled;
372                        }
373                    }
374                } else {
375                    let mut act_state = action_states.get_mut(thinker_ent).expect("???");
376                    debug!("No current thinker action. Wrapping up Thinker as Succeeded.");
377                    *act_state = ActionState::Success;
378                }
379            }
380            ActionState::Executing => {
381                #[cfg(feature = "trace")]
382                trace!("Thinker is executing. Thinking...");
383                if let Some(choice) = thinker.picker.pick(&thinker.choices, &scores) {
384                    // Think about what action we're supposed to be taking. We do this
385                    // every tick, because we might change our mind.
386                    // ...and then execute it (details below).
387                    #[cfg(feature = "trace")]
388                    trace!("Action picked. Executing picked action.");
389                    let action = choice.action.clone();
390                    let scorer = choice.scorer;
391                    let score = scores.get(choice.scorer.0).expect("Where is it?");
392                    exec_picked_action(
393                        &mut cmd,
394                        *actor,
395                        &mut thinker,
396                        &action,
397                        &mut action_states,
398                        &action_spans,
399                        Some((&scorer, score)),
400                        &scorer_spans,
401                        true,
402                    );
403                } else if should_schedule_action(&mut thinker, &mut action_states) {
404                    debug!("Spawning scheduled action.");
405                    let action = thinker
406                        .scheduled_actions
407                        .pop_front()
408                        .expect("we literally just checked if it was there.");
409                    let new_action = actions::spawn_action(action.1.as_ref(), &mut cmd, *actor);
410                    thinker.current_action = Some((Action(new_action), action.clone()));
411                    thinker.current_action_label = Some(action.1.label().map(|s| s.into()));
412                } else if let Some(default_action_ent) = &thinker.otherwise {
413                    // Otherwise, let's just execute the default one! (if it's there)
414                    let default_action_ent = default_action_ent.clone();
415                    exec_picked_action(
416                        &mut cmd,
417                        *actor,
418                        &mut thinker,
419                        &default_action_ent,
420                        &mut action_states,
421                        &action_spans,
422                        None,
423                        &scorer_spans,
424                        false,
425                    );
426                } else if let Some((action_ent, _)) = &thinker.current_action {
427                    let action_span = action_spans.get(action_ent.0).expect("Where is it?");
428                    let _guard = action_span.span.enter();
429                    let mut curr_action_state = action_states.get_mut(action_ent.0).expect("Couldn't find a component corresponding to the current action. This is definitely a bug.");
430                    let previous_done = matches!(
431                        *curr_action_state,
432                        ActionState::Success | ActionState::Failure
433                    );
434                    if previous_done {
435                        debug!(
436                            "Action completed and nothing was picked. Despawning action entity.",
437                        );
438                        // Despawn the action itself.
439                        if let Some(ent) = cmd.get_entity(action_ent.0) {
440                            ent.despawn_recursive();
441                        }
442                        thinker.current_action = None;
443                    } else if *curr_action_state == ActionState::Init {
444                        *curr_action_state = ActionState::Requested;
445                    }
446                }
447            }
448        }
449        if iterations.index % 500 == 0 && start.elapsed() > iterations.max_duration {
450            return;
451        }
452    }
453    iterations.index = 0;
454}
455
456fn should_schedule_action(
457    thinker: &mut Mut<Thinker>,
458    states: &mut Query<&mut ActionState>,
459) -> bool {
460    #[cfg(feature = "trace")]
461    let thinker_span = thinker.span.clone();
462    #[cfg(feature = "trace")]
463    let _thinker_span_guard = thinker_span.enter();
464    if thinker.scheduled_actions.is_empty() {
465        #[cfg(feature = "trace")]
466        trace!("No scheduled actions. Not scheduling anything.");
467        false
468    } else if let Some((action_ent, _)) = &mut thinker.current_action {
469        let curr_action_state = states.get_mut(action_ent.0).expect("Couldn't find a component corresponding to the current action. This is definitely a bug.");
470
471        let action_done = matches!(
472            *curr_action_state,
473            ActionState::Success | ActionState::Failure
474        );
475
476        #[cfg(feature = "trace")]
477        if action_done {
478            trace!("Current action is already done. Can schedule.");
479        } else {
480            trace!("Current action is still executing. Not scheduling anything.");
481        }
482
483        action_done
484    } else {
485        #[cfg(feature = "trace")]
486        trace!("No current action actions. Can schedule.");
487        true
488    }
489}
490
491#[allow(clippy::too_many_arguments)]
492fn exec_picked_action(
493    cmd: &mut Commands,
494    actor: Entity,
495    thinker: &mut Mut<Thinker>,
496    picked_action: &ActionBuilderWrapper,
497    states: &mut Query<&mut ActionState>,
498    action_spans: &Query<&ActionSpan>,
499    scorer_info: Option<(&Scorer, &Score)>,
500    scorer_spans: &Query<&ScorerSpan>,
501    override_current: bool,
502) {
503    // If we do find one, then we need to grab the corresponding
504    // component for it. The "action" that `picker.pick()` returns
505    // is just a newtype for an Entity.
506    //
507
508    // Now we check the current action. We need to check if we picked the same one as the previous tick.
509    //
510    // TODO: I don't know where the right place to put this is
511    // (maybe not in this logic), but we do need some kind of
512    // oscillation protection so we're not just bouncing back and
513    // forth between the same couple of actions.
514    let thinker_span = thinker.span.clone();
515    let _thinker_span_guard = thinker_span.enter();
516    if let Some((action_ent, ActionBuilderWrapper(current_id, _))) = &mut thinker.current_action {
517        let mut curr_action_state = states.get_mut(action_ent.0).expect("Couldn't find a component corresponding to the current action. This is definitely a bug.");
518        let previous_done = matches!(
519            *curr_action_state,
520            ActionState::Success | ActionState::Failure
521        );
522        let action_span = action_spans.get(action_ent.0).expect("Where is it?");
523        let _guard = action_span.span.enter();
524        if (!Arc::ptr_eq(current_id, &picked_action.0) && override_current) || previous_done {
525            // So we've picked a different action than we were
526            // currently executing. Just like before, we grab the
527            // actual Action component (and we assume it exists).
528            // If the action is executing, or was requested, we
529            // need to cancel it to make sure it stops.
530            if !previous_done {
531                if override_current {
532                    #[cfg(feature = "trace")]
533                    trace!("Falling back to `otherwise` clause.",);
534                } else {
535                    #[cfg(feature = "trace")]
536                    trace!("Picked a different action than the current one.",);
537                }
538            }
539            match *curr_action_state {
540                ActionState::Executing | ActionState::Requested => {
541                    debug!("Previous action is still executing. Requesting action cancellation.",);
542                    *curr_action_state = ActionState::Cancelled;
543                }
544                ActionState::Init | ActionState::Success | ActionState::Failure => {
545                    debug!("Previous action already completed. Despawning action entity.",);
546                    // Despawn the action itself.
547                    if let Some(ent) = cmd.get_entity(action_ent.0) {
548                        ent.despawn_recursive();
549                    }
550                    if let Some((Scorer(ent), score)) = scorer_info {
551                        let scorer_span = scorer_spans.get(*ent).expect("Where is it?");
552                        let _guard = scorer_span.span.enter();
553                        debug!("Winning scorer chosen with score {}", score.get());
554                    }
555                    std::mem::drop(_guard);
556                    debug!("Spawning next action");
557                    let new_action =
558                        Action(actions::spawn_action(picked_action.1.as_ref(), cmd, actor));
559                    thinker.current_action = Some((new_action, picked_action.clone()));
560                    thinker.current_action_label = Some(picked_action.1.label().map(|s| s.into()));
561                }
562                ActionState::Cancelled => {
563                    #[cfg(feature = "trace")]
564                    trace!(
565                    "Cancellation already requested. Waiting for action to be marked as completed.",
566                )
567                }
568            };
569        } else {
570            // Otherwise, it turns out we want to keep executing
571            // the same action. Just in case, we go ahead and set
572            // it as Requested if for some reason it had finished
573            // but the Action System hasn't gotten around to
574            // cleaning it up.
575            if *curr_action_state == ActionState::Init {
576                *curr_action_state = ActionState::Requested;
577            }
578            #[cfg(feature = "trace")]
579            trace!("Continuing execution of current action.",)
580        }
581    } else {
582        #[cfg(feature = "trace")]
583        trace!("Falling back to `otherwise` clause.",);
584
585        // This branch arm is called when there's no
586        // current_action in the thinker. The logic here is pretty
587        // straightforward -- we set the action, Request it, and
588        // that's it.
589        if let Some((Scorer(ent), score)) = scorer_info {
590            let scorer_span = scorer_spans.get(*ent).expect("Where is it?");
591            let _guard = scorer_span.span.enter();
592            debug!("Winning scorer chosen with score {}", score.get());
593        }
594        debug!("No current action. Spawning new action.");
595        let new_action = actions::spawn_action(picked_action.1.as_ref(), cmd, actor);
596        thinker.current_action = Some((Action(new_action), picked_action.clone()));
597        thinker.current_action_label = Some(picked_action.1.label().map(|s| s.into()));
598    }
599}