From 2130c39e1fe3c2d413a68eec2adad8762525a922 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Wed, 10 Mar 2021 12:53:36 -0800
Subject: [PATCH 01/24] Make modelCheck have flavors of error messages

---
 .../Editor/BehaviorParametersEditor.cs        |  13 +-
 .../Inference/BarracudaModelExtensions.cs     |  59 +++++-
 .../Inference/BarracudaModelParamLoader.cs    | 198 ++++++++++++++----
 3 files changed, 218 insertions(+), 52 deletions(-)

diff --git a/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs b/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
index c178bc668f..b1f7c338cb 100644
--- a/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
+++ b/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
@@ -147,7 +147,18 @@ void DisplayFailedModelChecks()
                 {
                     if (check != null)
                     {
-                        EditorGUILayout.HelpBox(check, MessageType.Warning);
+                        if (check.CheckType == Inference.CheckType.Info)
+                        {
+                            EditorGUILayout.HelpBox(check.Message, MessageType.Info);
+                        }
+                        if (check.CheckType == Inference.CheckType.Warning)
+                        {
+                            EditorGUILayout.HelpBox(check.Message, MessageType.Warning);
+                        }
+                        if (check.CheckType == Inference.CheckType.Error)
+                        {
+                            EditorGUILayout.HelpBox(check.Message, MessageType.Error);
+                        }
                     }
                 }
             }
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index c603477a59..73b42ba467 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -298,13 +298,19 @@ public static bool SupportsContinuousAndDiscrete(this Model model)
         /// <param name="failedModelChecks">Output list of failure messages</param>
         ///
         /// <returns>True if the model contains all the expected tensors.</returns>
-        public static bool CheckExpectedTensors(this Model model, List<string> failedModelChecks)
+        public static bool CheckExpectedTensors(this Model model, List<FailedCheck> failedModelChecks)
         {
             // Check the presence of model version
             var modelApiVersionTensor = model.GetTensorByName(TensorNames.VersionNumber);
             if (modelApiVersionTensor == null)
             {
-                failedModelChecks.Add($"Required constant \"{TensorNames.VersionNumber}\" was not found in the model file.");
+                failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
+                    $"Required constant \"{TensorNames.VersionNumber}\" was not found in the model file."
+                    });
                 return false;
             }
 
@@ -312,7 +318,13 @@ public static bool CheckExpectedTensors(this Model model, List<string> failedMod
             var memorySizeTensor = model.GetTensorByName(TensorNames.MemorySize);
             if (memorySizeTensor == null)
             {
-                failedModelChecks.Add($"Required constant \"{TensorNames.MemorySize}\" was not found in the model file.");
+                failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
+                    $"Required constant \"{TensorNames.MemorySize}\" was not found in the model file."
+                    });
                 return false;
             }
 
@@ -321,7 +333,13 @@ public static bool CheckExpectedTensors(this Model model, List<string> failedMod
                 !model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
                 !model.outputs.Contains(TensorNames.DiscreteActionOutput))
             {
-                failedModelChecks.Add("The model does not contain any Action Output Node.");
+                failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
+                    "The model does not contain any Action Output Node."
+                    });
                 return false;
             }
 
@@ -330,13 +348,24 @@ public static bool CheckExpectedTensors(this Model model, List<string> failedMod
             {
                 if (model.GetTensorByName(TensorNames.ActionOutputShapeDeprecated) == null)
                 {
-                    failedModelChecks.Add("The model does not contain any Action Output Shape Node.");
+                    failedModelChecks.Add(
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model does not contain any Action Output Shape Node."
+                        });
                     return false;
                 }
                 if (model.GetTensorByName(TensorNames.IsContinuousControlDeprecated) == null)
                 {
-                    failedModelChecks.Add($"Required constant \"{TensorNames.IsContinuousControlDeprecated}\" was not found in the model file. " +
-                        "This is only required for model that uses a deprecated model format.");
+                    failedModelChecks.Add(new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
+                    $"Required constant \"{TensorNames.IsContinuousControlDeprecated}\" was not found in the model file. " +
+                        "This is only required for model that uses a deprecated model format."
+                    });
                     return false;
                 }
             }
@@ -345,13 +374,25 @@ public static bool CheckExpectedTensors(this Model model, List<string> failedMod
                 if (model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
                     model.GetTensorByName(TensorNames.ContinuousActionOutputShape) == null)
                 {
-                    failedModelChecks.Add("The model uses continuous action but does not contain Continuous Action Output Shape Node.");
+                    failedModelChecks.Add(
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model uses continuous action but does not contain Continuous Action Output Shape Node."
+                        });
                     return false;
                 }
                 if (model.outputs.Contains(TensorNames.DiscreteActionOutput) &&
                     model.GetTensorByName(TensorNames.DiscreteActionOutputShape) == null)
                 {
-                    failedModelChecks.Add("The model uses discrete action but does not contain Discrete Action Output Shape Node.");
+                    failedModelChecks.Add(
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                    "The model uses discrete action but does not contain Discrete Action Output Shape Node."
+                        });
                     return false;
                 }
             }
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index d45e59246b..2d658fd50d 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -8,13 +8,28 @@
 
 namespace Unity.MLAgents.Inference
 {
+
+    internal enum CheckType
+    {
+        Info = 0,
+        Warning = 1,
+        Error = 2
+    }
+
+    internal class FailedCheck
+    {
+        public CheckType CheckType;
+        public string Message;
+    }
+
     /// <summary>
     /// Prepares the Tensors for the Learning Brain and exposes a list of failed checks if Model
     /// and BrainParameters are incompatible.
     /// </summary>
     internal class BarracudaModelParamLoader
     {
-        const long k_ApiVersion = 2;
+        const long k_ApiVersion = 3;
+        const long k_ApiVersionLegacy = 2;
 
         /// <summary>
         /// Factory for the ModelParamLoader : Creates a ModelParamLoader and runs the checks
@@ -31,12 +46,12 @@ internal class BarracudaModelParamLoader
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
         /// <param name="behaviorType">BehaviorType or the Agent to check.</param>
         /// <returns>The list the error messages of the checks that failed</returns>
-        public static IEnumerable<string> CheckModel(Model model, BrainParameters brainParameters,
+        public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters brainParameters,
             ISensor[] sensors, ActuatorComponent[] actuatorComponents,
             int observableAttributeTotalSize = 0,
             BehaviorType behaviorType = BehaviorType.Default)
         {
-            List<string> failedModelChecks = new List<string>();
+            List<FailedCheck> failedModelChecks = new List<FailedCheck>();
             if (model == null)
             {
                 var errorMsg = "There is no model for this Brain; cannot run inference. ";
@@ -48,7 +63,7 @@ public static IEnumerable<string> CheckModel(Model model, BrainParameters brainP
                 {
                     errorMsg += "(But can still train)";
                 }
-                failedModelChecks.Add(errorMsg);
+                failedModelChecks.Add(new FailedCheck { CheckType = CheckType.Warning, Message = errorMsg });
                 return failedModelChecks;
             }
 
@@ -62,22 +77,51 @@ public static IEnumerable<string> CheckModel(Model model, BrainParameters brainP
             if (modelApiVersion == -1)
             {
                 failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
                     "Model was not trained using the right version of ML-Agents. " +
-                    "Cannot use this model.");
+                    "Cannot use this model."
+                    });
                 return failedModelChecks;
             }
-            if (modelApiVersion != k_ApiVersion)
+            if (modelApiVersion != k_ApiVersion && modelApiVersion != k_ApiVersionLegacy)
             {
                 failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Error,
+                        Message =
+                    $"Version of the trainer the model was trained with ({modelApiVersion}) " +
+                    $"is not compatible with the Brain's version ({k_ApiVersion})."
+                    });
+                return failedModelChecks;
+            }
+            if (modelApiVersion == k_ApiVersionLegacy)
+            {
+                failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Info,
+                        Message =
                     $"Version of the trainer the model was trained with ({modelApiVersion}) " +
-                    $"is not compatible with the Brain's version ({k_ApiVersion}).");
+                    $"is not the same as the the Brain's version ({k_ApiVersion}). " +
+                    "The model should still work"
+                    });
                 return failedModelChecks;
             }
 
             var memorySize = (int)model.GetTensorByName(TensorNames.MemorySize)[0];
             if (memorySize == -1)
             {
-                failedModelChecks.Add($"Missing node in the model provided : {TensorNames.MemorySize}");
+                failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
+                    $"Missing node in the model provided : {TensorNames.MemorySize}"
+                    });
                 return failedModelChecks;
             }
 
@@ -113,14 +157,14 @@ public static IEnumerable<string> CheckModel(Model model, BrainParameters brainP
         /// <returns>
         /// A IEnumerable of string corresponding to the failed input presence checks.
         /// </returns>
-        static IEnumerable<string> CheckInputTensorPresence(
+        static IEnumerable<FailedCheck> CheckInputTensorPresence(
             Model model,
             BrainParameters brainParameters,
             int memory,
             ISensor[] sensors
         )
         {
-            var failedModelChecks = new List<string>();
+            var failedModelChecks = new List<FailedCheck>();
             var tensorsNames = model.GetInputNames();
 
             // If there is no Vector Observation Input but the Brain Parameters expect one.
@@ -128,8 +172,13 @@ ISensor[] sensors
                 (!tensorsNames.Contains(TensorNames.VectorObservationPlaceholder)))
             {
                 failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
                     "The model does not contain a Vector Observation Placeholder Input. " +
-                    "You must set the Vector Observation Space Size to 0.");
+                    "You must set the Vector Observation Space Size to 0."
+                    });
             }
 
             // If there are not enough Visual Observation Input compared to what the
@@ -144,8 +193,13 @@ ISensor[] sensors
                         TensorNames.GetVisualObservationName(visObsIndex)))
                     {
                         failedModelChecks.Add(
+                            new FailedCheck
+                            {
+                                CheckType = CheckType.Warning,
+                                Message =
                             "The model does not contain a Visual Observation Placeholder Input " +
-                            $"for sensor component {visObsIndex} ({sensor.GetType().Name}).");
+                            $"for sensor component {visObsIndex} ({sensor.GetType().Name})."
+                            });
                     }
                     visObsIndex++;
                 }
@@ -155,8 +209,13 @@ ISensor[] sensors
                         TensorNames.GetObservationName(sensorIndex)))
                     {
                         failedModelChecks.Add(
+                            new FailedCheck
+                            {
+                                CheckType = CheckType.Warning,
+                                Message =
                             "The model does not contain an Observation Placeholder Input " +
-                            $"for sensor component {sensorIndex} ({sensor.GetType().Name}).");
+                            $"for sensor component {sensorIndex} ({sensor.GetType().Name})."
+                            });
                     }
                 }
 
@@ -167,8 +226,13 @@ ISensor[] sensors
             if (expectedVisualObs > visObsIndex)
             {
                 failedModelChecks.Add(
+                    new FailedCheck
+                    {
+                        CheckType = CheckType.Warning,
+                        Message =
                     $"The model expects {expectedVisualObs} visual inputs," +
                     $" but only found {visObsIndex} visual sensors."
+                    }
                 );
             }
 
@@ -179,7 +243,12 @@ ISensor[] sensors
                     !tensorsNames.Any(x => x.EndsWith("_c")))
                 {
                     failedModelChecks.Add(
-                        "The model does not contain a Recurrent Input Node but has memory_size.");
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model does not contain a Recurrent Input Node but has memory_size."
+                        });
                 }
             }
 
@@ -189,7 +258,12 @@ ISensor[] sensors
                 if (!tensorsNames.Contains(TensorNames.ActionMaskPlaceholder))
                 {
                     failedModelChecks.Add(
-                        "The model does not contain an Action Mask but is using Discrete Control.");
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model does not contain an Action Mask but is using Discrete Control."
+                        });
                 }
             }
             return failedModelChecks;
@@ -206,9 +280,9 @@ ISensor[] sensors
         /// <returns>
         /// A IEnumerable of string corresponding to the failed output presence checks.
         /// </returns>
-        static IEnumerable<string> CheckOutputTensorPresence(Model model, int memory)
+        static IEnumerable<FailedCheck> CheckOutputTensorPresence(Model model, int memory)
         {
-            var failedModelChecks = new List<string>();
+            var failedModelChecks = new List<FailedCheck>();
 
             // If there is no Recurrent Output but the model is Recurrent.
             if (memory > 0)
@@ -219,7 +293,12 @@ static IEnumerable<string> CheckOutputTensorPresence(Model model, int memory)
                     !memOutputs.Any(x => x.EndsWith("_c")))
                 {
                     failedModelChecks.Add(
-                        "The model does not contain a Recurrent Output Node but has memory_size.");
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model does not contain a Recurrent Output Node but has memory_size."
+                        });
                 }
             }
             return failedModelChecks;
@@ -234,7 +313,7 @@ static IEnumerable<string> CheckOutputTensorPresence(Model model, int memory)
         /// If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.
         /// </returns>
-        static string CheckVisualObsShape(
+        static FailedCheck CheckVisualObsShape(
             TensorProxy tensorProxy, ISensor sensor)
         {
             var shape = sensor.GetObservationShape();
@@ -246,9 +325,14 @@ static string CheckVisualObsShape(
             var pixelT = tensorProxy.Channels;
             if ((widthBp != widthT) || (heightBp != heightT) || (pixelBp != pixelT))
             {
-                return $"The visual Observation of the model does not match. " +
+                return new FailedCheck
+                {
+                    CheckType = CheckType.Warning,
+                    Message =
+                    $"The visual Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{widthBp}x{heightBp}x{pixelBp}] but " +
-                    $"was expecting [?x{widthT}x{heightT}x{pixelT}].";
+                    $"was expecting [?x{widthT}x{heightT}x{pixelT}]."
+                };
             }
             return null;
         }
@@ -262,7 +346,7 @@ static string CheckVisualObsShape(
         /// If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.
         /// </returns>
-        static string CheckRankTwoObsShape(
+        static FailedCheck CheckRankTwoObsShape(
             TensorProxy tensorProxy, ISensor sensor)
         {
             var shape = sensor.GetObservationShape();
@@ -272,9 +356,14 @@ static string CheckRankTwoObsShape(
             var dim2T = tensorProxy.Width;
             if ((dim1Bp != dim1T) || (dim2Bp != dim2T))
             {
-                return $"An Observation of the model does not match. " +
+                return new FailedCheck
+                {
+                    CheckType = CheckType.Warning,
+                    Message =
+                    $"An Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{dim1Bp}x{dim2Bp}] but " +
-                    $"was expecting [?x{dim1T}x{dim2T}].";
+                    $"was expecting [?x{dim1T}x{dim2T}]."
+                };
             }
             return null;
         }
@@ -292,13 +381,13 @@ static string CheckRankTwoObsShape(
         /// <param name="sensors">Attached sensors</param>
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
         /// <returns>The list the error messages of the checks that failed</returns>
-        static IEnumerable<string> CheckInputTensorShape(
+        static IEnumerable<FailedCheck> CheckInputTensorShape(
             Model model, BrainParameters brainParameters, ISensor[] sensors,
             int observableAttributeTotalSize)
         {
-            var failedModelChecks = new List<string>();
+            var failedModelChecks = new List<FailedCheck>();
             var tensorTester =
-                new Dictionary<string, Func<BrainParameters, TensorProxy, ISensor[], int, string>>()
+                new Dictionary<string, Func<BrainParameters, TensorProxy, ISensor[], int, FailedCheck>>()
             {
                 {TensorNames.VectorObservationPlaceholder, CheckVectorObsShape},
                 {TensorNames.PreviousActionPlaceholder, CheckPreviousActionShape},
@@ -339,7 +428,12 @@ static IEnumerable<string> CheckInputTensorShape(
                     if (!tensor.name.Contains("visual_observation"))
                     {
                         failedModelChecks.Add(
-                            "Model requires an unknown input named : " + tensor.name);
+                            new FailedCheck
+                            {
+                                CheckType = CheckType.Warning,
+                                Message =
+                            "Model requires an unknown input named : " + tensor.name
+                            });
                     }
                 }
                 else
@@ -369,7 +463,7 @@ static IEnumerable<string> CheckInputTensorShape(
         /// If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.
         /// </returns>
-        static string CheckVectorObsShape(
+        static FailedCheck CheckVectorObsShape(
             BrainParameters brainParameters, TensorProxy tensorProxy, ISensor[] sensors,
             int observableAttributeTotalSize)
         {
@@ -406,11 +500,16 @@ static string CheckVectorObsShape(
                 }
 
                 sensorSizes += "]";
-                return $"Vector Observation Size of the model does not match. Was expecting {totalVecObsSizeT} " +
+                return new FailedCheck
+                {
+                    CheckType = CheckType.Warning,
+                    Message =
+                $"Vector Observation Size of the model does not match. Was expecting {totalVecObsSizeT} " +
                     $"but received: \n" +
                     $"Vector observations: {vecObsSizeBp} x {numStackedVector}\n" +
                     $"Total [Observable] attributes: {observableAttributeTotalSize}\n" +
-                    $"Sensor sizes: {sensorSizes}.";
+                    $"Sensor sizes: {sensorSizes}."
+                };
             }
             return null;
         }
@@ -427,7 +526,7 @@ static string CheckVectorObsShape(
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes (unused).</param>
         /// <returns>If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.</returns>
-        static string CheckPreviousActionShape(
+        static FailedCheck CheckPreviousActionShape(
             BrainParameters brainParameters, TensorProxy tensorProxy,
             ISensor[] sensors, int observableAttributeTotalSize)
         {
@@ -435,8 +534,13 @@ static string CheckPreviousActionShape(
             var numberActionsT = tensorProxy.shape[tensorProxy.shape.Length - 1];
             if (numberActionsBp != numberActionsT)
             {
-                return "Previous Action Size of the model does not match. " +
-                    $"Received {numberActionsBp} but was expecting {numberActionsT}.";
+                return new FailedCheck
+                {
+                    CheckType = CheckType.Warning,
+                    Message =
+                "Previous Action Size of the model does not match. " +
+                    $"Received {numberActionsBp} but was expecting {numberActionsT}."
+                };
             }
             return null;
         }
@@ -456,12 +560,12 @@ static string CheckPreviousActionShape(
         /// A IEnumerable of string corresponding to the incompatible shapes between model
         /// and BrainParameters.
         /// </returns>
-        static IEnumerable<string> CheckOutputTensorShape(
+        static IEnumerable<FailedCheck> CheckOutputTensorShape(
             Model model,
             BrainParameters brainParameters,
             ActuatorComponent[] actuatorComponents)
         {
-            var failedModelChecks = new List<string>();
+            var failedModelChecks = new List<FailedCheck>();
 
             // If the model expects an output but it is not in this list
             var modelContinuousActionSize = model.ContinuousOutputSize();
@@ -494,7 +598,7 @@ static IEnumerable<string> CheckOutputTensorShape(
         /// If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.
         /// </returns>
-        static string CheckDiscreteActionOutputShape(
+        static FailedCheck CheckDiscreteActionOutputShape(
             BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, int modelSumDiscreteBranchSizes)
         {
             // TODO: check each branch size instead of sum of branch sizes
@@ -508,8 +612,13 @@ static string CheckDiscreteActionOutputShape(
 
             if (modelSumDiscreteBranchSizes != sumOfDiscreteBranchSizes)
             {
-                return "Discrete Action Size of the model does not match. The BrainParameters expect " +
-                    $"{sumOfDiscreteBranchSizes} but the model contains {modelSumDiscreteBranchSizes}.";
+                return new FailedCheck
+                {
+                    CheckType = CheckType.Warning,
+                    Message =
+                "Discrete Action Size of the model does not match. The BrainParameters expect " +
+                    $"{sumOfDiscreteBranchSizes} but the model contains {modelSumDiscreteBranchSizes}."
+                };
             }
             return null;
         }
@@ -527,7 +636,7 @@ static string CheckDiscreteActionOutputShape(
         /// </param>
         /// <returns>If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.</returns>
-        static string CheckContinuousActionOutputShape(
+        static FailedCheck CheckContinuousActionOutputShape(
             BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, int modelContinuousActionSize)
         {
             var numContinuousActions = brainParameters.ActionSpec.NumContinuousActions;
@@ -540,8 +649,13 @@ static string CheckContinuousActionOutputShape(
 
             if (modelContinuousActionSize != numContinuousActions)
             {
-                return "Continuous Action Size of the model does not match. The BrainParameters and ActuatorComponents expect " +
-                    $"{numContinuousActions} but the model contains {modelContinuousActionSize}.";
+                return new FailedCheck
+                {
+                    CheckType = CheckType.Warning,
+                    Message =
+                "Continuous Action Size of the model does not match. The BrainParameters and ActuatorComponents expect " +
+                    $"{numContinuousActions} but the model contains {modelContinuousActionSize}."
+                };
             }
             return null;
         }

From a03e8a7316c13de1d120376c6476869421bf3909 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Wed, 10 Mar 2021 16:38:18 -0800
Subject: [PATCH 02/24] ONNX exporter v3

---
 .../Runtime/Inference/ApplierImpl.cs          |  44 +++-
 .../Inference/BarracudaModelExtensions.cs     |  12 +
 .../Inference/BarracudaModelParamLoader.cs    | 248 +++++++++++++++++-
 .../Runtime/Inference/ModelRunner.cs          |   5 +-
 .../Runtime/Inference/TensorApplier.cs        |  11 +-
 .../Runtime/Inference/TensorGenerator.cs      |  94 ++++---
 .../Editor/DiscreteActionOutputApplierTest.cs |   4 +-
 .../EditModeTestInternalBrainTensorApplier.cs |   6 +-
 ...ditModeTestInternalBrainTensorGenerator.cs |   2 +-
 .../mlagents/trainers/torch/distributions.py  |   2 +-
 .../trainers/torch/model_serialization.py     |  61 ++---
 ml-agents/mlagents/trainers/torch/networks.py |  39 +--
 12 files changed, 384 insertions(+), 144 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs b/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs
index 403d1b666b..2daac37eeb 100644
--- a/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs
+++ b/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs
@@ -46,11 +46,51 @@ public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int,
         }
     }
 
+    /// <summary>
+    /// The Applier for the Discrete Action output tensor.
+    /// </summary>
+    internal class DiscreteActionOutputApplier : TensorApplier.IApplier
+    {
+        readonly ActionSpec m_ActionSpec;
+
+
+        public DiscreteActionOutputApplier(ActionSpec actionSpec, int seed, ITensorAllocator allocator)
+        {
+            m_ActionSpec = actionSpec;
+        }
+
+        public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
+        {
+            var agentIndex = 0;
+            var actionSize = tensorProxy.shape[tensorProxy.shape.Length - 1];
+            for (var i = 0; i < actionIds.Count; i++)
+            {
+                var agentId = actionIds[i];
+                if (lastActions.ContainsKey(agentId))
+                {
+                    var actionBuffer = lastActions[agentId];
+                    if (actionBuffer.IsEmpty())
+                    {
+                        actionBuffer = new ActionBuffers(m_ActionSpec);
+                        lastActions[agentId] = actionBuffer;
+                    }
+                    var discreteBuffer = actionBuffer.DiscreteActions;
+                    for (var j = 0; j < actionSize; j++)
+                    {
+                        discreteBuffer[j] = (int)tensorProxy.data[agentIndex, j];
+                    }
+                }
+                agentIndex++;
+            }
+        }
+    }
+
+
     /// <summary>
     /// The Applier for the Discrete Action output tensor. Uses multinomial to sample discrete
     /// actions from the logits contained in the tensor.
     /// </summary>
-    internal class DiscreteActionOutputApplier : TensorApplier.IApplier
+    internal class DiscreteActionWithMultiNomialOutputApplier : TensorApplier.IApplier
     {
         readonly int[] m_ActionSize;
         readonly Multinomial m_Multinomial;
@@ -59,7 +99,7 @@ internal class DiscreteActionOutputApplier : TensorApplier.IApplier
         readonly float[] m_CdfBuffer;
 
 
-        public DiscreteActionOutputApplier(ActionSpec actionSpec, int seed, ITensorAllocator allocator)
+        public DiscreteActionWithMultiNomialOutputApplier(ActionSpec actionSpec, int seed, ITensorAllocator allocator)
         {
             m_ActionSize = actionSpec.BranchSizes;
             m_Multinomial = new Multinomial(seed);
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index 73b42ba467..53a2373329 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -38,6 +38,18 @@ public static string[] GetInputNames(this Model model)
             return names.ToArray();
         }
 
+        /// <summary>
+        /// Get the version of the model.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>The api version of the model</returns>
+        public static int GetVersion(this Model model)
+        {
+            return (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
+        }
+
         /// <summary>
         /// Generates the Tensor inputs that are expected to be present in the Model.
         /// </summary>
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 2d658fd50d..eae159a51d 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -28,8 +28,8 @@ internal class FailedCheck
     /// </summary>
     internal class BarracudaModelParamLoader
     {
-        const long k_ApiVersion = 3;
-        const long k_ApiVersionLegacy = 2;
+        public const long k_ApiVersion = 3;
+        public const long k_ApiVersionLegacy = 2;
 
         /// <summary>
         /// Factory for the ModelParamLoader : Creates a ModelParamLoader and runs the checks
@@ -125,18 +125,35 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                 return failedModelChecks;
             }
 
+            var modelVersion = model.GetVersion();
+            if (modelVersion == k_ApiVersionLegacy)
+            {
+                failedModelChecks.AddRange(
+                    CheckInputTensorPresenceLegacy(model, brainParameters, memorySize, sensors)
+                );
+                failedModelChecks.AddRange(
+                    CheckInputTensorShapeLegacy(model, brainParameters, sensors, observableAttributeTotalSize)
+                );
+            }
+            if (modelVersion == k_ApiVersion)
+            {
+                failedModelChecks.AddRange(
+                    CheckInputTensorPresence(model, brainParameters, memorySize, sensors)
+                );
+                failedModelChecks.AddRange(
+                    CheckInputTensorShape(model, brainParameters, sensors, observableAttributeTotalSize)
+                );
+            }
+
+
+
             failedModelChecks.AddRange(
-                CheckInputTensorPresence(model, brainParameters, memorySize, sensors)
+                CheckOutputTensorShape(model, brainParameters, actuatorComponents)
             );
+
             failedModelChecks.AddRange(
                 CheckOutputTensorPresence(model, memorySize)
             );
-            failedModelChecks.AddRange(
-                CheckInputTensorShape(model, brainParameters, sensors, observableAttributeTotalSize)
-            );
-            failedModelChecks.AddRange(
-                CheckOutputTensorShape(model, brainParameters, actuatorComponents)
-            );
             return failedModelChecks;
         }
 
@@ -157,7 +174,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
         /// <returns>
         /// A IEnumerable of string corresponding to the failed input presence checks.
         /// </returns>
-        static IEnumerable<FailedCheck> CheckInputTensorPresence(
+        static IEnumerable<FailedCheck> CheckInputTensorPresenceLegacy(
             Model model,
             BrainParameters brainParameters,
             int memory,
@@ -269,6 +286,82 @@ ISensor[] sensors
             return failedModelChecks;
         }
 
+        /// <summary>
+        /// Generates failed checks that correspond to inputs expected by the model that are not
+        /// present in the BrainParameters.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters
+        /// </param>
+        /// <param name="brainParameters">
+        /// The BrainParameters that are used verify the compatibility with the InferenceEngine
+        /// </param>
+        /// <param name="memory">
+        /// The memory size that the model is expecting.
+        /// </param>
+        /// <param name="sensors">Array of attached sensor components</param>
+        /// <returns>
+        /// A IEnumerable of string corresponding to the failed input presence checks.
+        /// </returns>
+        static IEnumerable<FailedCheck> CheckInputTensorPresence(
+            Model model,
+            BrainParameters brainParameters,
+            int memory,
+            ISensor[] sensors
+        )
+        {
+            var failedModelChecks = new List<FailedCheck>();
+            var tensorsNames = model.GetInputNames();
+            for (var sensorIndex = 0; sensorIndex < sensors.Length; sensorIndex++)
+            {
+                if (!tensorsNames.Contains(
+                    TensorNames.GetObservationName(sensorIndex)))
+                {
+                    var sensor = sensors[sensorIndex];
+                    failedModelChecks.Add(
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model does not contain an Observation Placeholder Input " +
+                        $"for sensor component {sensorIndex} ({sensor.GetType().Name})."
+                        });
+                }
+            }
+
+            // 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")))
+                {
+                    failedModelChecks.Add(
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model does not contain a Recurrent Input Node but has memory_size."
+                        });
+                }
+            }
+
+            // If the model uses discrete control but does not have an input for action masks
+            if (model.HasDiscreteOutputs())
+            {
+                if (!tensorsNames.Contains(TensorNames.ActionMaskPlaceholder))
+                {
+                    failedModelChecks.Add(
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "The model does not contain an Action Mask but is using Discrete Control."
+                        });
+                }
+            }
+            return failedModelChecks;
+        }
+
         /// <summary>
         /// Generates failed checks that correspond to outputs expected by the model that are not
         /// present in the BrainParameters.
@@ -354,15 +447,61 @@ static FailedCheck CheckRankTwoObsShape(
             var dim2Bp = shape[1];
             var dim1T = tensorProxy.Channels;
             var dim2T = tensorProxy.Width;
+            var dim3T = tensorProxy.Height;
             if ((dim1Bp != dim1T) || (dim2Bp != dim2T))
             {
+                var proxyDimStr = $"[?x{dim1T}x{dim2T}]";
+                if (dim3T > 0)
+                {
+                    proxyDimStr = $"[?x{dim3T}x{dim2T}x{dim1T}]";
+                }
                 return new FailedCheck
                 {
                     CheckType = CheckType.Warning,
                     Message =
                     $"An Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{dim1Bp}x{dim2Bp}] but " +
-                    $"was expecting [?x{dim1T}x{dim2T}]."
+                    $"was expecting {proxyDimStr}."
+                };
+            }
+            return null;
+        }
+
+        /// <summary>
+        /// Checks that the shape of the rank 2 observation input placeholder is the same as the corresponding sensor.
+        /// </summary>
+        /// <param name="tensorProxy">The tensor that is expected by the model</param>
+        /// <param name="sensor">The sensor that produces the visual observation.</param>
+        /// <returns>
+        /// If the Check failed, returns a string containing information about why the
+        /// check failed. If the check passed, returns null.
+        /// </returns>
+        static FailedCheck CheckRankOneObsShape(
+            TensorProxy tensorProxy, ISensor sensor)
+        {
+            var shape = sensor.GetObservationShape();
+            var dim1Bp = shape[0];
+            var dim1T = tensorProxy.Channels;
+            var dim2T = tensorProxy.Width;
+            var dim3T = tensorProxy.Height;
+            if ((dim1Bp != dim1T))
+            {
+                var proxyDimStr = $"[?x{dim1T}]";
+                if (dim2T > 0)
+                {
+                    proxyDimStr = $"[?x{dim1T}x{dim2T}]";
+                }
+                if (dim3T > 0)
+                {
+                    proxyDimStr = $"[?x{dim3T}x{dim2T}x{dim1T}]";
+                }
+                return new FailedCheck
+                {
+                    CheckType = CheckType.Warning,
+                    Message =
+                    $"An Observation of the model does not match. " +
+                    $"Received TensorProxy of shape [?x{dim1Bp}] but " +
+                    $"was expecting {proxyDimStr}."
                 };
             }
             return null;
@@ -381,7 +520,7 @@ static FailedCheck CheckRankTwoObsShape(
         /// <param name="sensors">Attached sensors</param>
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
         /// <returns>The list the error messages of the checks that failed</returns>
-        static IEnumerable<FailedCheck> CheckInputTensorShape(
+        static IEnumerable<FailedCheck> CheckInputTensorShapeLegacy(
             Model model, BrainParameters brainParameters, ISensor[] sensors,
             int observableAttributeTotalSize)
         {
@@ -389,7 +528,7 @@ static IEnumerable<FailedCheck> CheckInputTensorShape(
             var tensorTester =
                 new Dictionary<string, Func<BrainParameters, TensorProxy, ISensor[], int, FailedCheck>>()
             {
-                {TensorNames.VectorObservationPlaceholder, CheckVectorObsShape},
+                {TensorNames.VectorObservationPlaceholder, CheckVectorObsShapeLegacy},
                 {TensorNames.PreviousActionPlaceholder, CheckPreviousActionShape},
                 {TensorNames.RandomNormalEpsilonPlaceholder, ((bp, tensor, scs, i) => null)},
                 {TensorNames.ActionMaskPlaceholder, ((bp, tensor, scs, i) => null)},
@@ -463,7 +602,7 @@ static IEnumerable<FailedCheck> CheckInputTensorShape(
         /// If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.
         /// </returns>
-        static FailedCheck CheckVectorObsShape(
+        static FailedCheck CheckVectorObsShapeLegacy(
             BrainParameters brainParameters, TensorProxy tensorProxy, ISensor[] sensors,
             int observableAttributeTotalSize)
         {
@@ -514,6 +653,87 @@ static FailedCheck CheckVectorObsShape(
             return null;
         }
 
+
+        /// <summary>
+        /// Generates failed checks that correspond to inputs shapes incompatibilities between
+        /// the model and the BrainParameters.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters
+        /// </param>
+        /// <param name="brainParameters">
+        /// The BrainParameters that are used verify the compatibility with the InferenceEngine
+        /// </param>
+        /// <param name="sensors">Attached sensors</param>
+        /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
+        /// <returns>The list the error messages of the checks that failed</returns>
+        static IEnumerable<FailedCheck> CheckInputTensorShape(
+            Model model, BrainParameters brainParameters, ISensor[] sensors,
+            int observableAttributeTotalSize)
+        {
+            var failedModelChecks = new List<FailedCheck>();
+            var tensorTester =
+                new Dictionary<string, Func<BrainParameters, TensorProxy, ISensor[], int, FailedCheck>>()
+            {
+                {TensorNames.PreviousActionPlaceholder, CheckPreviousActionShape},
+                {TensorNames.RandomNormalEpsilonPlaceholder, ((bp, tensor, scs, i) => null)},
+                {TensorNames.ActionMaskPlaceholder, ((bp, tensor, scs, i) => null)},
+                {TensorNames.SequenceLengthPlaceholder, ((bp, tensor, scs, i) => null)},
+                {TensorNames.RecurrentInPlaceholder, ((bp, tensor, scs, i) => null)},
+            };
+
+            foreach (var mem in model.memories)
+            {
+                tensorTester[mem.input] = ((bp, tensor, scs, i) => null);
+            }
+
+            for (var sensorIndex = 0; sensorIndex < sensors.Length; sensorIndex++)
+            {
+                var sens = sensors[sensorIndex];
+                if (sens.GetObservationShape().Length == 3)
+                {
+                    tensorTester[TensorNames.GetObservationName(sensorIndex)] =
+                        (bp, tensor, scs, i) => CheckVisualObsShape(tensor, sens);
+                }
+                if (sens.GetObservationShape().Length == 2)
+                {
+                    tensorTester[TensorNames.GetObservationName(sensorIndex)] =
+                        (bp, tensor, scs, i) => CheckRankTwoObsShape(tensor, sens);
+                }
+                if (sens.GetObservationShape().Length == 1)
+                {
+                    tensorTester[TensorNames.GetObservationName(sensorIndex)] =
+                        (bp, tensor, scs, i) => CheckRankOneObsShape(tensor, sens);
+                }
+
+            }
+
+            // If the model expects an input but it is not in this list
+            foreach (var tensor in model.GetInputTensors())
+            {
+                if (!tensorTester.ContainsKey(tensor.name))
+                {
+                    failedModelChecks.Add(
+                        new FailedCheck
+                        {
+                            CheckType = CheckType.Warning,
+                            Message =
+                        "Model requires an unknown input named : " + tensor.name
+                        });
+                }
+                else
+                {
+                    var tester = tensorTester[tensor.name];
+                    var error = tester.Invoke(brainParameters, tensor, sensors, observableAttributeTotalSize);
+                    if (error != null)
+                    {
+                        failedModelChecks.Add(error);
+                    }
+                }
+            }
+            return failedModelChecks;
+        }
+
         /// <summary>
         /// Checks that the shape of the Previous Vector Action input placeholder is the same in the
         /// model and in the Brain Parameters.
diff --git a/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs b/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
index 12a8188c52..4eecf84862 100644
--- a/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
+++ b/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
@@ -97,10 +97,11 @@ public ModelRunner(
 
             m_InferenceInputs = barracudaModel.GetInputTensors();
             m_OutputNames = barracudaModel.GetOutputNames();
+            var apiVersion = barracudaModel.GetVersion();
             m_TensorGenerator = new TensorGenerator(
-                seed, m_TensorAllocator, m_Memories, barracudaModel);
+                apiVersion, seed, m_TensorAllocator, m_Memories, barracudaModel);
             m_TensorApplier = new TensorApplier(
-                actionSpec, seed, m_TensorAllocator, m_Memories, barracudaModel);
+                apiVersion, actionSpec, seed, m_TensorAllocator, m_Memories, barracudaModel);
             m_InputsByName = new Dictionary<string, Tensor>();
             m_InferenceOutputs = new List<TensorProxy>();
         }
diff --git a/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs b/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
index d3d2d393f7..de74fb5b91 100644
--- a/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
+++ b/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
@@ -39,12 +39,14 @@ public interface IApplier
         /// <summary>
         /// Returns a new TensorAppliers object.
         /// </summary>
+        /// <param name="modelVersion"> The API version of the model.</param>
         /// <param name="actionSpec"> Description of the actions for the Agent.</param>
         /// <param name="seed"> The seed the Appliers will be initialized with.</param>
         /// <param name="allocator"> Tensor allocator</param>
         /// <param name="memories">Dictionary of AgentInfo.id to memory used to pass to the inference model.</param>
         /// <param name="barracudaModel"></param>
         public TensorApplier(
+            int modelVersion,
             ActionSpec actionSpec,
             int seed,
             ITensorAllocator allocator,
@@ -70,7 +72,14 @@ public TensorApplier(
             if (actionSpec.NumDiscreteActions > 0)
             {
                 var tensorName = model.DiscreteOutputName();
-                m_Dict[tensorName] = new DiscreteActionOutputApplier(actionSpec, seed, allocator);
+                if (modelVersion == BarracudaModelParamLoader.k_ApiVersionLegacy)
+                {
+                    m_Dict[tensorName] = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, seed, allocator);
+                }
+                if (modelVersion == BarracudaModelParamLoader.k_ApiVersion)
+                {
+                    m_Dict[tensorName] = new DiscreteActionOutputApplier(actionSpec, seed, allocator);
+                }
             }
             m_Dict[TensorNames.RecurrentOutput] = new MemoryOutputApplier(memories);
 
diff --git a/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs b/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs
index 7615de9d9d..e0f20a319b 100644
--- a/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs
+++ b/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs
@@ -35,20 +35,24 @@ void Generate(
         }
 
         readonly Dictionary<string, IGenerator> m_Dict = new Dictionary<string, IGenerator>();
+        int m_ApiVersion;
 
         /// <summary>
         /// Returns a new TensorGenerators object.
         /// </summary>
+        /// <param name="modelVersion"> The API version of the model.</param>
         /// <param name="seed"> The seed the Generators will be initialized with.</param>
         /// <param name="allocator"> Tensor allocator.</param>
         /// <param name="memories">Dictionary of AgentInfo.id to memory for use in the inference model.</param>
         /// <param name="barracudaModel"></param>
         public TensorGenerator(
+            int modelVersion,
             int seed,
             ITensorAllocator allocator,
             Dictionary<int, List<float>> memories,
             object barracudaModel = null)
         {
+            m_ApiVersion = modelVersion;
             // If model is null, no inference to run and exception is thrown before reaching here.
             if (barracudaModel == null)
             {
@@ -93,47 +97,61 @@ public TensorGenerator(
 
         public void InitializeObservations(List<ISensor> sensors, ITensorAllocator allocator)
         {
-            // Loop through the sensors on a representative agent.
-            // All vector observations use a shared ObservationGenerator since they are concatenated.
-            // All other observations use a unique ObservationInputGenerator
-            var visIndex = 0;
-            ObservationGenerator vecObsGen = null;
-            for (var sensorIndex = 0; sensorIndex < sensors.Count; sensorIndex++)
+            if (m_ApiVersion == BarracudaModelParamLoader.k_ApiVersionLegacy)
             {
-                var sensor = sensors[sensorIndex];
-                var shape = sensor.GetObservationShape();
-                var rank = shape.Length;
-                ObservationGenerator obsGen = null;
-                string obsGenName = null;
-                switch (rank)
+                // Loop through the sensors on a representative agent.
+                // All vector observations use a shared ObservationGenerator since they are concatenated.
+                // All other observations use a unique ObservationInputGenerator
+                var visIndex = 0;
+                ObservationGenerator vecObsGen = null;
+                for (var sensorIndex = 0; sensorIndex < sensors.Count; sensorIndex++)
                 {
-                    case 1:
-                        if (vecObsGen == null)
-                        {
-                            vecObsGen = new ObservationGenerator(allocator);
-                        }
-                        obsGen = vecObsGen;
-                        obsGenName = TensorNames.VectorObservationPlaceholder;
-                        break;
-                    case 2:
-                        // If the tensor is of rank 2, we use the index of the sensor
-                        // to create the name
-                        obsGen = new ObservationGenerator(allocator);
-                        obsGenName = TensorNames.GetObservationName(sensorIndex);
-                        break;
-                    case 3:
-                        // If the tensor is of rank 3, we use the "visual observation
-                        // index", which only counts the rank 3 sensors
-                        obsGen = new ObservationGenerator(allocator);
-                        obsGenName = TensorNames.GetVisualObservationName(visIndex);
-                        visIndex++;
-                        break;
-                    default:
-                        throw new UnityAgentsException(
-                            $"Sensor {sensor.GetName()} have an invalid rank {rank}");
+                    var sensor = sensors[sensorIndex];
+                    var shape = sensor.GetObservationShape();
+                    var rank = shape.Length;
+                    ObservationGenerator obsGen = null;
+                    string obsGenName = null;
+                    switch (rank)
+                    {
+                        case 1:
+                            if (vecObsGen == null)
+                            {
+                                vecObsGen = new ObservationGenerator(allocator);
+                            }
+                            obsGen = vecObsGen;
+                            obsGenName = TensorNames.VectorObservationPlaceholder;
+                            break;
+                        case 2:
+                            // If the tensor is of rank 2, we use the index of the sensor
+                            // to create the name
+                            obsGen = new ObservationGenerator(allocator);
+                            obsGenName = TensorNames.GetObservationName(sensorIndex);
+                            break;
+                        case 3:
+                            // If the tensor is of rank 3, we use the "visual observation
+                            // index", which only counts the rank 3 sensors
+                            obsGen = new ObservationGenerator(allocator);
+                            obsGenName = TensorNames.GetVisualObservationName(visIndex);
+                            visIndex++;
+                            break;
+                        default:
+                            throw new UnityAgentsException(
+                                $"Sensor {sensor.GetName()} have an invalid rank {rank}");
+                    }
+                    obsGen.AddSensorIndex(sensorIndex);
+                    m_Dict[obsGenName] = obsGen;
+                }
+            }
+            if (m_ApiVersion == BarracudaModelParamLoader.k_ApiVersion)
+            {
+                for (var sensorIndex = 0; sensorIndex < sensors.Count; sensorIndex++)
+                {
+                    var obsGen = new ObservationGenerator(allocator);
+                    var obsGenName = TensorNames.GetObservationName(sensorIndex);
+                    obsGen.AddSensorIndex(sensorIndex);
+                    m_Dict[obsGenName] = obsGen;
+
                 }
-                obsGen.AddSensorIndex(sensorIndex);
-                m_Dict[obsGenName] = obsGen;
             }
         }
 
diff --git a/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs b/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
index b431f4f740..6635b50b6e 100644
--- a/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
+++ b/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
@@ -6,7 +6,7 @@
 
 namespace Unity.MLAgents.Tests
 {
-    public class DiscreteActionOutputApplierTest
+    public class DiscreteActionWithMultiNomialOutputApplierTest
     {
         [Test]
         public void TestDiscreteApply()
@@ -30,7 +30,7 @@ public void TestDiscreteApply()
                 valueType = TensorProxy.TensorType.FloatingPoint
             };
 
-            var applier = new DiscreteActionOutputApplier(actionSpec, 2020, null);
+            var applier = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, 2020, null);
             var agentIds = new List<int> { 42, 1337 };
             var actionBuffers = new Dictionary<int, ActionBuffers>();
             actionBuffers[42] = new ActionBuffers(actionSpec);
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
index 6a3af40f9b..bd478be201 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
@@ -18,7 +18,7 @@ public void Construction()
             var actionSpec = new ActionSpec();
             var alloc = new TensorCachingAllocator();
             var mem = new Dictionary<int, List<float>>();
-            var tensorGenerator = new TensorApplier(actionSpec, 0, alloc, mem);
+            var tensorGenerator = new TensorApplier(2, actionSpec, 0, alloc, mem);
             Assert.IsNotNull(tensorGenerator);
             alloc.Dispose();
         }
@@ -64,7 +64,7 @@ public void ApplyDiscreteActionOutput()
                     new[] { 0.5f, 22.5f, 0.1f, 5f, 1f, 4f, 5f, 6f, 7f, 8f })
             };
             var alloc = new TensorCachingAllocator();
-            var applier = new DiscreteActionOutputApplier(actionSpec, 0, alloc);
+            var applier = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, 0, alloc);
 
             var agentIds = new List<int>() { 0, 1 };
             // Dictionary from AgentId to Action
@@ -100,7 +100,7 @@ public void ApplyHybridActionOutput()
             };
             var continuousApplier = new ContinuousActionOutputApplier(actionSpec);
             var alloc = new TensorCachingAllocator();
-            var discreteApplier = new DiscreteActionOutputApplier(actionSpec, 0, alloc);
+            var discreteApplier = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, 0, alloc);
 
             var agentIds = new List<int>() { 0, 1 };
             // Dictionary from AgentId to Action
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
index e1990c772b..4904d2bea2 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
@@ -68,7 +68,7 @@ public void Construction()
         {
             var alloc = new TensorCachingAllocator();
             var mem = new Dictionary<int, List<float>>();
-            var tensorGenerator = new TensorGenerator(0, alloc, mem);
+            var tensorGenerator = new TensorGenerator(0, 2, alloc, mem);
             Assert.IsNotNull(tensorGenerator);
             alloc.Dispose();
         }
diff --git a/ml-agents/mlagents/trainers/torch/distributions.py b/ml-agents/mlagents/trainers/torch/distributions.py
index 2a2bc7b2af..1f5960d10b 100644
--- a/ml-agents/mlagents/trainers/torch/distributions.py
+++ b/ml-agents/mlagents/trainers/torch/distributions.py
@@ -133,7 +133,7 @@ def entropy(self):
         ).unsqueeze(-1)
 
     def exported_model_output(self):
-        return self.all_log_prob()
+        return self.sample()
 
 
 class GaussianDistribution(nn.Module):
diff --git a/ml-agents/mlagents/trainers/torch/model_serialization.py b/ml-agents/mlagents/trainers/torch/model_serialization.py
index f924fae4b8..c7911e5008 100644
--- a/ml-agents/mlagents/trainers/torch/model_serialization.py
+++ b/ml-agents/mlagents/trainers/torch/model_serialization.py
@@ -1,3 +1,4 @@
+from typing import Tuple
 import threading
 from mlagents.torch_utils import torch
 
@@ -90,28 +91,13 @@ def __init__(self, policy):
         observation_specs = self.policy.behavior_spec.observation_specs
         batch_dim = [1]
         seq_len_dim = [1]
-        vec_obs_size = 0
-        for obs_spec in observation_specs:
-            if len(obs_spec.shape) == 1:
-                vec_obs_size += obs_spec.shape[0]
-        num_vis_obs = sum(
-            1 for obs_spec in observation_specs if len(obs_spec.shape) == 3
-        )
-        dummy_vec_obs = [torch.zeros(batch_dim + [vec_obs_size])]
-        # create input shape of NCHW
-        # (It's NHWC in observation_specs.shape)
-        dummy_vis_obs = [
+        num_obs = len(observation_specs)
+
+        dummy_obs = [
             torch.zeros(
-                batch_dim + [obs_spec.shape[2], obs_spec.shape[0], obs_spec.shape[1]]
+                batch_dim + list(ModelSerializer._get_onnx_shape(obs_spec.shape))
             )
             for obs_spec in observation_specs
-            if len(obs_spec.shape) == 3
-        ]
-
-        dummy_var_len_obs = [
-            torch.zeros(batch_dim + [obs_spec.shape[0], obs_spec.shape[1]])
-            for obs_spec in observation_specs
-            if len(obs_spec.shape) == 2
         ]
 
         dummy_masks = torch.ones(
@@ -121,20 +107,9 @@ def __init__(self, policy):
             batch_dim + seq_len_dim + [self.policy.export_memory_size]
         )
 
-        self.dummy_input = (
-            dummy_vec_obs,
-            dummy_vis_obs,
-            dummy_var_len_obs,
-            dummy_masks,
-            dummy_memories,
-        )
+        self.dummy_input = (dummy_obs, dummy_masks, dummy_memories)
 
-        self.input_names = [TensorNames.vector_observation_placeholder]
-        for i in range(num_vis_obs):
-            self.input_names.append(TensorNames.get_visual_observation_name(i))
-        for i, obs_spec in enumerate(observation_specs):
-            if len(obs_spec.shape) == 2:
-                self.input_names.append(TensorNames.get_observation_name(i))
+        self.input_names = [TensorNames.get_observation_name(i) for i in range(num_obs)]
         self.input_names += [
             TensorNames.action_mask_placeholder,
             TensorNames.recurrent_in_placeholder,
@@ -157,22 +132,20 @@ def __init__(self, policy):
                 TensorNames.discrete_action_output_shape,
             ]
             self.dynamic_axes.update({TensorNames.discrete_action_output: {0: "batch"}})
-        if (
-            self.policy.behavior_spec.action_spec.continuous_size == 0
-            or self.policy.behavior_spec.action_spec.discrete_size == 0
-        ):
-            self.output_names += [
-                TensorNames.action_output_deprecated,
-                TensorNames.is_continuous_control_deprecated,
-                TensorNames.action_output_shape_deprecated,
-            ]
-            self.dynamic_axes.update(
-                {TensorNames.action_output_deprecated: {0: "batch"}}
-            )
 
         if self.policy.export_memory_size > 0:
             self.output_names += [TensorNames.recurrent_output]
 
+    @staticmethod
+    def _get_onnx_shape(shape: Tuple[int, ...]) -> Tuple[int, ...]:
+        """
+        Converts the shape of an observation to be compatible with the NCHW format
+        of ONNX
+        """
+        if len(shape) == 3:
+            return shape[2], shape[0], shape[1]
+        return shape
+
     def export_policy_model(self, output_filepath: str) -> None:
         """
         Exports a Torch model for a Policy to .onnx format for Unity embedding.
diff --git a/ml-agents/mlagents/trainers/torch/networks.py b/ml-agents/mlagents/trainers/torch/networks.py
index eaf8e471eb..1c626455e5 100644
--- a/ml-agents/mlagents/trainers/torch/networks.py
+++ b/ml-agents/mlagents/trainers/torch/networks.py
@@ -286,9 +286,7 @@ def get_stats(
     @abc.abstractmethod
     def forward(
         self,
-        vec_inputs: List[torch.Tensor],
-        vis_inputs: List[torch.Tensor],
-        var_len_inputs: List[torch.Tensor],
+        inputs: List[torch.Tensor],
         masks: Optional[torch.Tensor] = None,
         memories: Optional[torch.Tensor] = None,
     ) -> Tuple[Union[int, torch.Tensor], ...]:
@@ -312,7 +310,7 @@ def __init__(
         super().__init__()
         self.action_spec = action_spec
         self.version_number = torch.nn.Parameter(
-            torch.Tensor([2.0]), requires_grad=False
+            torch.Tensor([3.0]), requires_grad=False
         )
         self.is_continuous_int_deprecated = torch.nn.Parameter(
             torch.Tensor([int(self.action_spec.is_continuous())]), requires_grad=False
@@ -387,9 +385,7 @@ def get_stats(
 
     def forward(
         self,
-        vec_inputs: List[torch.Tensor],
-        vis_inputs: List[torch.Tensor],
-        var_len_inputs: List[torch.Tensor],
+        inputs: List[torch.Tensor],
         masks: Optional[torch.Tensor] = None,
         memories: Optional[torch.Tensor] = None,
     ) -> Tuple[Union[int, torch.Tensor], ...]:
@@ -399,28 +395,6 @@ def forward(
         At this moment, torch.onnx.export() doesn't accept None as tensor to be exported,
         so the size of return tuple varies with action spec.
         """
-        # This code will convert the vec and vis obs into a list of inputs for the network
-        concatenated_vec_obs = vec_inputs[0]
-        inputs = []
-        start = 0
-        end = 0
-        vis_index = 0
-        var_len_index = 0
-        for i, enc in enumerate(self.network_body.processors):
-            if isinstance(enc, VectorInput):
-                # This is a vec_obs
-                vec_size = self.network_body.embedding_sizes[i]
-                end = start + vec_size
-                inputs.append(concatenated_vec_obs[:, start:end])
-                start = end
-            elif isinstance(enc, EntityEmbedding):
-                inputs.append(var_len_inputs[var_len_index])
-                var_len_index += 1
-            else:  # visual input
-                inputs.append(vis_inputs[vis_index])
-                vis_index += 1
-
-        # End of code to convert the vec and vis obs into a list of inputs for the network
         encoding, memories_out = self.network_body(
             inputs, memories=memories, sequence_length=1
         )
@@ -435,13 +409,6 @@ def forward(
             export_out += [cont_action_out, self.continuous_act_size_vector]
         if self.action_spec.discrete_size > 0:
             export_out += [disc_action_out, self.discrete_act_size_vector]
-        # Only export deprecated nodes with non-hybrid action spec
-        if self.action_spec.continuous_size == 0 or self.action_spec.discrete_size == 0:
-            export_out += [
-                action_out_deprecated,
-                self.is_continuous_int_deprecated,
-                self.act_size_vector_deprecated,
-            ]
         if self.network_body.memory_size > 0:
             export_out += [memories_out]
         return tuple(export_out)

From ad7fc3ab17da4868eeaa6ff8a37f1a5723a0aae4 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Wed, 10 Mar 2021 17:01:41 -0800
Subject: [PATCH 03/24] Using a better CheckType and a switch statement

---
 .../Editor/BehaviorParametersEditor.cs        | 23 +++++++++-------
 .../Inference/BarracudaModelExtensions.cs     |  2 ++
 .../Inference/BarracudaModelParamLoader.cs    | 27 ++++++++++---------
 3 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs b/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
index b1f7c338cb..5b003d9851 100644
--- a/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
+++ b/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
@@ -5,6 +5,7 @@
 using Unity.MLAgents.Policies;
 using Unity.MLAgents.Sensors;
 using Unity.MLAgents.Sensors.Reflection;
+using CheckType = Unity.MLAgents.Inference.BarracudaModelParamLoader.CheckType;
 
 namespace Unity.MLAgents.Editor
 {
@@ -147,17 +148,19 @@ void DisplayFailedModelChecks()
                 {
                     if (check != null)
                     {
-                        if (check.CheckType == Inference.CheckType.Info)
+                        switch (check.CheckType)
                         {
-                            EditorGUILayout.HelpBox(check.Message, MessageType.Info);
-                        }
-                        if (check.CheckType == Inference.CheckType.Warning)
-                        {
-                            EditorGUILayout.HelpBox(check.Message, MessageType.Warning);
-                        }
-                        if (check.CheckType == Inference.CheckType.Error)
-                        {
-                            EditorGUILayout.HelpBox(check.Message, MessageType.Error);
+                            case CheckType.Info:
+                                EditorGUILayout.HelpBox(check.Message, MessageType.Info);
+                                break;
+                            case CheckType.Warning:
+                                EditorGUILayout.HelpBox(check.Message, MessageType.Warning);
+                                break;
+                            case CheckType.Error:
+                                EditorGUILayout.HelpBox(check.Message, MessageType.Error);
+                                break;
+                            default:
+                                break;
                         }
                     }
                 }
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index 53a2373329..17e7103d83 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -1,6 +1,8 @@
 using System.Collections.Generic;
 using System.Linq;
 using Unity.Barracuda;
+using FailedCheck = Unity.MLAgents.Inference.BarracudaModelParamLoader.FailedCheck;
+using CheckType = Unity.MLAgents.Inference.BarracudaModelParamLoader.CheckType;
 
 namespace Unity.MLAgents.Inference
 {
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index eae159a51d..468d06b842 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -9,25 +9,26 @@
 namespace Unity.MLAgents.Inference
 {
 
-    internal enum CheckType
-    {
-        Info = 0,
-        Warning = 1,
-        Error = 2
-    }
-
-    internal class FailedCheck
-    {
-        public CheckType CheckType;
-        public string Message;
-    }
-
     /// <summary>
     /// Prepares the Tensors for the Learning Brain and exposes a list of failed checks if Model
     /// and BrainParameters are incompatible.
     /// </summary>
     internal class BarracudaModelParamLoader
     {
+
+        internal enum CheckType
+        {
+            Info = 0,
+            Warning = 1,
+            Error = 2
+        }
+
+        internal class FailedCheck
+        {
+            public CheckType CheckType;
+            public string Message;
+        }
+
         public const long k_ApiVersion = 3;
         public const long k_ApiVersionLegacy = 2;
 

From 59cd5f950ff97dbbe09af137022c36882f237513 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Wed, 10 Mar 2021 17:08:22 -0800
Subject: [PATCH 04/24] Removing unused message

---
 .../Runtime/Inference/BarracudaModelParamLoader.cs  | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 468d06b842..c63f3b3797 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -99,19 +99,6 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                     });
                 return failedModelChecks;
             }
-            if (modelApiVersion == k_ApiVersionLegacy)
-            {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Info,
-                        Message =
-                    $"Version of the trainer the model was trained with ({modelApiVersion}) " +
-                    $"is not the same as the the Brain's version ({k_ApiVersion}). " +
-                    "The model should still work"
-                    });
-                return failedModelChecks;
-            }
 
             var memorySize = (int)model.GetTensorByName(TensorNames.MemorySize)[0];
             if (memorySize == -1)

From 4dec29b75a37df4ff4231c930c827bfa7ea60bdb Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Wed, 10 Mar 2021 17:27:11 -0800
Subject: [PATCH 05/24] More tests

---
 .../Editor/DiscreteActionOutputApplierTest.cs | 39 +++++++++
 .../EditModeTestInternalBrainTensorApplier.cs | 84 +++++++++++++++++--
 ...ditModeTestInternalBrainTensorGenerator.cs |  7 +-
 3 files changed, 122 insertions(+), 8 deletions(-)

diff --git a/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs b/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
index 6635b50b6e..44a35a1e0c 100644
--- a/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
+++ b/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
@@ -6,6 +6,45 @@
 
 namespace Unity.MLAgents.Tests
 {
+
+    public class DiscreteActionOutputApplierTest
+    {
+        [Test]
+        public void TestDiscreteApply()
+        {
+            var actionSpec = ActionSpec.MakeDiscrete(3, 2);
+
+            var applier = new DiscreteActionOutputApplier(actionSpec, 2020, null);
+            var agentIds = new List<int> { 42, 1337 };
+            var actionBuffers = new Dictionary<int, ActionBuffers>();
+            actionBuffers[42] = new ActionBuffers(actionSpec);
+            actionBuffers[1337] = new ActionBuffers(actionSpec);
+
+            var actionTensor = new TensorProxy
+            {
+                data = new Tensor(
+                    2,
+                    2,
+                    new[]
+                    {
+                        2.0f, // Agent 0, branch 0
+                        1.0f, // Agent 0, branch 1
+                        0.0f, // Agent 1, branch 0
+                        0.0f  // Agent 1, branch 1
+                    }),
+                shape = new long[] { 2, 2 },
+                valueType = TensorProxy.TensorType.FloatingPoint
+            };
+
+            applier.Apply(actionTensor, agentIds, actionBuffers);
+            Assert.AreEqual(2, actionBuffers[42].DiscreteActions[0]);
+            Assert.AreEqual(1, actionBuffers[42].DiscreteActions[1]);
+
+            Assert.AreEqual(0, actionBuffers[1337].DiscreteActions[0]);
+            Assert.AreEqual(0, actionBuffers[1337].DiscreteActions[1]);
+        }
+    }
+
     public class DiscreteActionWithMultiNomialOutputApplierTest
     {
         [Test]
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
index bd478be201..da4a6840ad 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
@@ -12,13 +12,14 @@ class TestAgent : Agent
         {
         }
 
-        [Test]
-        public void Construction()
+        [TestCase(2)]
+        [TestCase(3)]
+        public void Construction(int apiVersion)
         {
             var actionSpec = new ActionSpec();
             var alloc = new TensorCachingAllocator();
             var mem = new Dictionary<int, List<float>>();
-            var tensorGenerator = new TensorApplier(2, actionSpec, 0, alloc, mem);
+            var tensorGenerator = new TensorApplier(apiVersion, actionSpec, 0, alloc, mem);
             Assert.IsNotNull(tensorGenerator);
             alloc.Dispose();
         }
@@ -52,7 +53,7 @@ public void ApplyContinuousActionOutput()
         }
 
         [Test]
-        public void ApplyDiscreteActionOutput()
+        public void ApplyDiscreteActionOutputLegacy()
         {
             var actionSpec = ActionSpec.MakeDiscrete(2, 3);
             var inputTensor = new TensorProxy()
@@ -82,7 +83,37 @@ public void ApplyDiscreteActionOutput()
         }
 
         [Test]
-        public void ApplyHybridActionOutput()
+        public void ApplyDiscreteActionOutput()
+        {
+            var actionSpec = ActionSpec.MakeDiscrete(2, 3);
+            var inputTensor = new TensorProxy()
+            {
+                shape = new long[] { 2, 2 },
+                data = new Tensor(
+                    2,
+                    2,
+                    new[] { 1f, 1f, 1f, 2f }),
+            };
+            var alloc = new TensorCachingAllocator();
+            var applier = new DiscreteActionOutputApplier(actionSpec, 0, alloc);
+
+            var agentIds = new List<int>() { 0, 1 };
+            // Dictionary from AgentId to Action
+            var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
+
+
+            applier.Apply(inputTensor, agentIds, actionDict);
+
+            Assert.AreEqual(actionDict[0].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[0].DiscreteActions[1], 1);
+
+            Assert.AreEqual(actionDict[1].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[1].DiscreteActions[1], 2);
+            alloc.Dispose();
+        }
+
+        [Test]
+        public void ApplyHybridActionOutputLegacy()
         {
             var actionSpec = new ActionSpec(3, new[] { 2, 3 });
             var continuousInputTensor = new TensorProxy()
@@ -107,6 +138,49 @@ public void ApplyHybridActionOutput()
             var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
 
 
+            continuousApplier.Apply(continuousInputTensor, agentIds, actionDict);
+            discreteApplier.Apply(discreteInputTensor, agentIds, actionDict);
+
+            Assert.AreEqual(actionDict[0].ContinuousActions[0], 1);
+            Assert.AreEqual(actionDict[0].ContinuousActions[1], 2);
+            Assert.AreEqual(actionDict[0].ContinuousActions[2], 3);
+            Assert.AreEqual(actionDict[0].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[0].DiscreteActions[1], 1);
+
+            Assert.AreEqual(actionDict[1].ContinuousActions[0], 4);
+            Assert.AreEqual(actionDict[1].ContinuousActions[1], 5);
+            Assert.AreEqual(actionDict[1].ContinuousActions[2], 6);
+            Assert.AreEqual(actionDict[1].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[1].DiscreteActions[1], 2);
+            alloc.Dispose();
+        }
+
+        [Test]
+        public void ApplyHybridActionOutput()
+        {
+            var actionSpec = new ActionSpec(3, new[] { 2, 3 });
+            var continuousInputTensor = new TensorProxy()
+            {
+                shape = new long[] { 2, 3 },
+                data = new Tensor(2, 3, new float[] { 1, 2, 3, 4, 5, 6 })
+            };
+            var discreteInputTensor = new TensorProxy()
+            {
+                shape = new long[] { 2, 2 },
+                data = new Tensor(
+                    2,
+                    2,
+                    new[] { 1f, 1f, 1f, 2f }),
+            };
+            var continuousApplier = new ContinuousActionOutputApplier(actionSpec);
+            var alloc = new TensorCachingAllocator();
+            var discreteApplier = new DiscreteActionOutputApplier(actionSpec, 0, alloc);
+
+            var agentIds = new List<int>() { 0, 1 };
+            // Dictionary from AgentId to Action
+            var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
+
+
             continuousApplier.Apply(continuousInputTensor, agentIds, actionDict);
             discreteApplier.Apply(discreteInputTensor, agentIds, actionDict);
 
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
index 4904d2bea2..792b4b487d 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
@@ -63,12 +63,13 @@ static List<TestAgent> GetFakeAgents(ObservableAttributeOptions observableAttrib
             return agents;
         }
 
-        [Test]
-        public void Construction()
+        [TestCase(2)]
+        [TestCase(3)]
+        public void Construction(int apiVersion)
         {
             var alloc = new TensorCachingAllocator();
             var mem = new Dictionary<int, List<float>>();
-            var tensorGenerator = new TensorGenerator(0, 2, alloc, mem);
+            var tensorGenerator = new TensorGenerator(apiVersion, 0, alloc, mem);
             Assert.IsNotNull(tensorGenerator);
             alloc.Dispose();
         }

From abfb2689adba3b0924b95936aea919eaf2e668f1 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 10:34:59 -0800
Subject: [PATCH 06/24] Use an enum for valid versions and use GetVersion on
 model directly

---
 .../Inference/BarracudaModelParamLoader.cs    | 20 ++++++++++++-------
 .../Runtime/Inference/ModelRunner.cs          |  5 ++---
 .../Runtime/Inference/TensorApplier.cs        |  7 +++----
 .../Runtime/Inference/TensorGenerator.cs      |  9 ++++-----
 .../EditModeTestInternalBrainTensorApplier.cs |  4 +---
 ...ditModeTestInternalBrainTensorGenerator.cs |  4 +---
 6 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index c63f3b3797..ace6260859 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -16,6 +16,14 @@ namespace Unity.MLAgents.Inference
     internal class BarracudaModelParamLoader
     {
 
+        internal enum ModelApiVersion
+        {
+            MLAgents1_0 = 2,
+            MLAgents2_0 = 3,
+            MinSupportedVersion = MLAgents1_0,
+            MaxSupportedVersion = MLAgents2_0
+        }
+
         internal enum CheckType
         {
             Info = 0,
@@ -29,9 +37,6 @@ internal class FailedCheck
             public string Message;
         }
 
-        public const long k_ApiVersion = 3;
-        public const long k_ApiVersionLegacy = 2;
-
         /// <summary>
         /// Factory for the ModelParamLoader : Creates a ModelParamLoader and runs the checks
         /// on it.
@@ -87,7 +92,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                     });
                 return failedModelChecks;
             }
-            if (modelApiVersion != k_ApiVersion && modelApiVersion != k_ApiVersionLegacy)
+            if (modelApiVersion < (int)ModelApiVersion.MinSupportedVersion || modelApiVersion > (int)ModelApiVersion.MaxSupportedVersion)
             {
                 failedModelChecks.Add(
                     new FailedCheck
@@ -95,7 +100,8 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                         CheckType = CheckType.Error,
                         Message =
                     $"Version of the trainer the model was trained with ({modelApiVersion}) " +
-                    $"is not compatible with the Brain's version ({k_ApiVersion})."
+                    $"is not compatible with the current range of supported versions:  " +
+                    $"({(int)ModelApiVersion.MinSupportedVersion} to {(int)ModelApiVersion.MaxSupportedVersion})."
                     });
                 return failedModelChecks;
             }
@@ -114,7 +120,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
             }
 
             var modelVersion = model.GetVersion();
-            if (modelVersion == k_ApiVersionLegacy)
+            if (modelVersion == (int)ModelApiVersion.MLAgents1_0)
             {
                 failedModelChecks.AddRange(
                     CheckInputTensorPresenceLegacy(model, brainParameters, memorySize, sensors)
@@ -123,7 +129,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                     CheckInputTensorShapeLegacy(model, brainParameters, sensors, observableAttributeTotalSize)
                 );
             }
-            if (modelVersion == k_ApiVersion)
+            if (modelVersion == (int)ModelApiVersion.MLAgents2_0)
             {
                 failedModelChecks.AddRange(
                     CheckInputTensorPresence(model, brainParameters, memorySize, sensors)
diff --git a/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs b/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
index 4eecf84862..12a8188c52 100644
--- a/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
+++ b/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
@@ -97,11 +97,10 @@ public ModelRunner(
 
             m_InferenceInputs = barracudaModel.GetInputTensors();
             m_OutputNames = barracudaModel.GetOutputNames();
-            var apiVersion = barracudaModel.GetVersion();
             m_TensorGenerator = new TensorGenerator(
-                apiVersion, seed, m_TensorAllocator, m_Memories, barracudaModel);
+                seed, m_TensorAllocator, m_Memories, barracudaModel);
             m_TensorApplier = new TensorApplier(
-                apiVersion, actionSpec, seed, m_TensorAllocator, m_Memories, barracudaModel);
+                actionSpec, seed, m_TensorAllocator, m_Memories, barracudaModel);
             m_InputsByName = new Dictionary<string, Tensor>();
             m_InferenceOutputs = new List<TensorProxy>();
         }
diff --git a/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs b/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
index de74fb5b91..d61bd1a361 100644
--- a/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
+++ b/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
@@ -39,14 +39,12 @@ public interface IApplier
         /// <summary>
         /// Returns a new TensorAppliers object.
         /// </summary>
-        /// <param name="modelVersion"> The API version of the model.</param>
         /// <param name="actionSpec"> Description of the actions for the Agent.</param>
         /// <param name="seed"> The seed the Appliers will be initialized with.</param>
         /// <param name="allocator"> Tensor allocator</param>
         /// <param name="memories">Dictionary of AgentInfo.id to memory used to pass to the inference model.</param>
         /// <param name="barracudaModel"></param>
         public TensorApplier(
-            int modelVersion,
             ActionSpec actionSpec,
             int seed,
             ITensorAllocator allocator,
@@ -72,11 +70,12 @@ public TensorApplier(
             if (actionSpec.NumDiscreteActions > 0)
             {
                 var tensorName = model.DiscreteOutputName();
-                if (modelVersion == BarracudaModelParamLoader.k_ApiVersionLegacy)
+                var modelVersion = model.GetVersion();
+                if (modelVersion == (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents1_0)
                 {
                     m_Dict[tensorName] = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, seed, allocator);
                 }
-                if (modelVersion == BarracudaModelParamLoader.k_ApiVersion)
+                if (modelVersion == (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
                 {
                     m_Dict[tensorName] = new DiscreteActionOutputApplier(actionSpec, seed, allocator);
                 }
diff --git a/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs b/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs
index e0f20a319b..1b1fe0c1de 100644
--- a/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs
+++ b/com.unity.ml-agents/Runtime/Inference/TensorGenerator.cs
@@ -40,19 +40,16 @@ void Generate(
         /// <summary>
         /// Returns a new TensorGenerators object.
         /// </summary>
-        /// <param name="modelVersion"> The API version of the model.</param>
         /// <param name="seed"> The seed the Generators will be initialized with.</param>
         /// <param name="allocator"> Tensor allocator.</param>
         /// <param name="memories">Dictionary of AgentInfo.id to memory for use in the inference model.</param>
         /// <param name="barracudaModel"></param>
         public TensorGenerator(
-            int modelVersion,
             int seed,
             ITensorAllocator allocator,
             Dictionary<int, List<float>> memories,
             object barracudaModel = null)
         {
-            m_ApiVersion = modelVersion;
             // If model is null, no inference to run and exception is thrown before reaching here.
             if (barracudaModel == null)
             {
@@ -60,6 +57,8 @@ public TensorGenerator(
             }
             var model = (Model)barracudaModel;
 
+            m_ApiVersion = model.GetVersion();
+
             // Generator for Inputs
             m_Dict[TensorNames.BatchSizePlaceholder] =
                 new BatchSizeGenerator(allocator);
@@ -97,7 +96,7 @@ public TensorGenerator(
 
         public void InitializeObservations(List<ISensor> sensors, ITensorAllocator allocator)
         {
-            if (m_ApiVersion == BarracudaModelParamLoader.k_ApiVersionLegacy)
+            if (m_ApiVersion == (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents1_0)
             {
                 // Loop through the sensors on a representative agent.
                 // All vector observations use a shared ObservationGenerator since they are concatenated.
@@ -142,7 +141,7 @@ public void InitializeObservations(List<ISensor> sensors, ITensorAllocator alloc
                     m_Dict[obsGenName] = obsGen;
                 }
             }
-            if (m_ApiVersion == BarracudaModelParamLoader.k_ApiVersion)
+            if (m_ApiVersion == (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
             {
                 for (var sensorIndex = 0; sensorIndex < sensors.Count; sensorIndex++)
                 {
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
index da4a6840ad..42b1cac786 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
@@ -12,14 +12,12 @@ class TestAgent : Agent
         {
         }
 
-        [TestCase(2)]
-        [TestCase(3)]
         public void Construction(int apiVersion)
         {
             var actionSpec = new ActionSpec();
             var alloc = new TensorCachingAllocator();
             var mem = new Dictionary<int, List<float>>();
-            var tensorGenerator = new TensorApplier(apiVersion, actionSpec, 0, alloc, mem);
+            var tensorGenerator = new TensorApplier(actionSpec, 0, alloc, mem);
             Assert.IsNotNull(tensorGenerator);
             alloc.Dispose();
         }
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
index 792b4b487d..70cc6c4f54 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
@@ -63,13 +63,11 @@ static List<TestAgent> GetFakeAgents(ObservableAttributeOptions observableAttrib
             return agents;
         }
 
-        [TestCase(2)]
-        [TestCase(3)]
         public void Construction(int apiVersion)
         {
             var alloc = new TensorCachingAllocator();
             var mem = new Dictionary<int, List<float>>();
-            var tensorGenerator = new TensorGenerator(apiVersion, 0, alloc, mem);
+            var tensorGenerator = new TensorGenerator(0, alloc, mem);
             Assert.IsNotNull(tensorGenerator);
             alloc.Dispose();
         }

From 93e18b5fa5977e9218c8300bc15080ea0f542f9e Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 10:43:43 -0800
Subject: [PATCH 07/24] Maybe the model export version a static constant in
 Python

---
 ml-agents/mlagents/trainers/torch/networks.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/ml-agents/mlagents/trainers/torch/networks.py b/ml-agents/mlagents/trainers/torch/networks.py
index 1c626455e5..2f9717c7a2 100644
--- a/ml-agents/mlagents/trainers/torch/networks.py
+++ b/ml-agents/mlagents/trainers/torch/networks.py
@@ -299,6 +299,8 @@ def forward(
 
 
 class SimpleActor(nn.Module, Actor):
+    MODEL_EXPORT_VERSION = 3
+
     def __init__(
         self,
         observation_specs: List[ObservationSpec],
@@ -310,7 +312,7 @@ def __init__(
         super().__init__()
         self.action_spec = action_spec
         self.version_number = torch.nn.Parameter(
-            torch.Tensor([3.0]), requires_grad=False
+            torch.Tensor([self.MODEL_EXPORT_VERSION]), requires_grad=False
         )
         self.is_continuous_int_deprecated = torch.nn.Parameter(
             torch.Tensor([int(self.action_spec.is_continuous())]), requires_grad=False

From 4afde1041a52f0412997231604d8f79b33f6f0c1 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 11:04:08 -0800
Subject: [PATCH 08/24] Use static constructor for FailedCheck

---
 .../Inference/BarracudaModelExtensions.cs     |  58 +----
 .../Inference/BarracudaModelParamLoader.cs    | 206 +++++-------------
 2 files changed, 67 insertions(+), 197 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index 17e7103d83..4bbf0628b7 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -318,13 +318,7 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
             var modelApiVersionTensor = model.GetTensorByName(TensorNames.VersionNumber);
             if (modelApiVersionTensor == null)
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    $"Required constant \"{TensorNames.VersionNumber}\" was not found in the model file."
-                    });
+                failedModelChecks.Add(FailedCheck.Warning($"Required constant \"{TensorNames.VersionNumber}\" was not found in the model file."));
                 return false;
             }
 
@@ -332,13 +326,7 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
             var memorySizeTensor = model.GetTensorByName(TensorNames.MemorySize);
             if (memorySizeTensor == null)
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    $"Required constant \"{TensorNames.MemorySize}\" was not found in the model file."
-                    });
+                failedModelChecks.Add(FailedCheck.Warning($"Required constant \"{TensorNames.MemorySize}\" was not found in the model file."));
                 return false;
             }
 
@@ -347,13 +335,7 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
                 !model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
                 !model.outputs.Contains(TensorNames.DiscreteActionOutput))
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    "The model does not contain any Action Output Node."
-                    });
+                failedModelChecks.Add(FailedCheck.Warning("The model does not contain any Action Output Node."));
                 return false;
             }
 
@@ -362,24 +344,14 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
             {
                 if (model.GetTensorByName(TensorNames.ActionOutputShapeDeprecated) == null)
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model does not contain any Action Output Shape Node."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain any Action Output Shape Node."));
                     return false;
                 }
                 if (model.GetTensorByName(TensorNames.IsContinuousControlDeprecated) == null)
                 {
-                    failedModelChecks.Add(new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    $"Required constant \"{TensorNames.IsContinuousControlDeprecated}\" was not found in the model file. " +
+                    failedModelChecks.Add(FailedCheck.Warning($"Required constant \"{TensorNames.IsContinuousControlDeprecated}\" was not found in the model file. " +
                         "This is only required for model that uses a deprecated model format."
-                    });
+                    ));
                     return false;
                 }
             }
@@ -388,25 +360,15 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
                 if (model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
                     model.GetTensorByName(TensorNames.ContinuousActionOutputShape) == null)
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model uses continuous action but does not contain Continuous Action Output Shape Node."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model uses continuous action but does not contain Continuous Action Output Shape Node."
+                        ));
                     return false;
                 }
                 if (model.outputs.Contains(TensorNames.DiscreteActionOutput) &&
                     model.GetTensorByName(TensorNames.DiscreteActionOutputShape) == null)
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                    "The model uses discrete action but does not contain Discrete Action Output Shape Node."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model uses discrete action but does not contain Discrete Action Output Shape Node."
+                        ));
                     return false;
                 }
             }
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index ace6260859..4a5d02721b 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -35,6 +35,18 @@ internal class FailedCheck
         {
             public CheckType CheckType;
             public string Message;
+            public static FailedCheck Info(string message)
+            {
+                return new FailedCheck { CheckType = CheckType.Info, Message = message };
+            }
+            public static FailedCheck Warning(string message)
+            {
+                return new FailedCheck { CheckType = CheckType.Warning, Message = message };
+            }
+            public static FailedCheck Error(string message)
+            {
+                return new FailedCheck { CheckType = CheckType.Error, Message = message };
+            }
         }
 
         /// <summary>
@@ -69,7 +81,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                 {
                     errorMsg += "(But can still train)";
                 }
-                failedModelChecks.Add(new FailedCheck { CheckType = CheckType.Warning, Message = errorMsg });
+                failedModelChecks.Add(FailedCheck.Warning(errorMsg));
                 return failedModelChecks;
             }
 
@@ -82,40 +94,25 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
             var modelApiVersion = (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
             if (modelApiVersion == -1)
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    "Model was not trained using the right version of ML-Agents. " +
+                failedModelChecks.Add(FailedCheck.Warning("Model was not trained using the right version of ML-Agents. " +
                     "Cannot use this model."
-                    });
+                    ));
                 return failedModelChecks;
             }
             if (modelApiVersion < (int)ModelApiVersion.MinSupportedVersion || modelApiVersion > (int)ModelApiVersion.MaxSupportedVersion)
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Error,
-                        Message =
-                    $"Version of the trainer the model was trained with ({modelApiVersion}) " +
+                failedModelChecks.Add(FailedCheck.Warning($"Version of the trainer the model was trained with ({modelApiVersion}) " +
                     $"is not compatible with the current range of supported versions:  " +
                     $"({(int)ModelApiVersion.MinSupportedVersion} to {(int)ModelApiVersion.MaxSupportedVersion})."
-                    });
+                    ));
                 return failedModelChecks;
             }
 
             var memorySize = (int)model.GetTensorByName(TensorNames.MemorySize)[0];
             if (memorySize == -1)
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    $"Missing node in the model provided : {TensorNames.MemorySize}"
-                    });
+                failedModelChecks.Add(FailedCheck.Warning($"Missing node in the model provided : {TensorNames.MemorySize}"
+                    ));
                 return failedModelChecks;
             }
 
@@ -182,14 +179,9 @@ ISensor[] sensors
             if ((brainParameters.VectorObservationSize != 0) &&
                 (!tensorsNames.Contains(TensorNames.VectorObservationPlaceholder)))
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    "The model does not contain a Vector Observation Placeholder Input. " +
+                failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Vector Observation Placeholder Input. " +
                     "You must set the Vector Observation Space Size to 0."
-                    });
+                    ));
             }
 
             // If there are not enough Visual Observation Input compared to what the
@@ -203,14 +195,9 @@ ISensor[] sensors
                     if (!tensorsNames.Contains(
                         TensorNames.GetVisualObservationName(visObsIndex)))
                     {
-                        failedModelChecks.Add(
-                            new FailedCheck
-                            {
-                                CheckType = CheckType.Warning,
-                                Message =
-                            "The model does not contain a Visual Observation Placeholder Input " +
+                        failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Visual Observation Placeholder Input " +
                             $"for sensor component {visObsIndex} ({sensor.GetType().Name})."
-                            });
+                            ));
                     }
                     visObsIndex++;
                 }
@@ -219,14 +206,9 @@ ISensor[] sensors
                     if (!tensorsNames.Contains(
                         TensorNames.GetObservationName(sensorIndex)))
                     {
-                        failedModelChecks.Add(
-                            new FailedCheck
-                            {
-                                CheckType = CheckType.Warning,
-                                Message =
-                            "The model does not contain an Observation Placeholder Input " +
+                        failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Observation Placeholder Input " +
                             $"for sensor component {sensorIndex} ({sensor.GetType().Name})."
-                            });
+                            ));
                     }
                 }
 
@@ -236,15 +218,9 @@ ISensor[] sensors
             // Check if there's not enough visual sensors (too many would be handled above)
             if (expectedVisualObs > visObsIndex)
             {
-                failedModelChecks.Add(
-                    new FailedCheck
-                    {
-                        CheckType = CheckType.Warning,
-                        Message =
-                    $"The model expects {expectedVisualObs} visual inputs," +
+                failedModelChecks.Add(FailedCheck.Warning($"The model expects {expectedVisualObs} visual inputs," +
                     $" but only found {visObsIndex} visual sensors."
-                    }
-                );
+                ));
             }
 
             // If the model has a non-negative memory size but requires a recurrent input
@@ -253,13 +229,8 @@ ISensor[] sensors
                 if (!tensorsNames.Any(x => x.EndsWith("_h")) ||
                     !tensorsNames.Any(x => x.EndsWith("_c")))
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model does not contain a Recurrent Input Node but has memory_size."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size."
+                        ));
                 }
             }
 
@@ -268,13 +239,8 @@ ISensor[] sensors
             {
                 if (!tensorsNames.Contains(TensorNames.ActionMaskPlaceholder))
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model does not contain an Action Mask but is using Discrete Control."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Action Mask but is using Discrete Control."
+                        ));
                 }
             }
             return failedModelChecks;
@@ -312,14 +278,9 @@ ISensor[] sensors
                     TensorNames.GetObservationName(sensorIndex)))
                 {
                     var sensor = sensors[sensorIndex];
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model does not contain an Observation Placeholder Input " +
+                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Observation Placeholder Input " +
                         $"for sensor component {sensorIndex} ({sensor.GetType().Name})."
-                        });
+                        ));
                 }
             }
 
@@ -329,13 +290,8 @@ ISensor[] sensors
                 if (!tensorsNames.Any(x => x.EndsWith("_h")) ||
                     !tensorsNames.Any(x => x.EndsWith("_c")))
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model does not contain a Recurrent Input Node but has memory_size."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size."
+                        ));
                 }
             }
 
@@ -344,13 +300,8 @@ ISensor[] sensors
             {
                 if (!tensorsNames.Contains(TensorNames.ActionMaskPlaceholder))
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model does not contain an Action Mask but is using Discrete Control."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Action Mask but is using Discrete Control."
+                        ));
                 }
             }
             return failedModelChecks;
@@ -379,13 +330,8 @@ static IEnumerable<FailedCheck> CheckOutputTensorPresence(Model model, int memor
                 if (!memOutputs.Any(x => x.EndsWith("_h")) ||
                     !memOutputs.Any(x => x.EndsWith("_c")))
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "The model does not contain a Recurrent Output Node but has memory_size."
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Recurrent Output Node but has memory_size."
+                        ));
                 }
             }
             return failedModelChecks;
@@ -412,14 +358,10 @@ static FailedCheck CheckVisualObsShape(
             var pixelT = tensorProxy.Channels;
             if ((widthBp != widthT) || (heightBp != heightT) || (pixelBp != pixelT))
             {
-                return new FailedCheck
-                {
-                    CheckType = CheckType.Warning,
-                    Message =
-                    $"The visual Observation of the model does not match. " +
+                return FailedCheck.Warning($"The visual Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{widthBp}x{heightBp}x{pixelBp}] but " +
                     $"was expecting [?x{widthT}x{heightT}x{pixelT}]."
-                };
+                );
             }
             return null;
         }
@@ -449,14 +391,10 @@ static FailedCheck CheckRankTwoObsShape(
                 {
                     proxyDimStr = $"[?x{dim3T}x{dim2T}x{dim1T}]";
                 }
-                return new FailedCheck
-                {
-                    CheckType = CheckType.Warning,
-                    Message =
-                    $"An Observation of the model does not match. " +
+                return FailedCheck.Warning($"An Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{dim1Bp}x{dim2Bp}] but " +
                     $"was expecting {proxyDimStr}."
-                };
+                );
             }
             return null;
         }
@@ -489,14 +427,10 @@ static FailedCheck CheckRankOneObsShape(
                 {
                     proxyDimStr = $"[?x{dim3T}x{dim2T}x{dim1T}]";
                 }
-                return new FailedCheck
-                {
-                    CheckType = CheckType.Warning,
-                    Message =
-                    $"An Observation of the model does not match. " +
+                return FailedCheck.Warning($"An Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{dim1Bp}] but " +
                     $"was expecting {proxyDimStr}."
-                };
+                );
             }
             return null;
         }
@@ -560,13 +494,8 @@ static IEnumerable<FailedCheck> CheckInputTensorShapeLegacy(
                 {
                     if (!tensor.name.Contains("visual_observation"))
                     {
-                        failedModelChecks.Add(
-                            new FailedCheck
-                            {
-                                CheckType = CheckType.Warning,
-                                Message =
-                            "Model requires an unknown input named : " + tensor.name
-                            });
+                        failedModelChecks.Add(FailedCheck.Warning("Model requires an unknown input named : " + tensor.name
+                            ));
                     }
                 }
                 else
@@ -633,16 +562,12 @@ static FailedCheck CheckVectorObsShapeLegacy(
                 }
 
                 sensorSizes += "]";
-                return new FailedCheck
-                {
-                    CheckType = CheckType.Warning,
-                    Message =
-                $"Vector Observation Size of the model does not match. Was expecting {totalVecObsSizeT} " +
+                return FailedCheck.Warning($"Vector Observation Size of the model does not match. Was expecting {totalVecObsSizeT} " +
                     $"but received: \n" +
                     $"Vector observations: {vecObsSizeBp} x {numStackedVector}\n" +
                     $"Total [Observable] attributes: {observableAttributeTotalSize}\n" +
                     $"Sensor sizes: {sensorSizes}."
-                };
+                );
             }
             return null;
         }
@@ -707,13 +632,8 @@ static IEnumerable<FailedCheck> CheckInputTensorShape(
             {
                 if (!tensorTester.ContainsKey(tensor.name))
                 {
-                    failedModelChecks.Add(
-                        new FailedCheck
-                        {
-                            CheckType = CheckType.Warning,
-                            Message =
-                        "Model requires an unknown input named : " + tensor.name
-                        });
+                    failedModelChecks.Add(FailedCheck.Warning("Model requires an unknown input named : " + tensor.name
+                    ));
                 }
                 else
                 {
@@ -748,13 +668,9 @@ static FailedCheck CheckPreviousActionShape(
             var numberActionsT = tensorProxy.shape[tensorProxy.shape.Length - 1];
             if (numberActionsBp != numberActionsT)
             {
-                return new FailedCheck
-                {
-                    CheckType = CheckType.Warning,
-                    Message =
-                "Previous Action Size of the model does not match. " +
+                return FailedCheck.Warning("Previous Action Size of the model does not match. " +
                     $"Received {numberActionsBp} but was expecting {numberActionsT}."
-                };
+                );
             }
             return null;
         }
@@ -826,13 +742,9 @@ static FailedCheck CheckDiscreteActionOutputShape(
 
             if (modelSumDiscreteBranchSizes != sumOfDiscreteBranchSizes)
             {
-                return new FailedCheck
-                {
-                    CheckType = CheckType.Warning,
-                    Message =
-                "Discrete Action Size of the model does not match. The BrainParameters expect " +
+                return FailedCheck.Warning("Discrete Action Size of the model does not match. The BrainParameters expect " +
                     $"{sumOfDiscreteBranchSizes} but the model contains {modelSumDiscreteBranchSizes}."
-                };
+                );
             }
             return null;
         }
@@ -863,13 +775,9 @@ static FailedCheck CheckContinuousActionOutputShape(
 
             if (modelContinuousActionSize != numContinuousActions)
             {
-                return new FailedCheck
-                {
-                    CheckType = CheckType.Warning,
-                    Message =
-                "Continuous Action Size of the model does not match. The BrainParameters and ActuatorComponents expect " +
+                return FailedCheck.Warning("Continuous Action Size of the model does not match. The BrainParameters and ActuatorComponents expect " +
                     $"{numContinuousActions} but the model contains {modelContinuousActionSize}."
-                };
+                );
             }
             return null;
         }

From 393f14683e367efb7131ed5f7b9e22d7eb276e07 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 11:09:07 -0800
Subject: [PATCH 09/24] Use static constructor for FailedCheck

---
 .../Inference/BarracudaModelExtensions.cs     | 34 +++++---
 .../Inference/BarracudaModelParamLoader.cs    | 85 +++++++++++--------
 2 files changed, 73 insertions(+), 46 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index 4bbf0628b7..1a6342a1fb 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -318,7 +318,9 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
             var modelApiVersionTensor = model.GetTensorByName(TensorNames.VersionNumber);
             if (modelApiVersionTensor == null)
             {
-                failedModelChecks.Add(FailedCheck.Warning($"Required constant \"{TensorNames.VersionNumber}\" was not found in the model file."));
+                failedModelChecks.Add(
+                    FailedCheck.Warning($"Required constant \"{TensorNames.VersionNumber}\" was not found in the model file.")
+                    );
                 return false;
             }
 
@@ -326,7 +328,9 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
             var memorySizeTensor = model.GetTensorByName(TensorNames.MemorySize);
             if (memorySizeTensor == null)
             {
-                failedModelChecks.Add(FailedCheck.Warning($"Required constant \"{TensorNames.MemorySize}\" was not found in the model file."));
+                failedModelChecks.Add(
+                    FailedCheck.Warning($"Required constant \"{TensorNames.MemorySize}\" was not found in the model file.")
+                    );
                 return false;
             }
 
@@ -335,7 +339,9 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
                 !model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
                 !model.outputs.Contains(TensorNames.DiscreteActionOutput))
             {
-                failedModelChecks.Add(FailedCheck.Warning("The model does not contain any Action Output Node."));
+                failedModelChecks.Add(
+                    FailedCheck.Warning("The model does not contain any Action Output Node.")
+                    );
                 return false;
             }
 
@@ -344,14 +350,18 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
             {
                 if (model.GetTensorByName(TensorNames.ActionOutputShapeDeprecated) == null)
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain any Action Output Shape Node."));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model does not contain any Action Output Shape Node.")
+                        );
                     return false;
                 }
                 if (model.GetTensorByName(TensorNames.IsContinuousControlDeprecated) == null)
                 {
-                    failedModelChecks.Add(FailedCheck.Warning($"Required constant \"{TensorNames.IsContinuousControlDeprecated}\" was not found in the model file. " +
-                        "This is only required for model that uses a deprecated model format."
-                    ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning($"Required constant \"{TensorNames.IsContinuousControlDeprecated}\" was " +
+                        "not found in the model file. " +
+                        "This is only required for model that uses a deprecated model format.")
+                        );
                     return false;
                 }
             }
@@ -360,15 +370,17 @@ public static bool CheckExpectedTensors(this Model model, List<FailedCheck> fail
                 if (model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
                     model.GetTensorByName(TensorNames.ContinuousActionOutputShape) == null)
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model uses continuous action but does not contain Continuous Action Output Shape Node."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model uses continuous action but does not contain Continuous Action Output Shape Node.")
+                        );
                     return false;
                 }
                 if (model.outputs.Contains(TensorNames.DiscreteActionOutput) &&
                     model.GetTensorByName(TensorNames.DiscreteActionOutputShape) == null)
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model uses discrete action but does not contain Discrete Action Output Shape Node."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model uses discrete action but does not contain Discrete Action Output Shape Node.")
+                        );
                     return false;
                 }
             }
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 4a5d02721b..7cf2a04474 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -94,17 +94,19 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
             var modelApiVersion = (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
             if (modelApiVersion == -1)
             {
-                failedModelChecks.Add(FailedCheck.Warning("Model was not trained using the right version of ML-Agents. " +
-                    "Cannot use this model."
-                    ));
+                failedModelChecks.Add(
+                    FailedCheck.Warning("Model was not trained using the right version of ML-Agents. " +
+                    "Cannot use this model.")
+                    );
                 return failedModelChecks;
             }
             if (modelApiVersion < (int)ModelApiVersion.MinSupportedVersion || modelApiVersion > (int)ModelApiVersion.MaxSupportedVersion)
             {
-                failedModelChecks.Add(FailedCheck.Warning($"Version of the trainer the model was trained with ({modelApiVersion}) " +
+                failedModelChecks.Add(
+                    FailedCheck.Warning($"Version of the trainer the model was trained with ({modelApiVersion}) " +
                     $"is not compatible with the current range of supported versions:  " +
-                    $"({(int)ModelApiVersion.MinSupportedVersion} to {(int)ModelApiVersion.MaxSupportedVersion})."
-                    ));
+                    $"({(int)ModelApiVersion.MinSupportedVersion} to {(int)ModelApiVersion.MaxSupportedVersion}).")
+                    );
                 return failedModelChecks;
             }
 
@@ -179,9 +181,10 @@ ISensor[] sensors
             if ((brainParameters.VectorObservationSize != 0) &&
                 (!tensorsNames.Contains(TensorNames.VectorObservationPlaceholder)))
             {
-                failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Vector Observation Placeholder Input. " +
-                    "You must set the Vector Observation Space Size to 0."
-                    ));
+                failedModelChecks.Add(
+                    FailedCheck.Warning("The model does not contain a Vector Observation Placeholder Input. " +
+                    "You must set the Vector Observation Space Size to 0.")
+                    );
             }
 
             // If there are not enough Visual Observation Input compared to what the
@@ -195,9 +198,10 @@ ISensor[] sensors
                     if (!tensorsNames.Contains(
                         TensorNames.GetVisualObservationName(visObsIndex)))
                     {
-                        failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Visual Observation Placeholder Input " +
-                            $"for sensor component {visObsIndex} ({sensor.GetType().Name})."
-                            ));
+                        failedModelChecks.Add(
+                            FailedCheck.Warning("The model does not contain a Visual Observation Placeholder Input " +
+                            $"for sensor component {visObsIndex} ({sensor.GetType().Name}).")
+                            );
                     }
                     visObsIndex++;
                 }
@@ -206,9 +210,10 @@ ISensor[] sensors
                     if (!tensorsNames.Contains(
                         TensorNames.GetObservationName(sensorIndex)))
                     {
-                        failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Observation Placeholder Input " +
-                            $"for sensor component {sensorIndex} ({sensor.GetType().Name})."
-                            ));
+                        failedModelChecks.Add(
+                            FailedCheck.Warning("The model does not contain an Observation Placeholder Input " +
+                            $"for sensor component {sensorIndex} ({sensor.GetType().Name}).")
+                            );
                     }
                 }
 
@@ -218,9 +223,10 @@ ISensor[] sensors
             // Check if there's not enough visual sensors (too many would be handled above)
             if (expectedVisualObs > visObsIndex)
             {
-                failedModelChecks.Add(FailedCheck.Warning($"The model expects {expectedVisualObs} visual inputs," +
-                    $" but only found {visObsIndex} visual sensors."
-                ));
+                failedModelChecks.Add(
+                    FailedCheck.Warning($"The model expects {expectedVisualObs} visual inputs," +
+                    $" but only found {visObsIndex} visual sensors.")
+                    );
             }
 
             // If the model has a non-negative memory size but requires a recurrent input
@@ -229,8 +235,9 @@ ISensor[] sensors
                 if (!tensorsNames.Any(x => x.EndsWith("_h")) ||
                     !tensorsNames.Any(x => x.EndsWith("_c")))
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size.")
+                    );
                 }
             }
 
@@ -239,8 +246,9 @@ ISensor[] sensors
             {
                 if (!tensorsNames.Contains(TensorNames.ActionMaskPlaceholder))
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Action Mask but is using Discrete Control."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model does not contain an Action Mask but is using Discrete Control.")
+                        );
                 }
             }
             return failedModelChecks;
@@ -278,9 +286,10 @@ ISensor[] sensors
                     TensorNames.GetObservationName(sensorIndex)))
                 {
                     var sensor = sensors[sensorIndex];
-                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Observation Placeholder Input " +
-                        $"for sensor component {sensorIndex} ({sensor.GetType().Name})."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model does not contain an Observation Placeholder Input " +
+                            $"for sensor component {sensorIndex} ({sensor.GetType().Name}).")
+                        );
                 }
             }
 
@@ -290,8 +299,9 @@ ISensor[] sensors
                 if (!tensorsNames.Any(x => x.EndsWith("_h")) ||
                     !tensorsNames.Any(x => x.EndsWith("_c")))
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model does not contain a Recurrent Input Node but has memory_size.")
+                        );
                 }
             }
 
@@ -300,8 +310,9 @@ ISensor[] sensors
             {
                 if (!tensorsNames.Contains(TensorNames.ActionMaskPlaceholder))
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain an Action Mask but is using Discrete Control."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model does not contain an Action Mask but is using Discrete Control.")
+                        );
                 }
             }
             return failedModelChecks;
@@ -330,8 +341,9 @@ static IEnumerable<FailedCheck> CheckOutputTensorPresence(Model model, int memor
                 if (!memOutputs.Any(x => x.EndsWith("_h")) ||
                     !memOutputs.Any(x => x.EndsWith("_c")))
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("The model does not contain a Recurrent Output Node but has memory_size."
-                        ));
+                    failedModelChecks.Add(
+                        FailedCheck.Warning("The model does not contain a Recurrent Output Node but has memory_size.")
+                        );
                 }
             }
             return failedModelChecks;
@@ -494,8 +506,9 @@ static IEnumerable<FailedCheck> CheckInputTensorShapeLegacy(
                 {
                     if (!tensor.name.Contains("visual_observation"))
                     {
-                        failedModelChecks.Add(FailedCheck.Warning("Model requires an unknown input named : " + tensor.name
-                            ));
+                        failedModelChecks.Add(
+                            FailedCheck.Warning("Model requires an unknown input named : " + tensor.name)
+                            );
                     }
                 }
                 else
@@ -562,7 +575,8 @@ static FailedCheck CheckVectorObsShapeLegacy(
                 }
 
                 sensorSizes += "]";
-                return FailedCheck.Warning($"Vector Observation Size of the model does not match. Was expecting {totalVecObsSizeT} " +
+                return FailedCheck.Warning(
+                    $"Vector Observation Size of the model does not match. Was expecting {totalVecObsSizeT} " +
                     $"but received: \n" +
                     $"Vector observations: {vecObsSizeBp} x {numStackedVector}\n" +
                     $"Total [Observable] attributes: {observableAttributeTotalSize}\n" +
@@ -775,7 +789,8 @@ static FailedCheck CheckContinuousActionOutputShape(
 
             if (modelContinuousActionSize != numContinuousActions)
             {
-                return FailedCheck.Warning("Continuous Action Size of the model does not match. The BrainParameters and ActuatorComponents expect " +
+                return FailedCheck.Warning(
+                    "Continuous Action Size of the model does not match. The BrainParameters and ActuatorComponents expect " +
                     $"{numContinuousActions} but the model contains {modelContinuousActionSize}."
                 );
             }

From d305b5c54948015045d15f9b1662e7e8a3fa9065 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 11:37:39 -0800
Subject: [PATCH 10/24] Modifying the docstrings

---
 .../Inference/BarracudaModelExtensions.cs      |  3 ++-
 .../Inference/BarracudaModelParamLoader.cs     | 18 +++++++++---------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index 1a6342a1fb..acc99ec342 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -49,7 +49,8 @@ public static string[] GetInputNames(this Model model)
         /// <returns>The api version of the model</returns>
         public static int GetVersion(this Model model)
         {
-            return (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
+            return 3;
+            // return (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
         }
 
         /// <summary>
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 7cf2a04474..19d758817b 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -63,7 +63,7 @@ public static FailedCheck Error(string message)
         /// <param name="actuatorComponents">Attached actuator components</param>
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
         /// <param name="behaviorType">BehaviorType or the Agent to check.</param>
-        /// <returns>The list the error messages of the checks that failed</returns>
+        /// <returns>A IEnumerable of the checks that failed</returns>
         public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters brainParameters,
             ISensor[] sensors, ActuatorComponent[] actuatorComponents,
             int observableAttributeTotalSize = 0,
@@ -81,7 +81,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                 {
                     errorMsg += "(But can still train)";
                 }
-                failedModelChecks.Add(FailedCheck.Warning(errorMsg));
+                failedModelChecks.Add(FailedCheck.Info(errorMsg));
                 return failedModelChecks;
             }
 
@@ -91,7 +91,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                 return failedModelChecks;
             }
 
-            var modelApiVersion = (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
+            var modelApiVersion = model.GetVersion();
             if (modelApiVersion == -1)
             {
                 failedModelChecks.Add(
@@ -165,7 +165,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
         /// </param>
         /// <param name="sensors">Array of attached sensor components</param>
         /// <returns>
-        /// A IEnumerable of string corresponding to the failed input presence checks.
+        /// A IEnumerable of the checks that failed
         /// </returns>
         static IEnumerable<FailedCheck> CheckInputTensorPresenceLegacy(
             Model model,
@@ -269,7 +269,7 @@ ISensor[] sensors
         /// </param>
         /// <param name="sensors">Array of attached sensor components</param>
         /// <returns>
-        /// A IEnumerable of string corresponding to the failed input presence checks.
+        /// A IEnumerable of the checks that failed
         /// </returns>
         static IEnumerable<FailedCheck> CheckInputTensorPresence(
             Model model,
@@ -327,7 +327,7 @@ ISensor[] sensors
         /// </param>
         /// <param name="memory">The memory size that the model is expecting/</param>
         /// <returns>
-        /// A IEnumerable of string corresponding to the failed output presence checks.
+        /// A IEnumerable of the checks that failed
         /// </returns>
         static IEnumerable<FailedCheck> CheckOutputTensorPresence(Model model, int memory)
         {
@@ -459,7 +459,7 @@ static FailedCheck CheckRankOneObsShape(
         /// </param>
         /// <param name="sensors">Attached sensors</param>
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
-        /// <returns>The list the error messages of the checks that failed</returns>
+        /// <returns>A IEnumerable of the checks that failed</returns>
         static IEnumerable<FailedCheck> CheckInputTensorShapeLegacy(
             Model model, BrainParameters brainParameters, ISensor[] sensors,
             int observableAttributeTotalSize)
@@ -599,7 +599,7 @@ static FailedCheck CheckVectorObsShapeLegacy(
         /// </param>
         /// <param name="sensors">Attached sensors</param>
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
-        /// <returns>The list the error messages of the checks that failed</returns>
+        /// <returns>A IEnumerable of the checks that failed</returns>
         static IEnumerable<FailedCheck> CheckInputTensorShape(
             Model model, BrainParameters brainParameters, ISensor[] sensors,
             int observableAttributeTotalSize)
@@ -701,7 +701,7 @@ static FailedCheck CheckPreviousActionShape(
         /// </param>
         /// <param name="actuatorComponents">Array of attached actuator components.</param>
         /// <returns>
-        /// A IEnumerable of string corresponding to the incompatible shapes between model
+        /// A IEnumerable of error messages corresponding to the incompatible shapes between model
         /// and BrainParameters.
         /// </returns>
         static IEnumerable<FailedCheck> CheckOutputTensorShape(

From 067933eab448dfbe43f08e15f52b0e4f019a4713 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 11:41:08 -0800
Subject: [PATCH 11/24] renaming LegacyDiscreteActionOutputApplier

---
 com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs          | 4 ++--
 com.unity.ml-agents/Runtime/Inference/TensorApplier.cs        | 2 +-
 .../Tests/Editor/DiscreteActionOutputApplierTest.cs           | 4 ++--
 .../Tests/Editor/EditModeTestInternalBrainTensorApplier.cs    | 4 ++--
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs b/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs
index 2daac37eeb..fe459ecc71 100644
--- a/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs
+++ b/com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs
@@ -90,7 +90,7 @@ public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int,
     /// The Applier for the Discrete Action output tensor. Uses multinomial to sample discrete
     /// actions from the logits contained in the tensor.
     /// </summary>
-    internal class DiscreteActionWithMultiNomialOutputApplier : TensorApplier.IApplier
+    internal class LegacyDiscreteActionOutputApplier : TensorApplier.IApplier
     {
         readonly int[] m_ActionSize;
         readonly Multinomial m_Multinomial;
@@ -99,7 +99,7 @@ internal class DiscreteActionWithMultiNomialOutputApplier : TensorApplier.IAppli
         readonly float[] m_CdfBuffer;
 
 
-        public DiscreteActionWithMultiNomialOutputApplier(ActionSpec actionSpec, int seed, ITensorAllocator allocator)
+        public LegacyDiscreteActionOutputApplier(ActionSpec actionSpec, int seed, ITensorAllocator allocator)
         {
             m_ActionSize = actionSpec.BranchSizes;
             m_Multinomial = new Multinomial(seed);
diff --git a/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs b/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
index d61bd1a361..8845a9da30 100644
--- a/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
+++ b/com.unity.ml-agents/Runtime/Inference/TensorApplier.cs
@@ -73,7 +73,7 @@ public TensorApplier(
                 var modelVersion = model.GetVersion();
                 if (modelVersion == (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents1_0)
                 {
-                    m_Dict[tensorName] = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, seed, allocator);
+                    m_Dict[tensorName] = new LegacyDiscreteActionOutputApplier(actionSpec, seed, allocator);
                 }
                 if (modelVersion == (int)BarracudaModelParamLoader.ModelApiVersion.MLAgents2_0)
                 {
diff --git a/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs b/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
index 44a35a1e0c..58e6c9a95b 100644
--- a/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
+++ b/com.unity.ml-agents/Tests/Editor/DiscreteActionOutputApplierTest.cs
@@ -45,7 +45,7 @@ public void TestDiscreteApply()
         }
     }
 
-    public class DiscreteActionWithMultiNomialOutputApplierTest
+    public class LegacyDiscreteActionOutputApplierTest
     {
         [Test]
         public void TestDiscreteApply()
@@ -69,7 +69,7 @@ public void TestDiscreteApply()
                 valueType = TensorProxy.TensorType.FloatingPoint
             };
 
-            var applier = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, 2020, null);
+            var applier = new LegacyDiscreteActionOutputApplier(actionSpec, 2020, null);
             var agentIds = new List<int> { 42, 1337 };
             var actionBuffers = new Dictionary<int, ActionBuffers>();
             actionBuffers[42] = new ActionBuffers(actionSpec);
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
index 42b1cac786..5f4a3cc326 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
@@ -63,7 +63,7 @@ public void ApplyDiscreteActionOutputLegacy()
                     new[] { 0.5f, 22.5f, 0.1f, 5f, 1f, 4f, 5f, 6f, 7f, 8f })
             };
             var alloc = new TensorCachingAllocator();
-            var applier = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, 0, alloc);
+            var applier = new LegacyDiscreteActionOutputApplier(actionSpec, 0, alloc);
 
             var agentIds = new List<int>() { 0, 1 };
             // Dictionary from AgentId to Action
@@ -129,7 +129,7 @@ public void ApplyHybridActionOutputLegacy()
             };
             var continuousApplier = new ContinuousActionOutputApplier(actionSpec);
             var alloc = new TensorCachingAllocator();
-            var discreteApplier = new DiscreteActionWithMultiNomialOutputApplier(actionSpec, 0, alloc);
+            var discreteApplier = new LegacyDiscreteActionOutputApplier(actionSpec, 0, alloc);
 
             var agentIds = new List<int>() { 0, 1 };
             // Dictionary from AgentId to Action

From 24af8abe3895247851810ca6ee61179361840240 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 11:42:31 -0800
Subject: [PATCH 12/24] removing testing code

---
 .../Runtime/Inference/BarracudaModelExtensions.cs              | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index acc99ec342..1a6342a1fb 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -49,8 +49,7 @@ public static string[] GetInputNames(this Model model)
         /// <returns>The api version of the model</returns>
         public static int GetVersion(this Model model)
         {
-            return 3;
-            // return (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
+            return (int)model.GetTensorByName(TensorNames.VersionNumber)[0];
         }
 
         /// <summary>

From 8a4caddcccdaad85eabd4a890ff566c201f0c954 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 11:56:48 -0800
Subject: [PATCH 13/24] better warning message

---
 .../Runtime/Inference/BarracudaModelParamLoader.cs   | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 19d758817b..3bbd1b7ee6 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -372,7 +372,7 @@ static FailedCheck CheckVisualObsShape(
             {
                 return FailedCheck.Warning($"The visual Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{widthBp}x{heightBp}x{pixelBp}] but " +
-                    $"was expecting [?x{widthT}x{heightT}x{pixelT}]."
+                    $"was expecting [?x{widthT}x{heightT}x{pixelT}] for the {sensor.GetName()} Sensor."
                 );
             }
             return null;
@@ -399,13 +399,13 @@ static FailedCheck CheckRankTwoObsShape(
             if ((dim1Bp != dim1T) || (dim2Bp != dim2T))
             {
                 var proxyDimStr = $"[?x{dim1T}x{dim2T}]";
-                if (dim3T > 0)
+                if (dim3T > 1)
                 {
                     proxyDimStr = $"[?x{dim3T}x{dim2T}x{dim1T}]";
                 }
                 return FailedCheck.Warning($"An Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{dim1Bp}x{dim2Bp}] but " +
-                    $"was expecting {proxyDimStr}."
+                    $"was expecting {proxyDimStr} for the {sensor.GetName()} Sensor."
                 );
             }
             return null;
@@ -431,17 +431,17 @@ static FailedCheck CheckRankOneObsShape(
             if ((dim1Bp != dim1T))
             {
                 var proxyDimStr = $"[?x{dim1T}]";
-                if (dim2T > 0)
+                if (dim2T > 1)
                 {
                     proxyDimStr = $"[?x{dim1T}x{dim2T}]";
                 }
-                if (dim3T > 0)
+                if (dim3T > 1)
                 {
                     proxyDimStr = $"[?x{dim3T}x{dim2T}x{dim1T}]";
                 }
                 return FailedCheck.Warning($"An Observation of the model does not match. " +
                     $"Received TensorProxy of shape [?x{dim1Bp}] but " +
-                    $"was expecting {proxyDimStr}."
+                    $"was expecting {proxyDimStr} for the {sensor.GetName()} Sensor."
                 );
             }
             return null;

From bb9ddd63f53991d81f389e966280f249c9a95a8f Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 12:08:04 -0800
Subject: [PATCH 14/24] Nest the CheckTypeEnum into the FailedCheck class

---
 .../Editor/BehaviorParametersEditor.cs        |  8 +++----
 .../Inference/BarracudaModelExtensions.cs     |  1 -
 .../Inference/BarracudaModelParamLoader.cs    | 21 +++++++++----------
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs b/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
index 5b003d9851..f2925341b0 100644
--- a/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
+++ b/com.unity.ml-agents/Editor/BehaviorParametersEditor.cs
@@ -5,7 +5,7 @@
 using Unity.MLAgents.Policies;
 using Unity.MLAgents.Sensors;
 using Unity.MLAgents.Sensors.Reflection;
-using CheckType = Unity.MLAgents.Inference.BarracudaModelParamLoader.CheckType;
+using CheckTypeEnum = Unity.MLAgents.Inference.BarracudaModelParamLoader.FailedCheck.CheckTypeEnum;
 
 namespace Unity.MLAgents.Editor
 {
@@ -150,13 +150,13 @@ void DisplayFailedModelChecks()
                     {
                         switch (check.CheckType)
                         {
-                            case CheckType.Info:
+                            case CheckTypeEnum.Info:
                                 EditorGUILayout.HelpBox(check.Message, MessageType.Info);
                                 break;
-                            case CheckType.Warning:
+                            case CheckTypeEnum.Warning:
                                 EditorGUILayout.HelpBox(check.Message, MessageType.Warning);
                                 break;
-                            case CheckType.Error:
+                            case CheckTypeEnum.Error:
                                 EditorGUILayout.HelpBox(check.Message, MessageType.Error);
                                 break;
                             default:
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index 1a6342a1fb..78926608dc 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -2,7 +2,6 @@
 using System.Linq;
 using Unity.Barracuda;
 using FailedCheck = Unity.MLAgents.Inference.BarracudaModelParamLoader.FailedCheck;
-using CheckType = Unity.MLAgents.Inference.BarracudaModelParamLoader.CheckType;
 
 namespace Unity.MLAgents.Inference
 {
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 3bbd1b7ee6..726a1c6b22 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -24,28 +24,27 @@ internal enum ModelApiVersion
             MaxSupportedVersion = MLAgents2_0
         }
 
-        internal enum CheckType
-        {
-            Info = 0,
-            Warning = 1,
-            Error = 2
-        }
-
         internal class FailedCheck
         {
-            public CheckType CheckType;
+            public enum CheckTypeEnum
+            {
+                Info = 0,
+                Warning = 1,
+                Error = 2
+            }
+            public CheckTypeEnum CheckType;
             public string Message;
             public static FailedCheck Info(string message)
             {
-                return new FailedCheck { CheckType = CheckType.Info, Message = message };
+                return new FailedCheck { CheckType = CheckTypeEnum.Info, Message = message };
             }
             public static FailedCheck Warning(string message)
             {
-                return new FailedCheck { CheckType = CheckType.Warning, Message = message };
+                return new FailedCheck { CheckType = CheckTypeEnum.Warning, Message = message };
             }
             public static FailedCheck Error(string message)
             {
-                return new FailedCheck { CheckType = CheckType.Error, Message = message };
+                return new FailedCheck { CheckType = CheckTypeEnum.Error, Message = message };
             }
         }
 

From dc4865b3edf67c6342c64420e2818d972e169b62 Mon Sep 17 00:00:00 2001
From: Vincent-Pierre BERGES <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 12:13:11 -0800
Subject: [PATCH 15/24] Update
 com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs

---
 .../Runtime/Inference/BarracudaModelParamLoader.cs               | 1 -
 1 file changed, 1 deletion(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 726a1c6b22..2e0f0c4ee0 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -8,7 +8,6 @@
 
 namespace Unity.MLAgents.Inference
 {
-
     /// <summary>
     /// Prepares the Tensors for the Learning Brain and exposes a list of failed checks if Model
     /// and BrainParameters are incompatible.

From 405b18a97547681b6cc946b508ecec5fea5887c7 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 12:16:23 -0800
Subject: [PATCH 16/24] Adding a line explaining that legacy tensor checks are
 for versions 1.X only

---
 .../Runtime/Inference/BarracudaModelParamLoader.cs          | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 2e0f0c4ee0..89d114bbb1 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -150,7 +150,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
 
         /// <summary>
         /// Generates failed checks that correspond to inputs expected by the model that are not
-        /// present in the BrainParameters.
+        /// present in the BrainParameters. Tests the models created with the API of version 1.X
         /// </summary>
         /// <param name="model">
         /// The Barracuda engine model for loading static parameters
@@ -447,7 +447,7 @@ static FailedCheck CheckRankOneObsShape(
 
         /// <summary>
         /// Generates failed checks that correspond to inputs shapes incompatibilities between
-        /// the model and the BrainParameters.
+        /// the model and the BrainParameters. Tests the models created with the API of version 1.X
         /// </summary>
         /// <param name="model">
         /// The Barracuda engine model for loading static parameters
@@ -524,7 +524,7 @@ static IEnumerable<FailedCheck> CheckInputTensorShapeLegacy(
 
         /// <summary>
         /// Checks that the shape of the Vector Observation input placeholder is the same in the
-        /// model and in the Brain Parameters.
+        /// model and in the Brain Parameters. Tests the models created with the API of version 1.X
         /// </summary>
         /// <param name="brainParameters">
         /// The BrainParameters that are used verify the compatibility with the InferenceEngine

From fb459a7a5fbc629a4f7c09471b98cc993cc2a96a Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 12:25:01 -0800
Subject: [PATCH 17/24] Modifying the changelog

---
 com.unity.ml-agents/CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/com.unity.ml-agents/CHANGELOG.md b/com.unity.ml-agents/CHANGELOG.md
index 345bd586b8..a69c83bba8 100755
--- a/com.unity.ml-agents/CHANGELOG.md
+++ b/com.unity.ml-agents/CHANGELOG.md
@@ -19,6 +19,8 @@ details. (#5060)
 
 ### 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 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)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 
 ### Bug Fixes

From b79d44e4becf9fc2702d3260f1403be015d8005c Mon Sep 17 00:00:00 2001
From: Vincent-Pierre BERGES <vincentpierre@unity3d.com>
Date: Thu, 11 Mar 2021 16:53:26 -0800
Subject: [PATCH 18/24] Exporting all the branches size instead of omly the sum
 (#5092)

---
 .../Inference/BarracudaModelExtensions.cs     | 16 ++++-
 .../Inference/BarracudaModelParamLoader.cs    | 61 ++++++++++++++++++-
 ml-agents/mlagents/trainers/torch/networks.py |  3 +-
 3 files changed, 73 insertions(+), 7 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index 78926608dc..c2e0067740 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -239,7 +239,7 @@ public static bool HasDiscreteOutputs(this Model model)
             else
             {
                 return model.outputs.Contains(TensorNames.DiscreteActionOutput) &&
-                    (int)model.GetTensorByName(TensorNames.DiscreteActionOutputShape)[0] > 0;
+                    (int)model.DiscreteOutputSize() > 0;
             }
         }
 
@@ -262,7 +262,19 @@ public static int DiscreteOutputSize(this Model model)
             else
             {
                 var discreteOutputShape = model.GetTensorByName(TensorNames.DiscreteActionOutputShape);
-                return discreteOutputShape == null ? 0 : (int)discreteOutputShape[0];
+                if (discreteOutputShape == null)
+                {
+                    return 0;
+                }
+                else
+                {
+                    int result = 0;
+                    for (int i = 0; i < discreteOutputShape.length; i++)
+                    {
+                        result += (int)discreteOutputShape[i];
+                    }
+                    return result;
+                }
             }
         }
 
diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 89d114bbb1..928f253060 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -716,8 +716,19 @@ static IEnumerable<FailedCheck> CheckOutputTensorShape(
             {
                 failedModelChecks.Add(continuousError);
             }
-            var modelSumDiscreteBranchSizes = model.DiscreteOutputSize();
-            var discreteError = CheckDiscreteActionOutputShape(brainParameters, actuatorComponents, modelSumDiscreteBranchSizes);
+            FailedCheck discreteError = null;
+            var modelApiVersion = model.GetVersion();
+            if (modelApiVersion == (int)ModelApiVersion.MLAgents1_0)
+            {
+                var modelSumDiscreteBranchSizes = model.DiscreteOutputSize();
+                discreteError = CheckDiscreteActionOutputShapeLegacy(brainParameters, actuatorComponents, modelSumDiscreteBranchSizes);
+            }
+            if (modelApiVersion == (int)ModelApiVersion.MLAgents2_0)
+            {
+                var modeDiscreteBranches = model.GetTensorByName(TensorNames.DiscreteActionOutputShape);
+                discreteError = CheckDiscreteActionOutputShape(brainParameters, actuatorComponents, modeDiscreteBranches);
+            }
+
             if (discreteError != null)
             {
                 failedModelChecks.Add(discreteError);
@@ -733,6 +744,50 @@ static IEnumerable<FailedCheck> CheckOutputTensorShape(
         /// The BrainParameters that are used verify the compatibility with the InferenceEngine
         /// </param>
         /// <param name="actuatorComponents">Array of attached actuator components.</param>
+        /// <param name="modelDiscreteBranches"> The Tensor of branch sizes.
+        /// </param>
+        /// <returns>
+        /// If the Check failed, returns a string containing information about why the
+        /// check failed. If the check passed, returns null.
+        /// </returns>
+        static FailedCheck CheckDiscreteActionOutputShape(
+            BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, Tensor modelDiscreteBranches)
+        {
+
+            var discreteActionBranches = brainParameters.ActionSpec.BranchSizes.ToList();
+            foreach (var actuatorComponent in actuatorComponents)
+            {
+                var actionSpec = actuatorComponent.ActionSpec;
+                discreteActionBranches.AddRange(actionSpec.BranchSizes);
+            }
+
+            if (modelDiscreteBranches.length != discreteActionBranches.Count)
+            {
+                return FailedCheck.Warning("Discrete Action Size of the model does not match. The BrainParameters expect " +
+                    $"{discreteActionBranches.Count} branches but the model contains {modelDiscreteBranches.length}."
+                );
+            }
+
+            for (int i = 0; i < modelDiscreteBranches.length; i++)
+            {
+                if (modelDiscreteBranches[i] != discreteActionBranches[i])
+                {
+                    return FailedCheck.Warning($"The number of Discrete Actions of branch {i} does not match. " +
+                    $"Was expecting {discreteActionBranches[i]} but the model contains {modelDiscreteBranches[i]} "
+                    );
+                }
+            }
+            return null;
+        }
+
+        /// <summary>
+        /// Checks that the shape of the discrete action output is the same in the
+        /// model and in the Brain Parameters. Tests the models created with the API of version 1.X
+        /// </summary>
+        /// <param name="brainParameters">
+        /// The BrainParameters that are used verify the compatibility with the InferenceEngine
+        /// </param>
+        /// <param name="actuatorComponents">Array of attached actuator components.</param>
         /// <param name="modelSumDiscreteBranchSizes">
         /// The size of the discrete action output that is expected by the model.
         /// </param>
@@ -740,7 +795,7 @@ static IEnumerable<FailedCheck> CheckOutputTensorShape(
         /// If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.
         /// </returns>
-        static FailedCheck CheckDiscreteActionOutputShape(
+        static FailedCheck CheckDiscreteActionOutputShapeLegacy(
             BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, int modelSumDiscreteBranchSizes)
         {
             // TODO: check each branch size instead of sum of branch sizes
diff --git a/ml-agents/mlagents/trainers/torch/networks.py b/ml-agents/mlagents/trainers/torch/networks.py
index 2f9717c7a2..341c0a0bd2 100644
--- a/ml-agents/mlagents/trainers/torch/networks.py
+++ b/ml-agents/mlagents/trainers/torch/networks.py
@@ -320,9 +320,8 @@ def __init__(
         self.continuous_act_size_vector = torch.nn.Parameter(
             torch.Tensor([int(self.action_spec.continuous_size)]), requires_grad=False
         )
-        # TODO: export list of branch sizes instead of sum
         self.discrete_act_size_vector = torch.nn.Parameter(
-            torch.Tensor([sum(self.action_spec.discrete_branches)]), requires_grad=False
+            torch.Tensor([self.action_spec.discrete_branches]), requires_grad=False
         )
         self.act_size_vector_deprecated = torch.nn.Parameter(
             torch.Tensor(

From fa22688a9b91befefacc88887efec121128be1fc Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Fri, 12 Mar 2021 13:46:13 -0800
Subject: [PATCH 19/24] addressing comments

---
 .../Inference/BarracudaModelParamLoader.cs    | 23 ++++++++-----------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 928f253060..29406fdb99 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -62,10 +62,14 @@ public static FailedCheck Error(string message)
         /// <param name="observableAttributeTotalSize">Sum of the sizes of all ObservableAttributes.</param>
         /// <param name="behaviorType">BehaviorType or the Agent to check.</param>
         /// <returns>A IEnumerable of the checks that failed</returns>
-        public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters brainParameters,
-            ISensor[] sensors, ActuatorComponent[] actuatorComponents,
+        public static IEnumerable<FailedCheck> CheckModel(
+            Model model,
+            BrainParameters brainParameters,
+            ISensor[] sensors,
+            ActuatorComponent[] actuatorComponents,
             int observableAttributeTotalSize = 0,
-            BehaviorType behaviorType = BehaviorType.Default)
+            BehaviorType behaviorType = BehaviorType.Default
+            )
         {
             List<FailedCheck> failedModelChecks = new List<FailedCheck>();
             if (model == null)
@@ -90,14 +94,6 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
             }
 
             var modelApiVersion = model.GetVersion();
-            if (modelApiVersion == -1)
-            {
-                failedModelChecks.Add(
-                    FailedCheck.Warning("Model was not trained using the right version of ML-Agents. " +
-                    "Cannot use this model.")
-                    );
-                return failedModelChecks;
-            }
             if (modelApiVersion < (int)ModelApiVersion.MinSupportedVersion || modelApiVersion > (int)ModelApiVersion.MaxSupportedVersion)
             {
                 failedModelChecks.Add(
@@ -116,8 +112,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                 return failedModelChecks;
             }
 
-            var modelVersion = model.GetVersion();
-            if (modelVersion == (int)ModelApiVersion.MLAgents1_0)
+            if (modelApiVersion == (int)ModelApiVersion.MLAgents1_0)
             {
                 failedModelChecks.AddRange(
                     CheckInputTensorPresenceLegacy(model, brainParameters, memorySize, sensors)
@@ -126,7 +121,7 @@ public static IEnumerable<FailedCheck> CheckModel(Model model, BrainParameters b
                     CheckInputTensorShapeLegacy(model, brainParameters, sensors, observableAttributeTotalSize)
                 );
             }
-            if (modelVersion == (int)ModelApiVersion.MLAgents2_0)
+            else if (modelApiVersion == (int)ModelApiVersion.MLAgents2_0)
             {
                 failedModelChecks.AddRange(
                     CheckInputTensorPresence(model, brainParameters, memorySize, sensors)

From f6317fdff0858f641ae809a554ba7550ce7803ab Mon Sep 17 00:00:00 2001
From: Vincent-Pierre BERGES <vincentpierre@unity3d.com>
Date: Fri, 12 Mar 2021 14:08:01 -0800
Subject: [PATCH 20/24] Update
 com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs

Co-authored-by: Chris Elion <chris.elion@unity3d.com>
---
 .../Runtime/Inference/BarracudaModelParamLoader.cs            | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 29406fdb99..994b1f17da 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -720,8 +720,8 @@ static IEnumerable<FailedCheck> CheckOutputTensorShape(
             }
             if (modelApiVersion == (int)ModelApiVersion.MLAgents2_0)
             {
-                var modeDiscreteBranches = model.GetTensorByName(TensorNames.DiscreteActionOutputShape);
-                discreteError = CheckDiscreteActionOutputShape(brainParameters, actuatorComponents, modeDiscreteBranches);
+                var modelDiscreteBranches = model.GetTensorByName(TensorNames.DiscreteActionOutputShape);
+                discreteError = CheckDiscreteActionOutputShape(brainParameters, actuatorComponents, modelDiscreteBranches);
             }
 
             if (discreteError != null)

From 1b48e4bd221c32a6517a18bac88f062344b11fa2 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Fri, 12 Mar 2021 14:08:20 -0800
Subject: [PATCH 21/24] readding tests

---
 .../Tests/Editor/EditModeTestInternalBrainTensorApplier.cs     | 3 ++-
 .../Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs   | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
index 5f4a3cc326..e679254908 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorApplier.cs
@@ -12,7 +12,8 @@ class TestAgent : Agent
         {
         }
 
-        public void Construction(int apiVersion)
+        [Test]
+        public void Construction()
         {
             var actionSpec = new ActionSpec();
             var alloc = new TensorCachingAllocator();
diff --git a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
index 70cc6c4f54..e1990c772b 100644
--- a/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
+++ b/com.unity.ml-agents/Tests/Editor/EditModeTestInternalBrainTensorGenerator.cs
@@ -63,7 +63,8 @@ static List<TestAgent> GetFakeAgents(ObservableAttributeOptions observableAttrib
             return agents;
         }
 
-        public void Construction(int apiVersion)
+        [Test]
+        public void Construction()
         {
             var alloc = new TensorCachingAllocator();
             var mem = new Dictionary<int, List<float>>();

From 7d1d3363c4402be2a065cff4a5a5ec2b34d05186 Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Mon, 15 Mar 2021 10:05:00 -0700
Subject: [PATCH 22/24] Adding a comment around the new DiscreteOutputSize
 method

---
 .../Runtime/Inference/BarracudaModelExtensions.cs         | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
index c2e0067740..5cb4d15c74 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelExtensions.cs
@@ -245,6 +245,14 @@ public static bool HasDiscreteOutputs(this Model model)
 
         /// <summary>
         /// Discrete action output size of the model. This is equal to the sum of the branch sizes.
+        /// This method gets the tensor representing the list of branch size and returns the
+        /// sum of all the elements in the Tensor.
+        ///  - In version 1.X this tensor contains a single number, the sum of all branch
+        /// size values.
+        ///  - In version 2.X this tensor contains a 1D Tensor with each element corresponding
+        /// to a branch size.
+        /// Since this method does the sum of all elements in the tensor, the output
+        /// will be the same on both 1.X and 2.X.
         /// </summary>
         /// <param name="model">
         /// The Barracuda engine model for loading static parameters.

From a124c96bb713a246872da23e8dad08f6258496cb Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Mon, 15 Mar 2021 10:13:24 -0700
Subject: [PATCH 23/24] Clearer warning : Model contains unexpected input >
 Model requires unknown input

---
 .../Runtime/Inference/BarracudaModelParamLoader.cs            | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 994b1f17da..6b9288a0b6 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -500,7 +500,7 @@ static IEnumerable<FailedCheck> CheckInputTensorShapeLegacy(
                     if (!tensor.name.Contains("visual_observation"))
                     {
                         failedModelChecks.Add(
-                            FailedCheck.Warning("Model requires an unknown input named : " + tensor.name)
+                            FailedCheck.Warning("Model contains an unexpected input named : " + tensor.name)
                             );
                     }
                 }
@@ -639,7 +639,7 @@ static IEnumerable<FailedCheck> CheckInputTensorShape(
             {
                 if (!tensorTester.ContainsKey(tensor.name))
                 {
-                    failedModelChecks.Add(FailedCheck.Warning("Model requires an unknown input named : " + tensor.name
+                    failedModelChecks.Add(FailedCheck.Warning("Model contains an unexpected input named : " + tensor.name
                     ));
                 }
                 else

From 61412ad513294239ad4ef29d432a011d2c635ebe Mon Sep 17 00:00:00 2001
From: vincentpierre <vincentpierre@unity3d.com>
Date: Mon, 15 Mar 2021 11:42:12 -0700
Subject: [PATCH 24/24] Fixing a bug in the case where the discrete action
 tensor does not exist

---
 .../Runtime/Inference/BarracudaModelParamLoader.cs         | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
index 6b9288a0b6..a21728f35b 100644
--- a/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
+++ b/com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs
@@ -756,14 +756,15 @@ static FailedCheck CheckDiscreteActionOutputShape(
                 discreteActionBranches.AddRange(actionSpec.BranchSizes);
             }
 
-            if (modelDiscreteBranches.length != discreteActionBranches.Count)
+            int modelDiscreteBranchesLength = modelDiscreteBranches?.length ?? 0;
+            if (modelDiscreteBranchesLength != discreteActionBranches.Count)
             {
                 return FailedCheck.Warning("Discrete Action Size of the model does not match. The BrainParameters expect " +
-                    $"{discreteActionBranches.Count} branches but the model contains {modelDiscreteBranches.length}."
+                    $"{discreteActionBranches.Count} branches but the model contains {modelDiscreteBranchesLength}."
                 );
             }
 
-            for (int i = 0; i < modelDiscreteBranches.length; i++)
+            for (int i = 0; i < modelDiscreteBranchesLength; i++)
             {
                 if (modelDiscreteBranches[i] != discreteActionBranches[i])
                 {