Adding support for State Tree Linked Asset Overrides when the Overriding State Tree matches the runtime Owner context

My company is currently using UE 5.6.1, but a quick glance at the 5.7.3 branch makes it seem like the behavior is the same there too, so I thought I’d make a post.

Right now, Linked Asset Overrides in State Trees (particularly with the StateTreeAIComponentSchema) is a bit limited in that the configured ContextActorClass and AIControllerClass must match. Consider the following scenario, where I have the following AI Controller classes:

AAIController_Enemy
- AAIController_Human (inherits from AAIController_Enemy)
  + AAIController_HumanGuard (inherits from AAIController_Human)
- AAIController_Alien (inherits from AAIController_Enemy)

Let’s say that I want all of my Enemies to follow a base high-level logic, such as:

ST_Main:
- Root
  + Non-Combat (Linked Asset)
  + Combat (Linked Asset)
  + Something Else (Linked Asset)

In this scenario, ST_Main would have the AIControllerClass set to “AAIController_Enemy”, because this is the routine I want all of my Enemies to follow.

Now, let’s say I have the following trees:

- “ST_NonCombat_Enemy” - generic, can be used by any Enemy

- “ST_NonCombat_Alien” - specific, can only be used by AAIController_Alien

During runtime, for a AAIController_Alien running ST_Main, if I try to call AddLinkedAssetOverride for Non-Combat override to ST_NonCombat_Alien, it will fail - even though the Owner running the StateTree is a AAIController_Alien class. It makes sense to fail if the Owner was an AAIController_Human, but feels limiting for a AAIController_Alien.

I will post a potential solution to this in the replies.

I’ve come up with a potential solution for this issue, and it seems to be working for my company (at least in our use-cases), so I thought I’d share it here. That said, State Trees are a pretty massive system, so it’s possible I’ve missed some cases or am going against some assumptions the State Tree engineers made for this system.

In FStateTreeExecutionContext, I added the function:

bool FStateTreeExecutionContext::IsContextDataCompatible(const UStateTree* OtherStateTree) const
{
	if (OtherStateTree == nullptr || OtherStateTree->GetSchema() == nullptr || RootStateTree.GetSchema() == nullptr)
	{
		return false;
	}
	
	// Verify that the OtherStateTree context actor/controller class are children of RootStateTree's contexts
	if (!RootStateTree.HasCompatibleContextData(*OtherStateTree))
	{
		return false;
	}
	
	const TConstArrayView<FStateTreeExternalDataDesc>& RootContextDescs = RootStateTree.ContextDataDescs;
	const TConstArrayView<FStateTreeExternalDataDesc>& ItemContextDescs = OtherStateTree->ContextDataDescs;
	
	if (RootContextDescs.Num() != ItemContextDescs.Num())
	{
		return false;
	}
	
	const int32 Num = RootContextDescs.Num();
	for (int32 Index = 0; Index < Num; Index++)
	{
		const FStateTreeExternalDataDesc& OtherDataDesc = ItemContextDescs[Index];
		const UStruct* OwnerData = ContextAndExternalDataViews[Index].GetStruct();
		
		// Verify that the Owner is compatible with the context data for the State Tree we're trying to override
		if (!OwnerData || !OwnerData ->IsChildOf(OtherDataDesc.Struct))
		{
			return false;
		}
	}
	
	return true;
}

Then, in FStateTreeExecutionContext::SetLinkedStateTreeOverrides, we want to make a modification to the code:

void FStateTreeExecutionContext::SetLinkedStateTreeOverrides(FStateTreeReferenceOverrides InLinkedStateTreeOverrides)
{
	// Leave everything the same until we get to the * * * comments
 
	bool bValid = true;
 
	// Confirms that the overrides schema matches.
	const TConstArrayView<FStateTreeReferenceOverrideItem> InOverrideItems = InLinkedStateTreeOverrides.GetOverrideItems();
	for (const FStateTreeReferenceOverrideItem& Item : InOverrideItems)
	{
		if (const UStateTree* ItemStateTree = Item.GetStateTreeReference().GetStateTree())
		{
			if (!ItemStateTree->IsReadyToRun())
			{
				STATETREE_LOG(Error, TEXT("%hs: '%s' using StateTree '%s' trying to set override '%s' but the tree is not initialized properly."),
					__FUNCTION__, *GetNameSafe(&Owner), *GetFullNameSafe(GetStateTree()), *GetFullNameSafe(ItemStateTree));
				bValid = false;
				break;
			}
 
			// * * * This is where the modification starts; the rest is there for reference * * * //
			if (!ItemStateTree->HasCompatibleContextData(RootStateTree)) // This was originally backwards, which I made another post about
			{
				// Allow compatibility if OverrideTree.ContextActor is a child of RootStateTree.ContextActor
				if (!IsContextDataCompatible(ItemStateTree))
				{
					STATETREE_LOG(Error, TEXT("%hs: '%s' using StateTree '%s' trying to set override '%s' but the tree context data is not compatible."),
						   __FUNCTION__, *GetNameSafe(&Owner), *GetFullNameSafe(GetStateTree()), *GetFullNameSafe(ItemStateTree));
					bValid = false;
					break;
				}
			}
			// * * * This is where the modification end; the rest is there for reference * * * //
 
		// Leave the rest of the code the same
		}
	}
}

And then we want to do the same in FStateTreeExecutionContext::SelectStateInternal:

bool FStateTreeExecutionContext::SelectStateInternal(
	const FStateTreeExecutionFrame* CurrentParentFrame,
	FStateTreeExecutionFrame& CurrentFrame,
	const FStateTreeExecutionFrame* CurrentFrameInActiveFrames,
	TConstArrayView<FStateTreeStateHandle> PathToNextState,
	FStateSelectionResult& OutSelectionResult,
	const FStateTreeSharedEvent* TransitionEvent)
{
	// Leave everything the same until we get to the * * * comments
	
	else if (NextState.Type == EStateTreeStateType::LinkedAsset)
	{
		if (NextLinkedStateAsset == nullptr || NextLinkedStateAsset->States.Num() == 0)
		{
			break;
		}
 
		if (OutSelectionResult.IsFull())
		{
			STATETREE_LOG(Error, TEXT("%hs: Reached max execution depth when trying to select state %s from '%s'.  '%s' using StateTree '%s'."),
					__FUNCTION__, *GetSafeStateName(CurrentFrame, NextStateHandle), *GetStateStatusString(Exec), *GetNameSafe(&Owner), *GetFullNameSafe(CurrentFrame.StateTree));
			break;
		}
 
		// The linked state tree should have compatible context requirements.
		if (!NextLinkedStateAsset->HasCompatibleContextData(RootStateTree)
				|| NextLinkedStateAsset->GetSchema()->GetClass() != RootStateTree.GetSchema()->GetClass())
		{
			// * * * This is where the modification starts; the rest is there for reference * * * //
			if (!IsContextDataCompatible(NextLinkedStateAsset))
			{
				STATETREE_LOG(Error, TEXT("%hs: The linked State Tree '%s' does not have compatible schema, trying to select state %s from '%s'.  '%s' using StateTree '%s'."),
					__FUNCTION__, *GetFullNameSafe(NextLinkedStateAsset), *GetSafeStateName(CurrentFrame, NextStateHandle), *GetStateStatusString(Exec), *GetNameSafe(&Owner), *GetFullNameSafe(CurrentFrame.StateTree));
				break;
			}
			// * * * This is where the modification ends; the rest is there for reference * * * //
		}
 
	// Leave everything else the same
	}
}

There is one case I’ve noticed where this doesn’t work: if the State Tree is not running (i.e. bStartAutomatically is false on the StateTreeComponent and you call AddLinkedStateOverrides before starting the tree), then the Owner’s context data gets reset. This appears to be due to the order-of-operations in UStateTreeComponent::SetContextRequirements. To fix this, I just swapped the order a bit:

bool UStateTreeComponent::SetContextRequirements(FStateTreeExecutionContext& Context, bool bLogErrors)
{
	Context.SetCollectExternalDataCallback(FOnCollectStateTreeExternalData::CreateUObject(this, &UStateTreeComponent::CollectExternalData));
	const bool bSuccess = UStateTreeComponentSchema::SetContextRequirements(*this, Context);
	Context.SetLinkedStateTreeOverrides(LinkedStateTreeOverrides);
	return bSuccess;
}

As stated before, this seems to work well for our cases in initial testing, but there are likely a lot of cases that are untested. Just thought I’d share in case others found it helpful as well.

Unfortunately, we do not currently support running linked assets that do not have the exact same schema setup. Our workaround has been fairly generic context data setups similar to how BT/BB was setup for using the Run Behavior Dynamic nodes.

While looking into this, and part of the reason for the delayed response, I found a couple bugs in the system.

  1. I found that it will execute a linked asset set in the StateTree asset directly when the linked asset uses different classes for its context properties if the class is a parent class of the root tree’s relevant context property. This works, but speaking with the feature owner, it is not intended to be allowed.
  2. The fact that you can call SetLinkedStateTreeOverrides and have it successfully set the linked asset to a StateTree which has its context property be a class that is a child of the root tree’s context. This should certainly be a fail when attempting to set and not failing from inside of SelectStateInternal.

It is certainly not ideal in its current form, and we fully agree that it could be better. We do have items in our JIRA to investigate making this run better while maintaining type safety requirements of StateTree. I will link this EPS thread to the JIRAs for investigating this added support as we may look back at some of the code changes you have made and posted here.

-James

#1 is a misunderstanding from when I discussed it with the team originally. I believe we should be able to support this. Entered the bug report and added a bit more info into your other thread. I believe using that change that this should work as you expect.

#2 works even without your changes, but it is set inside the StateTree editor as the linked asset rather than dynamically set with SetLinkedStateTreeOverrides. It does appear that the logic in the compatible context data function is backwards, and for that I say great catch!

-James

Hey James,

I apologize for the delay in my response. I was heads-down in work for a good bit and didn’t have the chance to respond until now.

For #1, could you please elaborate why it is not intended to be allowed? I may be biased since we have specific use-cases where we want this type of behavior in our project, but I do not understand why the system would want to prevent configurations where it should be safe to assume the owner can run both.

For #2, I was probably able to set it because of my proposed solution provided in [Potential bug when Linked State Tree Overrides are checking for ContextData compatibility with the [Content removed] , where I also swapped the direction of the compatibility check. Unless you were saying that even with that, it should fail earlier.

Thank you for the update!

Cool, happy to help. Thanks for the updates, James!