V2 staging new model version

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>
@@ -298,21 +311,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 +338,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 +349,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 +369,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;
                 }
             }