Failed load model to Learning-Agents

Development Programming & Scripting AI
, , ,
1

Hello everyone,

I’m trying to replace the internally generated Critic network in UE5’s Learning-Agents plugin with an external ONNX model exported from PyTorch. I’ve made some modifications but am running into a crash and would appreciate any guidance.

My Goal:
To pass a UNNEModelData asset (created from a PyTorch-exported ONNX file) into the ULearningAgentsCritic setup process, so it uses my custom model instead of building its own MLP.

My Approach & Modifications(Code Below):

  1. I created modified versions of MakeCritic and SetupCritic (MyMakeCritic, MySetupCritic) in ULearningAgentsCritic. I added a new parameter UNNEModelData* MyModelData to pass the external model asset.

  2. In MySetupCritic, I modified the critical section to use my model data instead of the built-in FModelBuilder

  3. In ULearningNeuralNetworkData, I created MyInit and MyUpdateNetwork functions to handle the external UNNEModelDatafrom MySetupCritic

My Code:

  • LearningAgentsCritic.cpp
ULearningAgentsCritic* ULearningAgentsCritic::MyMakeCritic(
  ULearningAgentsManager*& InManager,
  ULearningAgentsInteractor*& InInteractor,
  ULearningAgentsPolicy*& InPolicy,
  UNNEModelData* MyModelData,
  TSubclassOf Class,
  const FName Name,
  ULearningAgentsNeuralNetwork* CriticNeuralNetworkAsset,
  const bool bReinitializeCriticNetwork,
  const FLearningAgentsCriticSettings& CriticSettings,
  const int32 Seed)
{
  if (!InManager)
  {
    UE_LOG(LogLearning, Error, TEXT(“MakeCritic: InManager is       nullptr.”));
    return nullptr;
  }

  if (!Class)
  {
	UE_LOG(LogLearning, Error, TEXT("MakeCritic: Class is nullptr."));
	return nullptr;
  }

  const FName UniqueName = MakeUniqueObjectName(InManager, Class, Name, EUniqueObjectNameOptions::GloballyUnique);

  ULearningAgentsCritic* Critic = NewObject<ULearningAgentsCritic>(InManager, Class, UniqueName);
  if (!Critic) { return nullptr; }

  Critic->MySetupCritic(
	InManager,
	InInteractor,
	InPolicy,
	MyModelData,
	CriticNeuralNetworkAsset,
	bReinitializeCriticNetwork,
	CriticSettings,
	Seed);

  return Critic->IsSetup() ? Critic : nullptr;

}

void ULearningAgentsCritic::MySetupCritic(
	ULearningAgentsManager*& InManager, 
	ULearningAgentsInteractor*& InInteractor, 
	ULearningAgentsPolicy*& InPolicy, 
	UNNEModelData* MyModelData,
	ULearningAgentsNeuralNetwork* CriticNeuralNetworkAsset, 
	const bool bReinitializeCriticNetwork, 
	const FLearningAgentsCriticSettings& CriticSettings, 
	const int32 Seed)
{

    //...
    //...same as ULearningAgentsCritic::SetupCritic...//
    //...

	if (!CriticNetwork->NeuralNetworkData || bReinitializeCriticNetwork)
	{
		// Create New Neural Network Asset

		if (!CriticNetwork->NeuralNetworkData)
		{
			CriticNetwork->NeuralNetworkData = NewObject<ULearningNeuralNetworkData>(CriticNetwork);
		}

		//UE::NNE::RuntimeBasic::FModelBuilder Builder(UE::Learning::Random::Int(Seed ^ 0x2610fc8f));

		//TArray<uint8> FileData;
		//uint32 CriticInputSize, CriticOutputSize;
		//Builder.WriteFileDataAndReset(FileData, CriticInputSize, CriticOutputSize,
		//	Builder.MakeMLP(
		//		ObservationEncodedVectorSize + MemoryStateSize,
		//		1,
		//		CriticSettings.HiddenLayerSize,
		//		CriticSettings.HiddenLayerNum + 2, // Add 2 to account for input and output layers
		//		UE::Learning::Agents::Critic::Private::GetBuilderActivationFunction(CriticSettings.ActivationFunction)));

		//check(CriticInputSize == ObservationEncodedVectorSize + MemoryStateSize);
		//check(CriticOutputSize == 1);

		//CriticNetwork->NeuralNetworkData->Init(CriticInputSize, CriticOutputSize, CriticCompatibilityHash, FileData);
		CriticNetwork->NeuralNetworkData->MyInit(ObservationEncodedVectorSize + MemoryStateSize, 1, CriticCompatibilityHash, MyModelData);  //Modified
		CriticNetwork->ForceMarkDirty();
	}

	// Create Critic Object
	/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////  <---- ERROR
	CriticObject = MakeShared<UE::Learning::FNeuralNetworkCritic>(
		Manager->GetMaxAgentNum(),
		ObservationEncodedVectorSize,
		MemoryStateSize,
		CriticNetwork->NeuralNetworkData->GetNetwork());

	Returns.SetNumUninitialized({ Manager->GetMaxAgentNum() });
	UE::Learning::Array::Set<1, float>(Returns, FLT_MAX);

	ReturnsIteration.SetNumUninitialized({ Manager->GetMaxAgentNum() });
	UE::Learning::Array::Set<1, uint64>(ReturnsIteration, INDEX_NONE);

	bIsSetup = true;

	Manager->AddListener(this);
}
  • LearningNeuralNetworkData.cpp
void ULearningNeuralNetworkData::MyInit(const int32 InInputSize, const int32 InOutputSize, const int32 InCompatibilityHash, UNNEModelData* InMyModelData)
{
	InputSize = InInputSize;
	OutputSize = InOutputSize;
	CompatibilityHash = InCompatibilityHash;
	ModelData = *InMyModelData;
	MyUpdateNetwork();
}


void ULearningNeuralNetworkData::MyUpdateNetwork()
{
	if (FileData.Num() > 0)
	{
		if (ensureMsgf(ModelData, TEXT("Could not find requested ModelData"))) {
		}

		if (!Network)
		{
			Network = MakeShared<UE::Learning::FNeuralNetwork>();
		}

		  ensureMsgf(FModuleManager::Get().LoadModule(TEXT("NNERuntimeBasicCpu")), TEXT("Unable to load module for NNE runtime NNERuntimeBasicCpu."));

		TWeakInterfacePtr<INNERuntimeCPU> RuntimeCPU = UE::NNE::GetRuntime<INNERuntimeCPU>(TEXT("NNERuntimeBasicCpu"));

		TSharedPtr<UE::NNE::IModelCPU> UpdatedModel = nullptr;

		if (ensureMsgf(RuntimeCPU.IsValid(), TEXT("Could not find requested NNE Runtime")))
		{
			UpdatedModel = RuntimeCPU->CreateModelCPU(ModelData);
		}

		// Compute the content hash

		//ContentHash = CityHash32((const char*)(ModelData->GetFileData()).GetData(), (ModelData->GetFileData()).Num());

		// If we are not in the editor we can now clear the FileData and FileType since these will be
		// using additional memory and we are not going to save this asset and so don't require them.

#if !WITH_EDITOR
		ModelData->ClearFileDataAndFileType();
#endif

		Network->UpdateModel(UpdatedModel, InputSize, OutputSize);
	}
}
  • onnx_generator.py
import warnings
import torch
import torch.nn as nn

warnings.filterwarnings("ignore", message=".*legacy TorchScript-based ONNX export.*")


class NeuralNetwork(nn.Module):
    def __init__(self, input_dim=66, hidden_dim=128, output_dim=1):
        super(NeuralNetwork, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.model(x)

model = NeuralNetwork(input_dim=66, hidden_dim=128, output_dim=1)
model.eval()

dummy_input = torch.randn(1, 66) #inputsize=2  memsize=64

torch.onnx.export(
    model,
    dummy_input,
    "dynamic_batch_model.onnx",
    export_params=True,
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

The Problem / Crash:

The code crashes when MySetupCritic tries to create the FNeuralNetworkCritic

CriticObject = MakeShared<UE::Learning::FNeuralNetworkCritic>(
    Manager->GetMaxAgentNum(),
    ObservationEncodedVectorSize,
    MemoryStateSize,
    CriticNetwork->NeuralNetworkData->GetNetwork() // **CRASH HERE**
);

The crash occurs inside the TSharedPtr’s operator*, with a check(IsValid()) failure. This indicates that the TSharedPtr<UE::Learning::FNeuralNetwork> returned by GetNetwork() is null or invalid.

/**
 * Dereference operator returns a reference to the object this shared pointer points to
 *
 * @return  Reference to the object
 */
template <
	typename DummyObjectType = ObjectType
	UE_REQUIRES(UE_REQUIRES_EXPR(*(DummyObjectType*)nullptr)) // this construct means that operator* is only considered for overload resolution if T is dereferenceable
>
[[nodiscard]] FORCEINLINE DummyObjectType& operator*() const
{
	check( IsValid() );
	return *Object;
}

What I’ve Checked/Tried:

  • The ONNX model has an identical architecture to the original critic network.
the results from the original, unmodified plugin:
================Critic Network ID: 1================================
NeuralNetwork(
  (model): Sequential(
    (0): Linear()
    (1): ELU()
    (2): Linear()
    (3): ELU()
    (4): Linear()
  )
)
  • I’ve verified that the InputSize and OutputSize I’m passing match the expectations of the original code and my ONNX model.

Environment:

  • Unreal Engine 5.6 Release

I’m a student researcher focusing on Multi-Agent Reinforcement Learning (MARL). I’m using UE5 to create custom environments for my research, and the Learning-Agents plugin is central to this work.

I’m looking for a robust method to integrate externally trained ONNX models into UE5’s Learning-Agents framework. While I’d appreciate help with my current implementation issues, I’m equally interested in any alternative approaches that achieve this integration successfully.

Thanks for any insights!