diff --git a/game/entity_factory.c b/game/entity_factory.c
index 942d2a7..3d72b0e 100644
--- a/game/entity_factory.c
+++ b/game/entity_factory.c
@@ -53,11 +53,6 @@ ecs_entity_t entityCreateBaseUnit(const Position position, f32 size, Player play
             .elapsed = 0.0f,
         });
     }
-    ecs_set(ECS, e, BzBTState, {
-        .root = NULL,
-        .nodeStatePool = game->pools.btNodeState
-    });
-    ecs_set(ECS, e, AIBlackboard, {.entity = e});
     ecs_add_id(ECS, e, Selectable);
     Unit unitComp = {
             .attackElapsed = 0.0f,
diff --git a/game/systems/s_ai.c b/game/systems/s_ai.c
index 41f2898..29a80ad 100644
--- a/game/systems/s_ai.c
+++ b/game/systems/s_ai.c
@@ -35,20 +35,23 @@ void updateAISystem(ecs_iter_t *it) {
 void setAIBehaviour(ecs_entity_t entity, const BzBTNode *root,
                     const AIBlackboard *blackboard) {
     Game *game = ecs_singleton_get_mut(ECS, Game);
-    BZ_ASSERT(ecs_has(ECS, entity, BzBTState));
-    BZ_ASSERT(!blackboard || ecs_has(ECS, entity, AIBlackboard));
-
-    if (blackboard) {
-        AIBlackboard *b = ecs_get_mut(ECS, entity, AIBlackboard);
-        *b = *blackboard;
-        if (b->proximity < 2.0f)
-            b->proximity = 2.0f;
+    if (ecs_has(ECS, entity, BzBTState)) {
+        BzBTState *prevState = ecs_get_mut(ECS, entity, BzBTState);
+        bzBTDestroyState(prevState);
     }
 
-    BzBTState *state = ecs_get_mut(ECS, entity, BzBTState);
-    bzBTDestroyState(state);
-    *state = bzBTCreateState(&(BzBTStateDesc) {
+    ecs_set(ECS, entity, BzBTState, {
         .root = root,
-        .pool = game->pools.btNodeState
+        .nodeStatePool = game->pools.btNodeState,
     });
+    AIBlackboard defaultBB = {
+            .entity = 0
+    };
+    if (blackboard) {
+        defaultBB = *blackboard;
+    }
+    defaultBB.entity = entity;
+    if (defaultBB.proximity < 2.0f)
+        defaultBB.proximity = 2.0f;
+    ecs_set_ptr(ECS, entity, AIBlackboard, &defaultBB);
 }