Barracuda inference for hybrid actions

.yamato/com.unity.ml-agents-performance.yml

@@ -20,8 +20,10 @@ Run_Mac_Perfomance_Tests{{ editor.version }}:
     - ./utr --suite=editor --platform=StandaloneOSX --editor-location=.Editor --testproject=DevProject --artifacts_path=build/test-results --report-performance-data --performance-project-id=com.unity.ml-agents --zero-tests-are-ok=1
   triggers:
     cancel_old_ci: true
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "com.unity.ml-agents/**" OR

.yamato/com.unity.ml-agents-test.yml

@@ -92,8 +92,10 @@ test_{{ package.name }}_{{ platform.name }}_{{ editor.version }}:
   triggers:
     cancel_old_ci: true
     {% if platform.name == "mac" %}
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "com.unity.ml-agents/**" OR

.yamato/compressed-sensor-test.yml

@@ -18,8 +18,10 @@ test_compressed_obs_{{ editor.version }}:
     - .yamato/standalone-build-test.yml#test_mac_standalone_{{ editor.version }}
   triggers:
     cancel_old_ci: true
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "com.unity.ml-agents/**" OR

.yamato/gym-interface-test.yml

@@ -17,8 +17,10 @@ test_gym_interface_{{ editor.version }}:
     - .yamato/standalone-build-test.yml#test_mac_standalone_{{ editor.version }}
   triggers:
     cancel_old_ci: true
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "com.unity.ml-agents/**" OR

.yamato/protobuf-generation-test.yml

@@ -27,8 +27,10 @@ test_mac_protobuf_generation:
       git diff -- :/ ":(exclude,top)$CS_PROTO_PATH/*.meta" > artifacts/proto.patch; exit $GIT_ERR; }
   triggers:
     cancel_old_ci: true
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "protobuf-definitions/**" OR

.yamato/python-ll-api-test.yml

@@ -20,8 +20,10 @@ test_mac_ll_api_{{ editor.version }}:
     - .yamato/standalone-build-test.yml#test_mac_standalone_{{ editor.version }}
   triggers:
     cancel_old_ci: true
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "com.unity.ml-agents/**" OR

.yamato/standalone-build-test.yml

@@ -22,8 +22,10 @@ test_mac_standalone_{{ editor.version }}:
     - python -u -m ml-agents.tests.yamato.standalone_build_tests --scene=Assets/ML-Agents/TestScenes/TestCompressedTexture/TestTextureCompressed.unity
   triggers:
     cancel_old_ci: true
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "com.unity.ml-agents/**" OR

.yamato/training-int-tests.yml

@@ -24,8 +24,10 @@ test_mac_training_int_{{ editor.version }}:
     - .yamato/standalone-build-test.yml#test_mac_standalone_{{ editor.version }}
   triggers:
     cancel_old_ci: true
+    # TODO remove develop-hybrid trigger before merging to master
     expression: |
       (pull_request.target eq "master" OR
+      pull_request.target match "develop-hybrid.+" OR
       pull_request.target match "release.+") AND
       NOT pull_request.draft AND
       (pull_request.changes.any match "com.unity.ml-agents/**" OR

com.unity.ml-agents/Runtime/Actuators/ActionSegment.cs

@@ -107,6 +107,14 @@ public void Clear()
             System.Array.Clear(Array, Offset, Length);
         }
 
+        /// <summary>
+        /// Check if the segment is empty.
+        /// </summary>
+        public bool IsEmpty()
+        {
+            return Array.Length == 0;
+        }
+
         /// <inheritdoc/>
         IEnumerator<T> IEnumerable<T>.GetEnumerator()
         {

com.unity.ml-agents/Runtime/Actuators/IActionReceiver.cs

@@ -61,6 +61,15 @@ public ActionBuffers(ActionSegment<float> continuousActions, ActionSegment<int>
             DiscreteActions = discreteActions;
         }
 
+        /// <summary>
+        /// Construct an <see cref="ActionBuffers"/> instance with <see cref="ActionSpec"/>. All values are initialized to zeros.
+        /// /// </summary>
+        /// <param name="actionSpec">The <see cref="ActionSpec"/>  to send to an <see cref="IActionReceiver"/>.</param>
+        public ActionBuffers(ActionSpec actionSpec)
+            : this(new ActionSegment<float>(new float[actionSpec.NumContinuousActions]),
+            new ActionSegment<int>(new int[actionSpec.NumDiscreteActions]))
+        { }
+
         /// <summary>
         /// Create an <see cref="ActionBuffers"/> instance with ActionSpec and all actions stored as a float array.
         /// </summary>
@@ -110,6 +119,14 @@ public void Clear()
             DiscreteActions.Clear();
         }
 
+        /// <summary>
+        /// Check if the <see cref="ActionBuffers"/> is empty.
+        /// </summary>
+        public bool IsEmpty()
+        {
+            return ContinuousActions.IsEmpty() && DiscreteActions.IsEmpty();
+        }
+
         /// <inheritdoc/>
         public override bool Equals(object obj)
         {

com.unity.ml-agents/Runtime/Communicator/GrpcExtensions.cs

@@ -126,8 +126,11 @@ public static BrainParametersProto ToBrainParametersProto(this ActionSpec action
             {
                 NumContinuousActions = actionSpec.NumContinuousActions,
                 NumDiscreteActions = actionSpec.NumDiscreteActions,
-                DiscreteBranchSizes = { actionSpec.BranchSizes },
             };
+            if (actionSpec.BranchSizes != null)
+            {
+                actionSpecProto.DiscreteBranchSizes.AddRange(actionSpec.BranchSizes);
+            }
             brainParametersProto.ActionSpec = actionSpecProto;
 
             var supportHybrid = Academy.Instance.TrainerCapabilities == null || Academy.Instance.TrainerCapabilities.HybridActions;

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

@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using Unity.MLAgents.Inference.Utils;
+using Unity.MLAgents.Actuators;
 using Unity.Barracuda;
 using UnityEngine;
 
@@ -13,23 +14,31 @@ namespace Unity.MLAgents.Inference
     /// </summary>
     internal class ContinuousActionOutputApplier : TensorApplier.IApplier
     {
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, float[]> lastActions)
+        readonly ActionSpec m_ActionSpec;
+
+        public ContinuousActionOutputApplier(ActionSpec actionSpec)
+        {
+            m_ActionSpec = actionSpec;
+        }
+
+        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
         {
             var actionSize = tensorProxy.shape[tensorProxy.shape.Length - 1];
             var agentIndex = 0;
             foreach (int agentId in actionIds)
             {
                 if (lastActions.ContainsKey(agentId))
                 {
-                    var actionValue = lastActions[agentId];
-                    if (actionValue == null)
+                    var actionBuffer = lastActions[agentId];
+                    if (actionBuffer.IsEmpty())
                     {
-                        actionValue = new float[actionSize];
-                        lastActions[agentId] = actionValue;
+                        actionBuffer = new ActionBuffers(m_ActionSpec);
+                        lastActions[agentId] = actionBuffer;
                     }
+                    var continuousBuffer = actionBuffer.ContinuousActions;
                     for (var j = 0; j < actionSize; j++)
                     {
-                        actionValue[j] = tensorProxy.data[agentIndex, j];
+                        continuousBuffer[j] = tensorProxy.data[agentIndex, j];
                     }
                 }
                 agentIndex++;
@@ -46,15 +55,17 @@ internal class DiscreteActionOutputApplier : TensorApplier.IApplier
         readonly int[] m_ActionSize;
         readonly Multinomial m_Multinomial;
         readonly ITensorAllocator m_Allocator;
+        readonly ActionSpec m_ActionSpec;
 
-        public DiscreteActionOutputApplier(int[] actionSize, int seed, ITensorAllocator allocator)
+        public DiscreteActionOutputApplier(ActionSpec actionSpec, int seed, ITensorAllocator allocator)
         {
-            m_ActionSize = actionSize;
+            m_ActionSize = actionSpec.BranchSizes;
             m_Multinomial = new Multinomial(seed);
             m_Allocator = allocator;
+            m_ActionSpec = actionSpec;
         }
 
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, float[]> lastActions)
+        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
         {
             //var tensorDataProbabilities = tensorProxy.Data as float[,];
             var idActionPairList = actionIds as List<int> ?? actionIds.ToList();
@@ -103,15 +114,16 @@ public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionar
             {
                 if (lastActions.ContainsKey(agentId))
                 {
-                    var actionVal = lastActions[agentId];
-                    if (actionVal == null)
+                    var actionBuffer = lastActions[agentId];
+                    if (actionBuffer.IsEmpty())
                     {
-                        actionVal = new float[m_ActionSize.Length];
-                        lastActions[agentId] = actionVal;
+                        actionBuffer = new ActionBuffers(m_ActionSpec);
+                        lastActions[agentId] = actionBuffer;
                     }
+                    var discreteBuffer = actionBuffer.DiscreteActions;
                     for (var j = 0; j < m_ActionSize.Length; j++)
                     {
-                        actionVal[j] = actionValues[agentIndex, j];
+                        discreteBuffer[j] = (int)actionValues[agentIndex, j];
                     }
                 }
                 agentIndex++;
@@ -197,7 +209,7 @@ public MemoryOutputApplier(
             m_Memories = memories;
         }
 
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, float[]> lastActions)
+        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
         {
             var agentIndex = 0;
             var memorySize = (int)tensorProxy.shape[tensorProxy.shape.Length - 1];
@@ -234,7 +246,7 @@ public BarracudaMemoryOutputApplier(
             m_Memories = memories;
         }
 
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, float[]> lastActions)
+        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
         {
             var agentIndex = 0;
             var memorySize = (int)tensorProxy.shape[tensorProxy.shape.Length - 1];

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

@@ -0,0 +1,141 @@
+using Unity.Barracuda;
+
+namespace Unity.MLAgents.Inference
+{
+    /// <summary>
+    /// Barracuda Model extension methods.
+    /// </summary>
+    internal static class BarracudaModelExtensions
+    {
+        /// <summary>
+        /// Check if the model has continuous action outputs.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>True if the model has continuous action outputs.</returns>
+        public static bool HasContinuousOutputs(this Model model)
+        {
+            if (model.UseDeprecated())
+            {
+                return (int)model.GetTensorByName(TensorNames.IsContinuousControlDeprecated)[0] > 0;
+            }
+            else
+            {
+                return model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
+                    (int)model.GetTensorByName(TensorNames.ContinuousActionOutputShape)[0] > 0;
+            }
+        }
+
+        /// <summary>
+        /// Continuous action output size of the model.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>Size of continuous action output.</returns>
+        public static int ContinuousOutputSize(this Model model)
+        {
+            if (model.UseDeprecated())
+            {
+                return (int)model.GetTensorByName(TensorNames.IsContinuousControlDeprecated)[0] > 0 ?
+                    (int)model.GetTensorByName(TensorNames.ActionOutputShapeDeprecated)[0] : 0;
+            }
+            else
+            {
+                return (int)model.GetTensorByName(TensorNames.ContinuousActionOutputShape)[0];
+            }
+        }
+
+        /// <summary>
+        /// Continuous action output tensor name of the model.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>Tensor name of continuous action output.</returns>
+        public static string ContinuousOutputName(this Model model)
+        {
+            if (model.UseDeprecated())
+            {
+                return TensorNames.ActionOutputDeprecated;
+            }
+            else
+            {
+                return TensorNames.ContinuousActionOutput;
+            }
+        }
+
+        /// <summary>
+        /// Check if the model has discrete action outputs.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>True if the model has discrete action outputs.</returns>
+        public static bool HasDiscreteOutputs(this Model model)
+        {
+            if (model.UseDeprecated())
+            {
+                return (int)model.GetTensorByName(TensorNames.IsContinuousControlDeprecated)[0] == 0;
+            }
+            else
+            {
+                return model.outputs.Contains(TensorNames.DiscreteActionOutput) &&
+                    (int)model.GetTensorByName(TensorNames.DiscreteActionOutputShape)[0] > 0;
+            }
+        }
+
+        /// <summary>
+        /// Discrete action output size of the model. This is equal to the sum of the branch sizes.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>Size of discrete action output.</returns>
+        public static int DiscreteOutputSize(this Model model)
+        {
+            if (model.UseDeprecated())
+            {
+                return (int)model.GetTensorByName(TensorNames.IsContinuousControlDeprecated)[0] > 0 ?
+                    0 : (int)model.GetTensorByName(TensorNames.ActionOutputShapeDeprecated)[0];
+            }
+            else
+            {
+                return (int)model.GetTensorByName(TensorNames.DiscreteActionOutputShape)[0];
+            }
+        }
+
+        /// <summary>
+        /// Discrete action output tensor name of the model.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>Tensor name of discrete action output.</returns>
+        public static string DiscreteOutputName(this Model model)
+        {
+            if (model.UseDeprecated())
+            {
+                return TensorNames.ActionOutputDeprecated;
+            }
+            else
+            {
+                return TensorNames.DiscreteActionOutput;
+            }
+        }
+
+        /// <summary>
+        /// Check if the model uses deprecated output fields and should be handled differently.
+        /// </summary>
+        /// <param name="model">
+        /// The Barracuda engine model for loading static parameters.
+        /// </param>
+        /// <returns>True if the model uses deprecated output fields.</returns>
+        public static bool UseDeprecated(this Model model)
+        {
+            return !model.outputs.Contains(TensorNames.ContinuousActionOutput) &&
+                !model.outputs.Contains(TensorNames.DiscreteActionOutput);
+        }
+    }
+}