V2 staging new model version (#5080)

* Make modelCheck have flavors of error messages

* ONNX exporter v3

* Using a better CheckType and a switch statement

* Removing unused message

* More tests

* Use an enum for valid versions and use GetVersion on model directly

* Maybe the model export version a static constant in Python

* Use static constructor for FailedCheck

* Use static constructor for FailedCheck

* Modifying the docstrings

* renaming LegacyDiscreteActionOutputApplier

* removing testing code

* better warning message

* Nest the CheckTypeEnum into the FailedCheck class

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

* Adding a line explaining that legacy tensor checks are for versions 1.X only

* Modifying the changelog

* Exporting all the branches size instead of omly the sum (#5092)

* addressing comments

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

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

* readding tests

* Adding a comment around the new DiscreteOutputSize method

* Clearer warning : Model contains unexpected input > Model requires unknown input

* Fixing a bug in the case where the discrete action tensor does not exist

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

Author: vincentpierre

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

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 CheckTypeEnum = Unity.MLAgents.Inference.BarracudaModelParamLoader.FailedCheck.CheckTypeEnum;
 
 namespace Unity.MLAgents.Editor
 {
@@ -147,7 +148,20 @@ void DisplayFailedModelChecks()
                 {
                     if (check != null)
                     {
-                        EditorGUILayout.HelpBox(check, MessageType.Warning);
+                        switch (check.CheckType)
+                        {
+                            case CheckTypeEnum.Info:
+                                EditorGUILayout.HelpBox(check.Message, MessageType.Info);
+                                break;
+                            case CheckTypeEnum.Warning:
+                                EditorGUILayout.HelpBox(check.Message, MessageType.Warning);
+                                break;
+                            case CheckTypeEnum.Error:
+                                EditorGUILayout.HelpBox(check.Message, MessageType.Error);
+                                break;
+                            default:
+                                break;
+                        }
                     }
                 }
             }

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 LegacyDiscreteActionOutputApplier : 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 LegacyDiscreteActionOutputApplier(ActionSpec actionSpec, int seed, ITensorAllocator allocator)
         {
             m_ActionSize = actionSpec.BranchSizes;
             m_Multinomial = new Multinomial(seed);

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

@@ -1,6 +1,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using Unity.Barracuda;
+using FailedCheck = Unity.MLAgents.Inference.BarracudaModelParamLoader.FailedCheck;
 
 namespace Unity.MLAgents.Inference
 {
@@ -38,6 +39,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>
@@ -226,12 +239,20 @@ 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;
             }
         }
 
         /// <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.
@@ -249,7 +270,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;
+                }
             }
         }
 
@@ -298,21 +331,25 @@ 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(
+                    FailedCheck.Warning($"Required constant \"{TensorNames.VersionNumber}\" was not found in the model file.")
+                    );
                 return false;
             }
 
             // Check the presence of memory size
             var memorySizeTensor = model.GetTensorByName(TensorNames.MemorySize);
             if (memorySizeTensor == null)
             {
-                failedModelChecks.Add($"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;
             }
 
@@ -321,7 +358,9 @@ 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(
+                    FailedCheck.Warning("The model does not contain any Action Output Node.")
+                    );
                 return false;
             }
 
@@ -330,13 +369,18 @@ 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(
+                        FailedCheck.Warning("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(
+                        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;
                 }
             }
@@ -345,13 +389,17 @@ 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(
+                        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("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;
                 }
             }