Extending the Learning Agents plugin

_YAF_Lightbringer · August 3, 2025, 7:51pm

Hi,
we’re currently looking into using LA to train FPS style bots that use raycast information, e.g. a vision grid. Typically, I would feed that into a 2D convolutional layer to extract information, but as far as a I can see, the LA plugin doesn’t let me do that (please advise if I missed something). Hence my question, which isn’t really limited to convolutional layers: which parts do I need to touch to expose additional pytorch functionality?

nne_runtime_basic_cpu_pytorch.py seems to handle the conversion between what the engine sent/receives and the pytorch API
LearningObservation.cpp builds the data that is send to the trainer process
LearningAgentsObservations.cpp exposes that to the user code

Did I miss anything?

(I can of course edit the plugin code, but it would be great if the plugin had some extension points, so I could make my implementation available to people who do not fancy recompiling the engine )

Deathcalibur · August 4, 2025, 1:42pm

1d and 2d conv are available on UE-main (from github). And will be coming to UE 5.7

PedroGriffith · August 4, 2025, 1:55pm

How can we effectively extend the Learning Agents plugin to support more complex tasks?

_YAF_Lightbringer · August 4, 2025, 2:53pm

Sweat, thanks! Will check it out! I had it actually half-way implemented it myself, but I can abandon that now.

General question remains though. LA plugin and expecially the BasicCPU runtime module feel very closed off and everything is hidden away behind interfaces or private utility functions. Though I can swap out the model used in the python trainer with my own script, and produce an ONNX artefact for NNE, but that defeats a bit the purpose of the LA plugin to architectect models from C++.

If I could wish for something, I’d like to have a “custom” layer, where I can manually specify what’s serialized back and forth, which python module to use for training, and provide an evaluation function for the runtime.

Deathcalibur · August 4, 2025, 5:47pm

What would you like to do that you can’t? LA’s lower level public API is pretty flexible and you can mix in extra stuff fairly easily, e.g. calling GatherObservations etc. instead of RunTraining.

@_YAF_Lightbringer
Something like that sounds good. We would need to work out the details but it might be feasible.

_YAF_Lightbringer · August 4, 2025, 8:05pm

I just tried the Conv2d layer on ue-main (b1a5dcf49535bbdc388b3f1dd81bd72fe7c37d74), seems to work. Thanks.

There’s a compilation error though that needs to be fixed. In LearningTrainer.cpp:721 ff. the function UE::Learning::Trainer::IsObservationSchemaSubsetCompatible is missing a case for Conv2. It’s simple enough to add (I can open a PR if that’s easier for you):

		case Observation::EType::Conv2d:
		{
			const Observation::FSchemaConv2dParameters Parameters = ObservationSchema.GetConv2d(ObservationSchemaElement);

			if (Object->GetStringField(TEXT("Type")) != TEXT("Conv2d"))
			{
				return false;
			}

			return IsObservationSchemaSubsetCompatible(Object->GetObjectField(TEXT("Element")).ToSharedRef(), ObservationSchema, Parameters.Element);
		}

There was also the fluke that my code didn’t run anymore until I set CustomTrainerModulePath in the trainer process settings to an absolute path, despite me not having moved any files around. But that may also be due to the fact that I’m configuring everything via code and not assets. So, I might have missed a new UPROPERTY initalizer of something.

Deathcalibur · August 4, 2025, 8:48pm

Good catch. Weirdly enough it compiles fine for us but we’ll have a fix up shortly

Jeebius · September 29, 2025, 6:57pm

Hey @Deathcalibur I was wondering if you could give me a quick tip on to get the 2d convolution observation working in 5.7, what do we input for the make observation? Appreciate if you’re too busy to answer right, but excited to try some new things.

Deathcalibur · October 7, 2025, 3:01pm

For the make Conv2d, you pass in another observation which contains the data you want to convolve. For example, you can use MakeStaticArrayObservation() with a bunch of Raycasts inside. The shape is determined by the FConv2dObservationParams passed during the Specify.

Hope that helps.

_YAF_Lightbringer · November 26, 2025, 11:07am

Hey @Deathcalibur,

I’m using the EOS version of UE 5.7 and created a simple setup: Float Observation → Static Array Observation → Conv2D Observation.

Filling the recording works fine, the schema looks fine, but (imitation) training immediately runs into the “not implemented” exception at the end of schema_noise_mask_observation in train_common.py (obs_type = obs_schema[‘Type’] evaluates to “Conv2d”).

I also don’t see the word “Conv” mentioned anywhere in that file.

Was the wrong file published or something?

_YAF_Lightbringer · December 2, 2025, 12:43pm

Below is the relevant log output

LogLearning: Display: Subprocess: INFO: {
LogLearning: Display: Subprocess:     "SharedMemory": {
LogLearning: Display: Subprocess:         "NetworkGuids": [
LogLearning: Display: Subprocess:             {
LogLearning: Display: Subprocess:                 "NetworkId": 0,
LogLearning: Display: Subprocess:                 "Guid": "{24586EFF-46CC-CD4A-33D8-33AD8D872117}"
LogLearning: Display: Subprocess:             },
LogLearning: Display: Subprocess:             {
LogLearning: Display: Subprocess:                 "NetworkId": 1,
LogLearning: Display: Subprocess:                 "Guid": "{481C10A2-4656-A920-7C5A-3F806D30E9F8}"
LogLearning: Display: Subprocess:             },
LogLearning: Display: Subprocess:             {
LogLearning: Display: Subprocess:                 "NetworkId": 2,
LogLearning: Display: Subprocess:                 "Guid": "{DF1829C9-414E-0476-FF27-C3A9D4FD493D}"
LogLearning: Display: Subprocess:             }
LogLearning: Display: Subprocess:         ],
LogLearning: Display: Subprocess:         "ReplayBuffers": [
LogLearning: Display: Subprocess:             {
LogLearning: Display: Subprocess:                 "EpisodeStartsGuid": "{34F8E3AA-4324-48D7-6E58-B591FD04AB44}",
LogLearning: Display: Subprocess:                 "EpisodeLengthsGuid": "{4CA50097-4941-8426-8898-589F22EE9011}",
LogLearning: Display: Subprocess:                 "EpisodeCompletionModesGuid": "{00000000-0000-0000-0000-000000000000}",
LogLearning: Display: Subprocess:                 "EpisodeFinalObservationsGuids": [],
LogLearning: Display: Subprocess:                 "EpisodeFinalMemoryStatesGuids": [],
LogLearning: Display: Subprocess:                 "ObservationsGuids": [
LogLearning: Display: Subprocess:                     "{44367CC1-405C-0CA1-4182-9CBF664F1858}"
LogLearning: Display: Subprocess:                 ],
LogLearning: Display: Subprocess:                 "ActionsGuids": [
LogLearning: Display: Subprocess:                     "{B1887B17-4CB8-CCA8-0E41-2BBE9B517CFB}"
LogLearning: Display: Subprocess:                 ],
LogLearning: Display: Subprocess:                 "ActionModifiersGuids": [],
LogLearning: Display: Subprocess:                 "MemoryStatesGuids": [],
LogLearning: Display: Subprocess:                 "RewardsGuids": []
LogLearning: Display: Subprocess:             }
LogLearning: Display: Subprocess:         ]
LogLearning: Display: Subprocess:     },
LogLearning: Display: Subprocess:     "Networks": [
LogLearning: Display: Subprocess:         {
LogLearning: Display: Subprocess:             "Id": 0,
LogLearning: Display: Subprocess:             "Name": "DA_Policy",
LogLearning: Display: Subprocess:             "MaxByteNum": 470584
LogLearning: Display: Subprocess:         },
LogLearning: Display: Subprocess:         {
LogLearning: Display: Subprocess:             "Id": 1,
LogLearning: Display: Subprocess:             "Name": "DA_Encoder",
LogLearning: Display: Subprocess:             "MaxByteNum": 548,
LogLearning: Display: Subprocess:             "InputSchemaId": 0
LogLearning: Display: Subprocess:         },
LogLearning: Display: Subprocess:         {
LogLearning: Display: Subprocess:             "Id": 2,
LogLearning: Display: Subprocess:             "Name": "DA_Decoder",
LogLearning: Display: Subprocess:             "MaxByteNum": 176,
LogLearning: Display: Subprocess:             "OutputSchemaId": 0
LogLearning: Display: Subprocess:         }
LogLearning: Display: Subprocess:     ],
LogLearning: Display: Subprocess:     "ReplayBuffers": [
LogLearning: Display: Subprocess:         {
LogLearning: Display: Subprocess:             "Id": 0,
LogLearning: Display: Subprocess:             "MaxEpisodeNum": 1,
LogLearning: Display: Subprocess:             "MaxStepNum": 65,
LogLearning: Display: Subprocess:             "HasCompletions": false,
LogLearning: Display: Subprocess:             "HasFinalObservations": false,
LogLearning: Display: Subprocess:             "HasFinalMemoryStates": false,
LogLearning: Display: Subprocess:             "Observations": [
LogLearning: Display: Subprocess:                 {
LogLearning: Display: Subprocess:                     "Id": 0,
LogLearning: Display: Subprocess:                     "Name": "Observations",
LogLearning: Display: Subprocess:                     "SchemaId": 0,
LogLearning: Display: Subprocess:                     "VectorDimensionNum": 8
LogLearning: Display: Subprocess:                 }
LogLearning: Display: Subprocess:             ],
LogLearning: Display: Subprocess:             "Actions": [
LogLearning: Display: Subprocess:                 {
LogLearning: Display: Subprocess:                     "Id": 0,
LogLearning: Display: Subprocess:                     "Name": "Actions",
LogLearning: Display: Subprocess:                     "SchemaId": 0,
LogLearning: Display: Subprocess:                     "VectorDimensionNum": 3
LogLearning: Display: Subprocess:                 }
LogLearning: Display: Subprocess:             ],
LogLearning: Display: Subprocess:             "ActionModifiers": [],
LogLearning: Display: Subprocess:             "MemoryStates": [],
LogLearning: Display: Subprocess:             "Rewards": []
LogLearning: Display: Subprocess:         }
LogLearning: Display: Subprocess:     ],
LogLearning: Display: Subprocess:     "Schemas": {
LogLearning: Display: Subprocess:         "Observations": [
LogLearning: Display: Subprocess:             {
LogLearning: Display: Subprocess:                 "Id": 0,
LogLearning: Display: Subprocess:                 "Name": "Default",
LogLearning: Display: Subprocess:                 "Schema": {
LogLearning: Display: Subprocess:                     "VectorSize": 8,
LogLearning: Display: Subprocess:                     "EncodedSize": 8,
LogLearning: Display: Subprocess:                     "Type": "And",
LogLearning: Display: Subprocess:                     "Elements": {
LogLearning: Display: Subprocess:                         "Direction": {
LogLearning: Display: Subprocess:                             "VectorSize": 3,
LogLearning: Display: Subprocess:                             "EncodedSize": 3,
LogLearning: Display: Subprocess:                             "Type": "Continuous",
LogLearning: Display: Subprocess:                             "Num": 3,
LogLearning: Display: Subprocess:                             "Index": 0
LogLearning: Display: Subprocess:                         },
LogLearning: Display: Subprocess:                         "Raycast": {
LogLearning: Display: Subprocess:                             "VectorSize": 5,
LogLearning: Display: Subprocess:                             "EncodedSize": 5,
LogLearning: Display: Subprocess:                             "Type": "Conv2d",
LogLearning: Display: Subprocess:                             "InputHeight": 1,
LogLearning: Display: Subprocess:                             "InputWidth": 5,
LogLearning: Display: Subprocess:                             "InChannels": 1,
LogLearning: Display: Subprocess:                             "OutChannels": 1,
LogLearning: Display: Subprocess:                             "KernelSize": 3,
LogLearning: Display: Subprocess:                             "Stride": 1,
LogLearning: Display: Subprocess:                             "Padding": 1,
LogLearning: Display: Subprocess:                             "PaddingMode": 0,
LogLearning: Display: Subprocess:                             "Element": {
LogLearning: Display: Subprocess:                                 "VectorSize": 5,
LogLearning: Display: Subprocess:                                 "EncodedSize": 5,
LogLearning: Display: Subprocess:                                 "Type": "Array",
LogLearning: Display: Subprocess:                                 "Num": 5,
LogLearning: Display: Subprocess:                                 "Element": {
LogLearning: Display: Subprocess:                                     "VectorSize": 1,
LogLearning: Display: Subprocess:                                     "EncodedSize": 1,
LogLearning: Display: Subprocess:                                     "Type": "Continuous",
LogLearning: Display: Subprocess:                                     "Num": 1
LogLearning: Display: Subprocess:                                 }
LogLearning: Display: Subprocess:                             },
LogLearning: Display: Subprocess:                             "Index": 1
LogLearning: Display: Subprocess:                         }
LogLearning: Display: Subprocess:                     }
LogLearning: Display: Subprocess:                 }
LogLearning: Display: Subprocess:             }
LogLearning: Display: Subprocess:         ],
LogLearning: Display: Subprocess:         "Actions": [
LogLearning: Display: Subprocess:             {
LogLearning: Display: Subprocess:                 "Id": 0,
LogLearning: Display: Subprocess:                 "Name": "Default",
LogLearning: Display: Subprocess:                 "Schema": {
LogLearning: Display: Subprocess:                     "VectorSize": 3,
LogLearning: Display: Subprocess:                     "DistributionSize": 6,
LogLearning: Display: Subprocess:                     "EncodedSize": 6,
LogLearning: Display: Subprocess:                     "ModifierSize": 7,
LogLearning: Display: Subprocess:                     "Type": "Continuous",
LogLearning: Display: Subprocess:                     "Num": 3
LogLearning: Display: Subprocess:                 }
LogLearning: Display: Subprocess:             }
LogLearning: Display: Subprocess:         ]
LogLearning: Display: Subprocess:     },
LogLearning: Display: Subprocess:     "TrainerMethod": "BehaviorCloning",
LogLearning: Display: Subprocess:     "TimeStamp": "2025-12-02_13-36-47",
LogLearning: Display: Subprocess:     "BehaviorCloningSettings": {
LogLearning: Display: Subprocess:         "IterationNum": 1000000,
LogLearning: Display: Subprocess:         "LearningRate": 0.0010000000474974513,
LogLearning: Display: Subprocess:         "LearningRateDecay": 1,
LogLearning: Display: Subprocess:         "LearningRateDecayStepSize": 1000,
LogLearning: Display: Subprocess:         "WeightDecay": 9.999999747378752e-05,
LogLearning: Display: Subprocess:         "BatchSize": 128,
LogLearning: Display: Subprocess:         "Window": 64,
LogLearning: Display: Subprocess:         "ActionRegularizationWeight": 0.0010000000474974513,
LogLearning: Display: Subprocess:         "ActionEntropyWeight": 0,
LogLearning: Display: Subprocess:         "ObservationNoiseScale": 0,
LogLearning: Display: Subprocess:         "Seed": 1234,
LogLearning: Display: Subprocess:         "SaveSnapshots": false,
LogLearning: Display: Subprocess:         "IterationsPerSnapshot": 1000,
LogLearning: Display: Subprocess:         "RunEvaluation": true,
LogLearning: Display: Subprocess:         "TrainEvalDatasetSplit": 0.10000000149011612,
LogLearning: Display: Subprocess:         "EvaluationFrequency": 50,
LogLearning: Display: Subprocess:         "BatchCountPerEvaluation": 10
LogLearning: Display: Subprocess:     },
LogLearning: Display: Subprocess:     "PolicySettings": {
LogLearning: Display: Subprocess:         "HiddenLayerNum": 1,
LogLearning: Display: Subprocess:         "HiddenLayerSize": 128,
LogLearning: Display: Subprocess:         "MemoryCell": "LearningAgentsGRU",
LogLearning: Display: Subprocess:         "MemoryStateSize": 64,
LogLearning: Display: Subprocess:         "InitialEncodedActionScale": 0.10000000149011612,
LogLearning: Display: Subprocess:         "ActivationFunction": "ELU"
LogLearning: Display: Subprocess:     },
LogLearning: Display: Subprocess:     "MemoryStateNum": 64,
LogLearning: Display: Subprocess:     "Device": "CPU",
LogLearning: Display: Subprocess:     "UseTensorBoard": false,
LogLearning: Display: Subprocess:     "UseMLflow": false,
LogLearning: Display: Subprocess:     "MLflowTrackingUri": "",
LogLearning: Display: Subprocess:     "TempDirectory": "C:/MyWorkspace/MyProject/Intermediate/LearningAgents",
LogLearning: Display: Subprocess:     "TaskName": "Training",
LogLearning: Display: Subprocess:     "TaskDirectory": "C:/MyWorkspace/MyProject/Intermediate/LearningAgents/Training14",
LogLearning: Display: Subprocess:     "CommunicationType": "SharedMemory"
LogLearning: Display: Subprocess: }
LogLearning: Display: Subprocess: INFO: Receiving Policy...
LogLearning: Display: Pushing network...
LogLearning: Display: Pushing network...
LogLearning: Display: Subprocess: INFO: Receiving Encoder...
LogLearning: Display: Pushing network...
LogLearning: Display: ImitationTrainer_5: Sending Experience...
LogLearning: Display: Pushing Experience...
LogLearning: Display: ImitationTrainer_5: Stopping training...
LogLearning: Display: Subprocess: INFO: Receiving Decoder...
LogLearning: Display: Subprocess: INFO: Profile| Waiting for Experience...      0ms   GPU Usage| Allocated: 0.00 GiB, Reserved: 0.00 GiB
LogLearning: Display: Subprocess: INFO: Creating Optimizer...
LogLearning: Display: Subprocess: INFO: Begin Training...
LogLearning: Display: Subprocess: INFO: Profile| Training                      15ms   GPU Usage| Allocated: 0.00 GiB, Reserved: 0.00 GiB
LogLearning: Display: Subprocess: Traceback (most recent call last):
LogLearning: Display: Subprocess:   File "C:\MyWorkspace\Engine\Plugins\Experimental\LearningAgents\Content\Python\train.py", line 103, in <module>
LogLearning: Display: Subprocess:     train(communicator.shared_memory_processes[0].config, communicator, get_experiment_trackers(config))
LogLearning: Display: Subprocess:   File "C:\MyWorkspace\Engine\Plugins\Experimental\LearningAgents\Content\Python\learning_agents\train_behavior_cloning.py", line 181, in train
LogLearning: Display: Subprocess:     batch_obs_noise_mask = schema_noise_mask_observation(observation_schema, batch_obs.reshape(-1, batch_obs.shape[-1])).view_as(batch_obs)
LogLearning: Display: Subprocess:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
LogLearning: Display: Subprocess:   File "C:\MyWorkspace\Engine\Plugins\Experimental\LearningAgents\Content\Python\learning_agents\train_common.py", line 799, in schema_noise_mask_observation
LogLearning: Display: Subprocess:     mask[:,smpl_offset:smpl_offset+smpl_size] = schema_noise_mask_observation(
LogLearning: Display: Subprocess:                                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
LogLearning: Display: Subprocess:   File "C:\MyWorkspace\Engine\Plugins\Experimental\LearningAgents\Content\Python\learning_agents\train_common.py", line 897, in schema_noise_mask_observation
LogLearning: Display: Subprocess:     raise Exception(obs_type + ' Not Implemented')
LogLearning: Display: Subprocess: Exception: Conv2d Not Implemented
LogLearning: Display: Subprocess: Exception ignored in: <function SharedMemory.__del__ at 0x000002353C7E3B00>
LogLearning: Display: Subprocess: Traceback (most recent call last):
LogLearning: Display: Subprocess:   File "C:\MyWorkspace\Engine\Plugins\Experimental\LearningAgents\Content\Python\learning_agents\communicators\shared_memory.py", line 193, in __del__
LogLearning: Display: Subprocess:   File "C:\MyWorkspace\Engine\Plugins\Experimental\LearningAgents\Content\Python\learning_agents\communicators\shared_memory.py", line 236, in close

Deathcalibur · December 2, 2025, 7:16pm

Sorry, I think this is something we missed.

We can get it added in the next week or two, and then I can send you the few lines of code you would need to add to that function in the meantime so you don’t have to wait for 5.8.

Thanks for finding this and reporting it!