UE 5.8 – NNERuntimeBasicCpu: Conv1d/Conv2d crash when the convolved observation is a stateless layer

Development Programming & Scripting
, , , ,
1

Summary

FConv1dInstance::SetMaxBatchSize and FConv2dInstance::SetMaxBatchSize call SubLayerInstance->SetMaxBatchSize(MaxBatchSize) unconditionally.

ILayer::MakeInstance() returns nullptr by default, so any stateless sub-layer has no instance and the call dereferences a null TSharedPtr. The editor asserts in SharedPointer.h (IsValid(), line 1133) as soon as the policy is created.

  • Engine version: 5.8
  • File: Engine/Plugins/Experimental/NNERuntimeBasicCpu/Source/NNERuntimeBasicCpu/Private/NNERuntimeBasicCpuModel.cpp
  • Affects: Learning Agents, SpecifyConv1dObservation / SpecifyConv2dObservation

Steps to reproduce

  1. In a Learning Agents interactor, wrap a plain continuous observation in a 2D convolution:
ULearningAgentsObservations::SpecifyConv2dObservation(
    Schema,
    ULearningAgentsObservations::SpecifyContinuousObservation(Schema, 12 * 24),
    FConv2dObservationParams(12, 24, 1, 8, 3, 2, 1));
  1. Start training, or anything else that calls ULearningAgentsPolicy::SetupPolicy.
  2. The editor crashes.

Callstack

Assertion failed: IsValid() [SharedPointer.h] [Line: 1133]

UE::NNE::RuntimeBasic::Private::FConv2dInstance::SetMaxBatchSize()
UE::NNE::RuntimeBasic::Private::FSequenceLayerInstance::SetMaxBatchSize()
UE::NNE::RuntimeBasic::Private::FConcatLayerInstance::SetMaxBatchSize()
UE::NNE::RuntimeBasic::FModelInstanceCPU::SetInputTensorShapes()
UE::Learning::FNeuralNetworkInference::ReloadModelInstances()
UE::Learning::FNeuralNetwork::CreateInferenceObject()
ULearningAgentsPolicy::SetupPolicy()

Cause

FConv2dInstance’s constructor stores Conv2dLayer.SubLayer->MakeInstance(), which is nullptr for a stateless layer:

FConv2dInstance::FConv2dInstance(const FConv2dLayer& InConv2dLayer) : Conv2dLayer(InConv2dLayer)
{
    SubLayerInstance = Conv2dLayer.SubLayer->MakeInstance();
}

SetMaxBatchSize then uses it without a check:

void FConv2dInstance::SetMaxBatchSize(const uint32 MaxBatchSize)
{
    SubLayerInstance->SetMaxBatchSize(MaxBatchSize); // null here
    SubLayerOutputBuffer.SetNumUninitialized(MaxBatchSize * Conv2dLayer.SubLayer->GetOutputSize(), EAllowShrinking::No);
}

Other wrapper layers in the same file already guard this case, for example the aggregate set layer:

if (SubLayerInstance) { SubLayerInstance->SetMaxBatchSize(MaxBatchSize * AggregateSetLayer.MaxElementNum); }

FConv1dInstance has the same unguarded call.

Suggested fix

Guard the call in both FConv1dInstance::SetMaxBatchSize and FConv2dInstance::SetMaxBatchSize:

if (SubLayerInstance) { SubLayerInstance->SetMaxBatchSize(MaxBatchSize); }

Verification

With this guard applied locally, the Conv2d observation above trains and runs for several hours without the crash.

Conv1d was patched the same way because the code is identical. We have not exercised it.