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):
-
I created modified versions of
MakeCriticandSetupCritic(MyMakeCritic,MySetupCritic) inULearningAgentsCritic. I added a new parameterUNNEModelData* MyModelDatato pass the external model asset. -
In
MySetupCritic, I modified the critical section to use my model data instead of the built-inFModelBuilder -
In
ULearningNeuralNetworkData, I createdMyInitandMyUpdateNetworkfunctions to handle the externalUNNEModelDatafromMySetupCritic
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
InputSizeandOutputSizeI’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!
