diff --git a/Source/Engine/AI/BehaviorKnowledge.cpp b/Source/Engine/AI/BehaviorKnowledge.cpp
index c08a09ffc..fc2504688 100644
--- a/Source/Engine/AI/BehaviorKnowledge.cpp
+++ b/Source/Engine/AI/BehaviorKnowledge.cpp
@@ -144,10 +144,17 @@ void BehaviorKnowledge::FreeMemory()
     {
         // Release any outstanding nodes state and memory
         ASSERT_LOW_LAYER(Tree);
+        BehaviorUpdateContext context;
+        context.Behavior = Behavior;
+        context.Knowledge = this;
+        context.Memory = Memory;
+        context.RelevantNodes = &RelevantNodes;
+        context.DeltaTime = 0.0f;
+        context.Time = 0.0f;
         for (const auto& node : Tree->Graph.Nodes)
         {
             if (node.Instance && node.Instance->_executionIndex != -1 && RelevantNodes[node.Instance->_executionIndex])
-                node.Instance->ReleaseState(Behavior, Memory);
+                node.Instance->ReleaseState(context);
         }
         Allocator::Free(Memory);
         Memory = nullptr;
diff --git a/Source/Engine/AI/BehaviorTreeNode.h b/Source/Engine/AI/BehaviorTreeNode.h
index 192db8f7e..3ea05977e 100644
--- a/Source/Engine/AI/BehaviorTreeNode.h
+++ b/Source/Engine/AI/BehaviorTreeNode.h
@@ -51,20 +51,18 @@ public:
     }
 
     /// <summary>
-    /// Initializes node instance state. Called when starting logic simulation for a given behavior.
+    /// Initializes node instance state. Called when starting logic simulation for a given behavior. Call constructor of the state container.
     /// </summary>
-    /// <param name="behavior">Behavior to simulate.</param>
-    /// <param name="memory">Pointer to pre-allocated memory for this node to use (call constructor of the state container).</param>
-    API_FUNCTION() virtual void InitState(Behavior* behavior, void* memory)
+    /// <param name="context">Behavior update context data.</param>
+    API_FUNCTION() virtual void InitState(const BehaviorUpdateContext& context)
     {
     }
 
     /// <summary>
-    /// Cleanups node instance state. Called when stopping logic simulation for a given behavior.
+    /// Cleanups node instance state. Called when stopping logic simulation for a given behavior. Call destructor of the state container.
     /// </summary>
-    /// <param name="behavior">Behavior to simulate.</param>
-    /// <param name="memory">Pointer to pre-allocated memory for this node to use (call destructor of the state container).</param>
-    API_FUNCTION() virtual void ReleaseState(Behavior* behavior, void* memory)
+    /// <param name="context">Behavior update context data.</param>
+    API_FUNCTION() virtual void ReleaseState(const BehaviorUpdateContext& context)
     {
     }
 
@@ -73,7 +71,7 @@ public:
     /// </summary>
     /// <param name="context">Behavior update context data.</param>
     /// <returns>Operation result enum.</returns>
-    API_FUNCTION() virtual BehaviorUpdateResult Update(BehaviorUpdateContext context)
+    API_FUNCTION() virtual BehaviorUpdateResult Update(const BehaviorUpdateContext& context)
     {
         return BehaviorUpdateResult::Success;
     }
@@ -111,7 +109,7 @@ API_CLASS(Abstract) class FLAXENGINE_API BehaviorTreeDecorator : public Behavior
     /// </summary>
     /// <param name="context">Behavior update context data.</param>
     /// <returns>True if can update, otherwise false to block it.</returns>
-    API_FUNCTION() virtual bool CanUpdate(BehaviorUpdateContext context)
+    API_FUNCTION() virtual bool CanUpdate(const BehaviorUpdateContext& context)
     {
         return true;
     }
@@ -121,7 +119,7 @@ API_CLASS(Abstract) class FLAXENGINE_API BehaviorTreeDecorator : public Behavior
     /// </summary>
     /// <param name="context">Behavior update context data.</param>
     /// <param name="result">The node update result. Can be modified by the decorator (eg. to force success).</param>
-    API_FUNCTION() virtual void PostUpdate(BehaviorUpdateContext context, API_PARAM(ref) BehaviorUpdateResult& result)
+    API_FUNCTION() virtual void PostUpdate(const BehaviorUpdateContext& context, API_PARAM(ref) BehaviorUpdateResult& result)
     {
     }
 };
diff --git a/Source/Engine/AI/BehaviorTreeNodes.cpp b/Source/Engine/AI/BehaviorTreeNodes.cpp
index c92331ce9..a134697fe 100644
--- a/Source/Engine/AI/BehaviorTreeNodes.cpp
+++ b/Source/Engine/AI/BehaviorTreeNodes.cpp
@@ -93,7 +93,7 @@ void BehaviorTreeNode::BecomeRelevant(const BehaviorUpdateContext& context)
     BitArray<>& relevantNodes = *(BitArray<>*)context.RelevantNodes;
     ASSERT_LOW_LAYER(relevantNodes.Get(_executionIndex) == false);
     relevantNodes.Set(_executionIndex, true);
-    InitState(context.Behavior, context.Memory);
+    InitState(context);
 }
 
 void BehaviorTreeNode::BecomeIrrelevant(const BehaviorUpdateContext& context)
@@ -102,7 +102,7 @@ void BehaviorTreeNode::BecomeIrrelevant(const BehaviorUpdateContext& context)
     BitArray<>& relevantNodes = *(BitArray<>*)context.RelevantNodes;
     ASSERT_LOW_LAYER(relevantNodes.Get(_executionIndex) == true);
     relevantNodes.Set(_executionIndex, false);
-    ReleaseState(context.Behavior, context.Memory);
+    ReleaseState(context);
 
     // Release decorators
     for (BehaviorTreeDecorator* decorator : _decorators)
@@ -136,7 +136,7 @@ void BehaviorTreeCompoundNode::Init(BehaviorTree* tree)
         child->Init(tree);
 }
 
-BehaviorUpdateResult BehaviorTreeCompoundNode::Update(BehaviorUpdateContext context)
+BehaviorUpdateResult BehaviorTreeCompoundNode::Update(const BehaviorUpdateContext& context)
 {
     auto result = BehaviorUpdateResult::Success;
     for (int32 i = 0; i < Children.Count() && result == BehaviorUpdateResult::Success; i++)
@@ -167,13 +167,13 @@ int32 BehaviorTreeSequenceNode::GetStateSize() const
     return sizeof(State);
 }
 
-void BehaviorTreeSequenceNode::InitState(Behavior* behavior, void* memory)
+void BehaviorTreeSequenceNode::InitState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
+    auto state = GetState<State>(context.Memory);
     new(state)State();
 }
 
-BehaviorUpdateResult BehaviorTreeSequenceNode::Update(BehaviorUpdateContext context)
+BehaviorUpdateResult BehaviorTreeSequenceNode::Update(const BehaviorUpdateContext& context)
 {
     auto state = GetState<State>(context.Memory);
 
@@ -204,13 +204,13 @@ int32 BehaviorTreeSelectorNode::GetStateSize() const
     return sizeof(State);
 }
 
-void BehaviorTreeSelectorNode::InitState(Behavior* behavior, void* memory)
+void BehaviorTreeSelectorNode::InitState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
+    auto state = GetState<State>(context.Memory);
     new(state)State();
 }
 
-BehaviorUpdateResult BehaviorTreeSelectorNode::Update(BehaviorUpdateContext context)
+BehaviorUpdateResult BehaviorTreeSelectorNode::Update(const BehaviorUpdateContext& context)
 {
     auto state = GetState<State>(context.Memory);
 
@@ -238,15 +238,15 @@ int32 BehaviorTreeDelayNode::GetStateSize() const
     return sizeof(State);
 }
 
-void BehaviorTreeDelayNode::InitState(Behavior* behavior, void* memory)
+void BehaviorTreeDelayNode::InitState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
-    if (!WaitTimeSelector.TryGet(behavior->GetKnowledge(), state->TimeLeft))
+    auto state = GetState<State>(context.Memory);
+    if (!WaitTimeSelector.TryGet(context.Knowledge, state->TimeLeft))
         state->TimeLeft = WaitTime;
     state->TimeLeft = Random::RandRange(Math::Max(state->TimeLeft - RandomDeviation, 0.0f), state->TimeLeft + RandomDeviation);
 }
 
-BehaviorUpdateResult BehaviorTreeDelayNode::Update(BehaviorUpdateContext context)
+BehaviorUpdateResult BehaviorTreeDelayNode::Update(const BehaviorUpdateContext& context)
 {
     auto state = GetState<State>(context.Memory);
     state->TimeLeft -= context.DeltaTime;
@@ -258,9 +258,9 @@ int32 BehaviorTreeSubTreeNode::GetStateSize() const
     return sizeof(State);
 }
 
-void BehaviorTreeSubTreeNode::InitState(Behavior* behavior, void* memory)
+void BehaviorTreeSubTreeNode::InitState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
+    auto state = GetState<State>(context.Memory);
     new(state)State();
     const BehaviorTree* tree = Tree.Get();
     if (!tree || tree->WaitForLoaded())
@@ -270,22 +270,27 @@ void BehaviorTreeSubTreeNode::InitState(Behavior* behavior, void* memory)
     state->RelevantNodes.SetAll(false);
 }
 
-void BehaviorTreeSubTreeNode::ReleaseState(Behavior* behavior, void* memory)
+void BehaviorTreeSubTreeNode::ReleaseState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
+    auto state = GetState<State>(context.Memory);
     const BehaviorTree* tree = Tree.Get();
     if (tree && tree->IsLoaded())
     {
+        // Override memory with custom one for the subtree
+        BehaviorUpdateContext subContext = context;
+        subContext.Memory = state->Memory.Get();
+        subContext.RelevantNodes = &state->RelevantNodes;
+
         for (const auto& node : tree->Graph.Nodes)
         {
             if (node.Instance && node.Instance->_executionIndex != -1 && state->RelevantNodes.HasItems() && state->RelevantNodes[node.Instance->_executionIndex])
-                node.Instance->ReleaseState(behavior, state->Memory.Get());
+                node.Instance->ReleaseState(subContext);
         }
     }
     state->~State();
 }
 
-BehaviorUpdateResult BehaviorTreeSubTreeNode::Update(BehaviorUpdateContext context)
+BehaviorUpdateResult BehaviorTreeSubTreeNode::Update(const BehaviorUpdateContext& context)
 {
     const BehaviorTree* tree = Tree.Get();
     if (!tree || !tree->Graph.Root)
@@ -306,20 +311,21 @@ BehaviorUpdateResult BehaviorTreeSubTreeNode::Update(BehaviorUpdateContext conte
 
     // Override memory with custom one for the subtree
     auto state = GetState<State>(context.Memory);
-    context.Memory = state->Memory.Get();
-    context.RelevantNodes = &state->RelevantNodes;
+    BehaviorUpdateContext subContext = context;
+    subContext.Memory = state->Memory.Get();
+    subContext.RelevantNodes = &state->RelevantNodes;
 
     // Run nested tree
-    return tree->Graph.Root->InvokeUpdate(context);
+    return tree->Graph.Root->InvokeUpdate(subContext);
 }
 
-BehaviorUpdateResult BehaviorTreeForceFinishNode::Update(BehaviorUpdateContext context)
+BehaviorUpdateResult BehaviorTreeForceFinishNode::Update(const BehaviorUpdateContext& context)
 {
     context.Behavior->StopLogic(Result);
     return Result;
 }
 
-void BehaviorTreeInvertDecorator::PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result)
+void BehaviorTreeInvertDecorator::PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result)
 {
     if (result == BehaviorUpdateResult::Success)
         result = BehaviorUpdateResult::Failed;
@@ -327,13 +333,13 @@ void BehaviorTreeInvertDecorator::PostUpdate(BehaviorUpdateContext context, Beha
         result = BehaviorUpdateResult::Success;
 }
 
-void BehaviorTreeForceSuccessDecorator::PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result)
+void BehaviorTreeForceSuccessDecorator::PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result)
 {
     if (result != BehaviorUpdateResult::Running)
         result = BehaviorUpdateResult::Success;
 }
 
-void BehaviorTreeForceFailedDecorator::PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result)
+void BehaviorTreeForceFailedDecorator::PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result)
 {
     if (result != BehaviorUpdateResult::Running)
         result = BehaviorUpdateResult::Failed;
@@ -344,14 +350,14 @@ int32 BehaviorTreeLoopDecorator::GetStateSize() const
     return sizeof(State);
 }
 
-void BehaviorTreeLoopDecorator::InitState(Behavior* behavior, void* memory)
+void BehaviorTreeLoopDecorator::InitState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
-    if (!LoopCountSelector.TryGet(behavior->GetKnowledge(), state->Loops))
+    auto state = GetState<State>(context.Memory);
+    if (!LoopCountSelector.TryGet(context.Knowledge, state->Loops))
         state->Loops = LoopCount;
 }
 
-void BehaviorTreeLoopDecorator::PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result)
+void BehaviorTreeLoopDecorator::PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result)
 {
     // Continue looping only if node succeeds
     if (result == BehaviorUpdateResult::Success)
@@ -375,15 +381,15 @@ int32 BehaviorTreeTimeLimitDecorator::GetStateSize() const
     return sizeof(State);
 }
 
-void BehaviorTreeTimeLimitDecorator::InitState(Behavior* behavior, void* memory)
+void BehaviorTreeTimeLimitDecorator::InitState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
-    if (!MaxDurationSelector.TryGet(behavior->GetKnowledge(), state->TimeLeft))
+    auto state = GetState<State>(context.Memory);
+    if (!MaxDurationSelector.TryGet(context.Knowledge, state->TimeLeft))
         state->TimeLeft = MaxDuration;
     state->TimeLeft = Random::RandRange(Math::Max(state->TimeLeft - RandomDeviation, 0.0f), state->TimeLeft + RandomDeviation);
 }
 
-BehaviorUpdateResult BehaviorTreeTimeLimitDecorator::Update(BehaviorUpdateContext context)
+BehaviorUpdateResult BehaviorTreeTimeLimitDecorator::Update(const BehaviorUpdateContext& context)
 {
     auto state = GetState<State>(context.Memory);
     state->TimeLeft -= context.DeltaTime;
@@ -395,26 +401,26 @@ int32 BehaviorTreeCooldownDecorator::GetStateSize() const
     return sizeof(State);
 }
 
-void BehaviorTreeCooldownDecorator::InitState(Behavior* behavior, void* memory)
+void BehaviorTreeCooldownDecorator::InitState(const BehaviorUpdateContext& context)
 {
-    auto state = GetState<State>(memory);
+    auto state = GetState<State>(context.Memory);
     state->EndTime = 0; // Allow to entry on start
 }
 
-void BehaviorTreeCooldownDecorator::ReleaseState(Behavior* behavior, void* memory)
+void BehaviorTreeCooldownDecorator::ReleaseState(const BehaviorUpdateContext& context)
 {
     // Preserve the decorator's state to keep cooldown
-    BitArray<>& relevantNodes = behavior->GetKnowledge()->RelevantNodes;
+    BitArray<>& relevantNodes = *(BitArray<>*)context.RelevantNodes;
     relevantNodes.Set(_executionIndex, true);
 }
 
-bool BehaviorTreeCooldownDecorator::CanUpdate(BehaviorUpdateContext context)
+bool BehaviorTreeCooldownDecorator::CanUpdate(const BehaviorUpdateContext& context)
 {
     auto state = GetState<State>(context.Memory);
     return state->EndTime <= context.Time;
 }
 
-void BehaviorTreeCooldownDecorator::PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result)
+void BehaviorTreeCooldownDecorator::PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result)
 {
     if (result != BehaviorUpdateResult::Running)
     {
diff --git a/Source/Engine/AI/BehaviorTreeNodes.h b/Source/Engine/AI/BehaviorTreeNodes.h
index 2cd2d6f25..10eca71ea 100644
--- a/Source/Engine/AI/BehaviorTreeNodes.h
+++ b/Source/Engine/AI/BehaviorTreeNodes.h
@@ -24,7 +24,7 @@ API_CLASS(Abstract) class FLAXENGINE_API BehaviorTreeCompoundNode : public Behav
 public:
     // [BehaviorTreeNode]
     void Init(BehaviorTree* tree) override;
-    BehaviorUpdateResult Update(BehaviorUpdateContext context) override;
+    BehaviorUpdateResult Update(const BehaviorUpdateContext& context) override;
 
 protected:
     // [BehaviorTreeNode]
@@ -41,8 +41,8 @@ API_CLASS() class FLAXENGINE_API BehaviorTreeSequenceNode : public BehaviorTreeC
 public:
     // [BehaviorTreeNode]
     int32 GetStateSize() const override;
-    void InitState(Behavior* behavior, void* memory) override;
-    BehaviorUpdateResult Update(BehaviorUpdateContext context) override;
+    void InitState(const BehaviorUpdateContext& context) override;
+    BehaviorUpdateResult Update(const BehaviorUpdateContext& context) override;
 
 private:
     struct State
@@ -61,8 +61,8 @@ API_CLASS() class FLAXENGINE_API BehaviorTreeSelectorNode : public BehaviorTreeC
 public:
     // [BehaviorTreeNode]
     int32 GetStateSize() const override;
-    void InitState(Behavior* behavior, void* memory) override;
-    BehaviorUpdateResult Update(BehaviorUpdateContext context) override;
+    void InitState(const BehaviorUpdateContext& context) override;
+    BehaviorUpdateResult Update(const BehaviorUpdateContext& context) override;
 
 private:
     struct State
@@ -111,8 +111,8 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeDelayNode : public BehaviorTr
 public:
     // [BehaviorTreeNode]
     int32 GetStateSize() const override;
-    void InitState(Behavior* behavior, void* memory) override;
-    BehaviorUpdateResult Update(BehaviorUpdateContext context) override;
+    void InitState(const BehaviorUpdateContext& context) override;
+    BehaviorUpdateResult Update(const BehaviorUpdateContext& context) override;
 
 private:
     struct State
@@ -136,9 +136,9 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeSubTreeNode : public Behavior
 public:
     // [BehaviorTreeNode]
     int32 GetStateSize() const override;
-    void InitState(Behavior* behavior, void* memory) override;
-    void ReleaseState(Behavior* behavior, void* memory) override;
-    BehaviorUpdateResult Update(BehaviorUpdateContext context) override;
+    void InitState(const BehaviorUpdateContext& context) override;
+    void ReleaseState(const BehaviorUpdateContext& context) override;
+    BehaviorUpdateResult Update(const BehaviorUpdateContext& context) override;
 
     struct State
     {
@@ -160,7 +160,7 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeForceFinishNode : public Beha
 
 public:
     // [BehaviorTreeNode]
-    BehaviorUpdateResult Update(BehaviorUpdateContext context) override;
+    BehaviorUpdateResult Update(const BehaviorUpdateContext& context) override;
 };
 
 /// <summary>
@@ -172,7 +172,7 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeInvertDecorator : public Beha
 
 public:
     // [BehaviorTreeNode]
-    void PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result) override;
+    void PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result) override;
 };
 
 /// <summary>
@@ -184,7 +184,7 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeForceSuccessDecorator : publi
 
 public:
     // [BehaviorTreeNode]
-    void PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result) override;
+    void PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result) override;
 };
 
 /// <summary>
@@ -196,7 +196,7 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeForceFailedDecorator : public
 
 public:
     // [BehaviorTreeNode]
-    void PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result) override;
+    void PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result) override;
 };
 
 /// <summary>
@@ -218,8 +218,8 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeLoopDecorator : public Behavi
 public:
     // [BehaviorTreeNode]
     int32 GetStateSize() const override;
-    void InitState(Behavior* behavior, void* memory) override;
-    void PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result) override;
+    void InitState(const BehaviorUpdateContext& context) override;
+    void PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result) override;
 
     struct State
     {
@@ -250,8 +250,8 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeTimeLimitDecorator : public B
 public:
     // [BehaviorTreeNode]
     int32 GetStateSize() const override;
-    void InitState(Behavior* behavior, void* memory) override;
-    BehaviorUpdateResult Update(BehaviorUpdateContext context) override;
+    void InitState(const BehaviorUpdateContext& context) override;
+    BehaviorUpdateResult Update(const BehaviorUpdateContext& context) override;
 
     struct State
     {
@@ -282,10 +282,10 @@ API_CLASS(Sealed) class FLAXENGINE_API BehaviorTreeCooldownDecorator : public Be
 public:
     // [BehaviorTreeNode]
     int32 GetStateSize() const override;
-    void InitState(Behavior* behavior, void* memory) override;
-    void ReleaseState(Behavior* behavior, void* memory) override;
-    bool CanUpdate(BehaviorUpdateContext context) override;
-    void PostUpdate(BehaviorUpdateContext context, BehaviorUpdateResult& result) override;
+    void InitState(const BehaviorUpdateContext& context) override;
+    void ReleaseState(const BehaviorUpdateContext& context) override;
+    bool CanUpdate(const BehaviorUpdateContext& context) override;
+    void PostUpdate(const BehaviorUpdateContext& context, BehaviorUpdateResult& result) override;
 
     struct State
     {