diff --git a/engine/breeze/ai/behaviour_tree.c b/engine/breeze/ai/behaviour_tree.c
index 4eea827..f8ec0d6 100644
--- a/engine/breeze/ai/behaviour_tree.c
+++ b/engine/breeze/ai/behaviour_tree.c
@@ -203,6 +203,11 @@ BzBTNode *bzBTNodeNext(const BzBTNode *node) {
     return node->next;
 }
 
+const BzBTNode *bzBTNodeStateGetNode(const BzBTNodeState *state) {
+    BZ_ASSERT(state);
+    return state->node;
+}
+
 const BzBTNodeState *bzBTNodeStateNext(const BzBTNodeState *state) {
     BZ_ASSERT(state);
     return state->next;
diff --git a/engine/breeze/ai/behaviour_tree.h b/engine/breeze/ai/behaviour_tree.h
index 4ac09e1..600559b 100644
--- a/engine/breeze/ai/behaviour_tree.h
+++ b/engine/breeze/ai/behaviour_tree.h
@@ -92,6 +92,8 @@ BzBTNodeType bzBTGetNodeType(const BzBTNode *node);
 BzBTNode *bzBTNodeChild(const BzBTNode *node);
 BzBTNode *bzBTNodeNext(const BzBTNode *node);
 
+const BzBTNode *bzBTNodeStateGetNode(const BzBTNodeState *state);
+
 const BzBTNodeState *bzBTNodeStateNext(const BzBTNodeState *state);
 bool bzBTNodeMatchesState(const BzBTNode *node, const BzBTNodeState *state);
 
diff --git a/game/components.c b/game/components.c
index 93cc3cc..01c703a 100644
--- a/game/components.c
+++ b/game/components.c
@@ -227,8 +227,7 @@ static const BzBTNodeState *findNodeState(const BzBTNode *node, const BzBTNodeSt
     }
     return NULL;
 }
-void igVisualizeBTState(const BzBTNode *node, const BzBTNodeState *state,
-                        bool isActive, bool sameLine, i32 depth)  {
+static void visualizeBTState(const BzBTNode *node, const BzBTNodeState *state, bool sameLine, bool isActive, i32 depth) {
     const BzBTNode *child = bzBTNodeChild(node);
     BzBTNodeType type = bzBTGetNodeType(node);
     char extraInfo[128];
@@ -236,7 +235,7 @@ void igVisualizeBTState(const BzBTNode *node, const BzBTNodeState *state,
 
     const BzBTNodeState *nodeState = findNodeState(node, state);
     bool hasState = nodeState != NULL;
-    isActive |= hasState;
+    isActive |= nodeState != NULL;
 
     switch (type) {
         case BZ_BT_DECOR_REPEAT:
@@ -249,7 +248,7 @@ void igVisualizeBTState(const BzBTNode *node, const BzBTNodeState *state,
             }
             break;
         case BZ_BT_DECOR_DELAY:
-            if (hasState) {
+            if (isActive) {
                 snprintf(extraInfo, sizeof(extraInfo), " (%.2f < %.2fms)",
                          bzBTDelayStateGetElapsed(nodeState),
                          bzBTDecorGetDelay(node));
@@ -271,35 +270,43 @@ void igVisualizeBTState(const BzBTNode *node, const BzBTNodeState *state,
     if (isActive)
         color = (ImVec4) {1.0f, 1.0f, 0.5f, 1.0f};
 
-    bool hasSingleChild = true;
-    if (child && bzBTNodeNext(child)) hasSingleChild = false;
+    bool hasSingleChild = child && bzBTNodeNext(child) == NULL;
 
-    const char *suffix = hasSingleChild ? " > " : ": ";
-
-    if (sameLine) {
-        igTextColored(color, "%s%s%s", bzBTNodeTypeToStr(type),
-                      extraInfo, suffix);
-    } else {
-        igTextColored(color, "%*s%s %s",
-                      depth * 2, "",
-                      bzBTNodeTypeToStr(type), extraInfo, suffix);
-        depth++;
-    }
+    i32 offset = depth * 2;
+    if (sameLine) offset = 0;
+    const char *prefix = sameLine ? " > " : "";
+    const char *postfix = child != NULL && !hasSingleChild ? ":" : "";
+    igTextColored(color, "%s%*s%s%s%s", prefix, offset, "", bzBTNodeTypeToStr(type), extraInfo, postfix);
 
     bool isComposite = type == BZ_BT_COMP_SELECTOR ||
                        type == BZ_BT_COMP_P_SELECTOR ||
                        type == BZ_BT_COMP_SEQUENCE ||
                        type == BZ_BT_COMP_P_SEQUENCE;
-
+    depth++;
     while (child) {
         if (hasSingleChild) igSameLine(0, 0);
-        bool childActive = isActive && hasSingleChild;
-        if (hasState && isComposite && !childActive)
-            childActive = bzBTCompStateGetRunningChild(state) == child;
-        igVisualizeBTState(child, state, childActive, hasSingleChild, depth);
+        bool childHasState = hasState && isComposite && bzBTCompStateGetRunningChild(nodeState) == child;
+        visualizeBTState(child, state, hasSingleChild, childHasState, depth);
         child = bzBTNodeNext(child);
     }
 }
+void igVisualizeBTState(const BzBTNode *node, const BzBTNodeState *state)  {
+    igSeparatorText("Tree State");
+    visualizeBTState(node, state, false, true, 0);
+
+    igSeparatorText("State stack");
+    size_t stateSize = 0;
+    const BzBTNodeState *pState = state;
+    while (pState) {
+        const BzBTNode *pNode = bzBTNodeStateGetNode(pState);
+        stateSize += bzBTGetNodeStateSize();
+        BzBTNodeType type = bzBTGetNodeType(pNode);
+        igText("%s", bzBTNodeTypeToStr(type));
+        pState = bzBTNodeStateNext(pState);
+    }
+    igNewLine();
+    igText("Total stack state size: %ld bytes.", stateSize);
+}
 void igBzBTState(ecs_world_t *ecs,
                  ecs_entity_t entity, ecs_entity_t comp) {
 
diff --git a/game/components.h b/game/components.h
index 93f14ae..10f9f87 100644
--- a/game/components.h
+++ b/game/components.h
@@ -235,8 +235,7 @@ void igArms(ecs_world_t *ecs,
 void igArm(ecs_world_t *ecs,
            ecs_entity_t entity, ecs_entity_t comp);
 
-void igVisualizeBTState(const BzBTNode *node, const BzBTNodeState *state,
-                        bool isActive, bool sameLine, i32 depth);
+void igVisualizeBTState(const BzBTNode *node, const BzBTNodeState *state);
 void igBzBTState(ecs_world_t *ecs,
                  ecs_entity_t entity, ecs_entity_t comp);
 void igAIBlackboard(ecs_world_t *ecs,
diff --git a/game/main.c b/game/main.c
index c498dd1..28602d9 100644
--- a/game/main.c
+++ b/game/main.c
@@ -624,7 +624,7 @@ static void renderMainMenu(Game *game, float dt) {
     if (uiMainMenuButton("Play")) {
         setScreen(game, SCREEN_GAME);
         unloadMap(game);
-        loadMap(game, "assets/maps/pathing_test.tmj");
+        loadMap(game, "assets/maps/tree_test.tmj");
     }
     if (uiMainMenuButton("Settings")) {
         setScreen(game, SCREEN_SETTINGS);
@@ -769,7 +769,7 @@ void igInspectWindow(ecs_entity_t entity, bool *open) {
             igCollapsingHeader_TreeNodeFlags("BehaviourTree", 0)) {
             const BzBTState *state = ecs_get(ECS, entity, BzBTState);
             if (state->root)
-                igVisualizeBTState(state->root, state->_first, true, false, 0);
+                igVisualizeBTState(state->root, state->_first);
             else
                 igTextColored((ImVec4) {1, 0, 0, 1}, "NONE");
         }