Update to barracuda 1.3.3 and changes to the model inputs and outputs for LSTM (#5236)

* Initial commit

* making it work with hallway. Added a new model version

* addressing feedback

* Adding a comment

* formatting

* modifying the changelog

* Adding descriptions on the model version descriptions

* Update com.unity.ml-agents/Runtime/Inference/GeneratorImpl.cs

Co-authored-by: Chris Elion <chris.elion@unity3d.com>

* addressing comments

Co-authored-by: Chris Elion <chris.elion@unity3d.com>

com.unity.ml-agents/CHANGELOG.md

 sizes and will need to be retrained. (#5181)
 - The `AbstractBoard` class for integration with Match-3 games was changed to make it easier to support boards with
 different sizes using the same model. For a summary of the interface changes, please see the Migration Guide. (##5189)
+- Updated the Barracuda package to version `1.3.3-preview`(#5236)
 
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 - The `--resume` flag now supports resuming experiments with additional reward providers or
 
 ### Minor Changes
 #### com.unity.ml-agents / com.unity.ml-agents.extensions (C#)
-- The `.onnx` models input names have changed. All input placeholders will now use the prefix `obs_` removing the distinction between visual and vector observations. Models created with this version will not be usable with previous versions of the package (#5080)
+- The `.onnx` models input names have changed. All input placeholders will now use the prefix `obs_` removing the distinction between visual and vector observations. In addition, the inputs and outputs of LSTM changed. Models created with this version will not be usable with previous versions of the package (#5080, #5236)
 - The `.onnx` models discrete action output now contains the discrete actions values and not the logits. Models created with this version will not be usable with previous versions of the package (#5080)
 - Added ML-Agents package settings. (#5027)
 - Make com.unity.modules.unityanalytics an optional dependency. (#5109)

com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs

         public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
         {
             var agentIndex = 0;
-            var memorySize = (int)tensorProxy.shape[tensorProxy.shape.Length - 1];
+            var memorySize = tensorProxy.data.width;
             for (var i = 0; i < actionIds.Count; i++)
             {
                 var agentId = actionIds[i];
                 {
                     memory = new List<float>();
                     memory.AddRange(Enumerable.Repeat(0f, memorySize));
+                }
+
+                for (var j = 0; j < memorySize; j++)
+                {
+                    memory[j] = tensorProxy.data[agentIndex, 0, j, 0];
                 }
 
                 m_Memories[agentId] = memory;

com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs

                 });
             }
 
-            foreach (var mem in model.memories)
+            var modelVersion = model.GetVersion();
+            if (modelVersion < (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
-                tensors.Add(new TensorProxy
+                foreach (var mem in model.memories)
-                    name = mem.input,
-                    valueType = TensorProxy.TensorType.FloatingPoint,
-                    data = null,
-                    shape = TensorUtils.TensorShapeFromBarracuda(mem.shape)
-                });
+                    tensors.Add(new TensorProxy
+                    {
+                        name = mem.input,
+                        valueType = TensorProxy.TensorType.FloatingPoint,
+                        data = null,
+                        shape = TensorUtils.TensorShapeFromBarracuda(mem.shape)
+                    });
+                }
             }
 
             tensors.Sort((el1, el2) => string.Compare(el1.name, el2.name, StringComparison.InvariantCulture));
                 names.Add(model.DiscreteOutputName());
             }
 
+            var modelVersion = model.GetVersion();
-                foreach (var mem in model.memories)
+                if (modelVersion < (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
-                    names.Add(mem.output);
+                    foreach (var mem in model.memories)
+                    {
+                        names.Add(mem.output);
+                    }
+                }
+                else
+                {
+                    names.Add(TensorNames.RecurrentOutput);
                 }
             }
 

com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs

 
         internal enum ModelApiVersion
         {
+            /// <summary>
+            /// ML-Agents model version for versions 1.x.y
+            /// The observations are split between vector and visual observations
+            /// There are legacy action outputs for discrete and continuous actions
+            /// LSTM inputs and outputs are handled by Barracuda
+            /// </summary>
+
+            /// <summary>
+            /// All observations are treated the same and named obs_{i} with i being
+            /// the sensor index
+            /// Legacy "action" output is no longer present
+            /// LSTM inputs and outputs are treated like regular inputs and outputs
+            /// and no longer managed by Barracuda
+            /// </summary>
             MLAgents2_0 = 3,
             MinSupportedVersion = MLAgents1_0,
             MaxSupportedVersion = MLAgents2_0
             // If the model has a non-negative memory size but requires a recurrent input
             if (memory > 0)
             {
-                if (!tensorsNames.Any(x => x.EndsWith("_h")) ||
-                    !tensorsNames.Any(x => x.EndsWith("_c")))
+                var modelVersion = model.GetVersion();
+                var netHasMemories = false;
+                if (modelVersion < (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
+                {
+                    netHasMemories = tensorsNames.Any(x => x.EndsWith("_h")) &&
+                        tensorsNames.Any(x => x.EndsWith("_c"));
+                }
+                else
+                {
+                    netHasMemories = tensorsNames.Any(x => x == TensorNames.RecurrentInPlaceholder);
+                }
+                if (!netHasMemories)
-                        FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size.")
-                        );
+                            FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size.")
+                            );
                 }
             }
 
             // If there is no Recurrent Output but the model is Recurrent.
             if (memory > 0)
             {
-                var memOutputs = model.memories.Select(x => x.output).ToList();
-                if (!memOutputs.Any(x => x.EndsWith("_h")) ||
-                    !memOutputs.Any(x => x.EndsWith("_c")))
+                var netHasMemories = false;
+                var modelVersion = model.GetVersion();
+                if (modelVersion < (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
+                {
+                    var memOutputs = model.memories.Select(x => x.output).ToList();
+                    netHasMemories = memOutputs.Any(x => x.EndsWith("_h")) &&
+                        memOutputs.Any(x => x.EndsWith("_c"));
+                }
+                else
+                {
+                    var allOutputs = model.GetOutputNames().ToList();
+                    netHasMemories = allOutputs.Any(x => x == TensorNames.RecurrentOutput);
+                }
+                if (!netHasMemories)
+
             }
             return failedModelChecks;
         }

com.unity.ml-agents/Runtime/Inference/GeneratorImpl.cs

         {
             TensorUtils.ResizeTensor(tensorProxy, batchSize, m_Allocator);
 
-            var memorySize = tensorProxy.shape[tensorProxy.shape.Length - 1];
+            var memorySize = tensorProxy.data.width;
+
             var agentIndex = 0;
             for (var infoIndex = 0; infoIndex < infos.Count; infoIndex++)
             {
                 {
                     for (var j = 0; j < memorySize; j++)
                     {
-                        tensorProxy.data[agentIndex, j] = 0;
+                        tensorProxy.data[agentIndex, 0, j, 0] = 0;
                     }
                     agentIndex++;
                     continue;
                     {
                         break;
                     }
-                    tensorProxy.data[agentIndex, j] = memory[j];
+                    tensorProxy.data[agentIndex, 0, j, 0] = memory[j];
                 }
                 agentIndex++;
             }

com.unity.ml-agents/Runtime/Inference/TensorApplier.cs

                 var tensorName = model.ContinuousOutputName();
                 m_Dict[tensorName] = new ContinuousActionOutputApplier(actionSpec);
             }
+            var modelVersion = model.GetVersion();
-                var modelVersion = model.GetVersion();
                 if (modelVersion == (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents1_0)
                 {
                     m_Dict[tensorName] = new LegacyDiscreteActionOutputApplier(actionSpec, seed, allocator);
             }
             m_Dict[TensorNames.RecurrentOutput] = new MemoryOutputApplier(memories);
 
-            for (var i = 0; i < model?.memories.Count; i++)
+            if (modelVersion < (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
-                m_Dict[model.memories[i].output] =
-                    new BarracudaMemoryOutputApplier(model.memories.Count, i, memories);
+                for (var i = 0; i < model?.memories.Count; i++)
+                {
+                    m_Dict[model.memories[i].output] =
+                        new BarracudaMemoryOutputApplier(model.memories.Count, i, memories);
+                }
             }
         }
 

com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs

             m_Dict[TensorNames.RecurrentInPlaceholder] =
                 new RecurrentInputGenerator(allocator, memories);
 
-            for (var i = 0; i < model.memories.Count; i++)
+            if (m_ApiVersion < (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
-                m_Dict[model.memories[i].input] =
-                    new BarracudaRecurrentInputGenerator(i, allocator, memories);
+                for (var i = 0; i < model.memories.Count; i++)
+                {
+                    m_Dict[model.memories[i].input] =
+                        new BarracudaRecurrentInputGenerator(i, allocator, memories);
+                }
             }
 
             m_Dict[TensorNames.PreviousActionPlaceholder] =

com.unity.ml-agents/package.json

   "unity": "2019.4",
   "description": "Use state-of-the-art machine learning to create intelligent character behaviors in any Unity environment (games, robotics, film, etc.).",
   "dependencies": {
-    "com.unity.barracuda": "1.3.2-preview",
+    "com.unity.barracuda": "1.3.3-preview",
     "com.unity.modules.imageconversion": "1.0.0",
     "com.unity.modules.jsonserialize": "1.0.0"
   }

ml-agents/mlagents/trainers/torch/layers.py

 import abc
 from typing import Tuple
 from enum import Enum
+from mlagents.trainers.torch.model_serialization import exporting_to_onnx
 
 
 class Swish(torch.nn.Module):
     def forward(
         self, input_tensor: torch.Tensor, memories: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        if exporting_to_onnx.is_exporting():
+            # This transpose is needed both at input and output of the LSTM when
+            # exporting because ONNX will expect (sequence_len, batch, memory_size)
+            # instead of (batch, sequence_len, memory_size)
+            memories = torch.transpose(memories, 0, 1)
+
+
+
+        if exporting_to_onnx.is_exporting():
+            output_mem = torch.transpose(output_mem, 0, 1)
+
         return lstm_out, output_mem

ml-agents/mlagents/trainers/torch/networks.py

 
 
 class SimpleActor(nn.Module, Actor):
-    MODEL_EXPORT_VERSION = 3
+    MODEL_EXPORT_VERSION = 3  # Corresponds to ModelApiVersion.MLAgents2_0
 
     def __init__(
         self,