[MLA-1724] Reduce use of IEnumerable during inference (#4887)

* improve allocations in inference

* Add IList overload for VectorSensor.AddObservation

* [skip ci] changelog

* [skip ci] migrating

Project/Assets/ML-Agents/Examples/SharedAssets/Scripts/SensorBase.cs

             float[] buffer = new float[numFloats];
             WriteObservation(buffer);
 
-            writer.AddRange(buffer);
+            writer.AddList(buffer);
 
             return numFloats;
         }

com.unity.ml-agents/CHANGELOG.md

 - Added the IHeuristicProvider interface to allow IActuators as well as Agent implement the Heuristic function to generate actions.
   Updated the Basic example and the Match3 Example to use Actuators.
   Changed the namespace and file names of classes in com.unity.ml-agents.extensions. (#4849)
+- Added `VectorSensor.AddObservation(IList<float>)`. `VectorSensor.AddObservation(IEnumerable<float>)`
+  is deprecated. The `IList` version is recommended, as it does not generate any
+  additional memory allocations. (#4887)
+- Added `ObservationWriter.AddList()` and deprecated `ObservationWriter.AddRange()`.
+  `AddList()` is recommended, as it does not generate any additional memory allocations. (#4887)
 
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 
 - CameraSensor now logs an error if the GraphicsDevice is null. (#4880)
+- Removed several memory allocations that happened during inference. On a test scene, this
+  reduced the amount of memory allocated by approximately 25%. (#4887)
+
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 - Fixed a bug that would cause an exception when `RunOptions` was deserialized via `pickle`. (#4842)
 - Fixed a bug that can cause a crash if a behavior can appear during training in multi-environment training. (#4872)

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

             m_ActionSpec = actionSpec;
         }
 
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
+        public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
-            foreach (int agentId in actionIds)
+            for (var i = 0; i < actionIds.Count; i++)
+                var agentId = actionIds[i];
                 if (lastActions.ContainsKey(agentId))
                 {
                     var actionBuffer = lastActions[agentId];
             m_ActionSpec = actionSpec;
         }
 
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
+        public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
         {
             //var tensorDataProbabilities = tensorProxy.Data as float[,];
             var idActionPairList = actionIds as List<int> ?? actionIds.ToList();
                 actionProbs.data.Dispose();
                 outputTensor.data.Dispose();
             }
+
-            foreach (int agentId in actionIds)
+            for (var i = 0; i < actionIds.Count; i++)
+                var agentId = actionIds[i];
                 if (lastActions.ContainsKey(agentId))
                 {
                     var actionBuffer = lastActions[agentId];
             m_Memories = memories;
         }
 
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
+        public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
-            foreach (int agentId in actionIds)
+            for (var i = 0; i < actionIds.Count; i++)
+                var agentId = actionIds[i];
                 List<float> memory;
                 if (!m_Memories.TryGetValue(agentId, out memory)
                     || memory.Count < memorySize)
             m_Memories = memories;
         }
 
-        public void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
+        public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
-            foreach (int agentId in actionIds)
+            for (var i = 0; i < actionIds.Count; i++)
+                var agentId = actionIds[i];
                 List<float> memory;
                 if (!m_Memories.TryGetValue(agentId, out memory)
                     || memory.Count < memorySize * m_MemoriesCount)

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

             m_Allocator = allocator;
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
         {
             TensorUtils.ResizeTensor(tensorProxy, batchSize, m_Allocator);
         }
             m_Allocator = allocator;
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
         {
             tensorProxy.data?.Dispose();
             tensorProxy.data = m_Allocator.Alloc(new TensorShape(1, 1));
             m_Allocator = allocator;
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
         {
             tensorProxy.shape = new long[0];
             tensorProxy.data?.Dispose();
         }
 
         public void Generate(
-            TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+            TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
-            foreach (var infoSensorPair in infos)
+            for (var infoIndex = 0; infoIndex < infos.Count; infoIndex++)
+                var infoSensorPair = infos[infoIndex];
                 var info = infoSensorPair.agentInfo;
                 List<float> memory;
 
             m_Memories = memories;
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
-            foreach (var infoSensorPair in infos)
+            for (var infoIndex = 0; infoIndex < infos.Count; infoIndex++)
+                var infoSensorPair = infos[infoIndex];
                 var info = infoSensorPair.agentInfo;
                 var offset = memorySize * m_MemoryIndex;
                 List<float> memory;
             m_Allocator = allocator;
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
-            foreach (var infoSensorPair in infos)
+            for (var infoIndex = 0; infoIndex < infos.Count; infoIndex++)
+                var infoSensorPair = infos[infoIndex];
                 var info = infoSensorPair.agentInfo;
                 var pastAction = info.storedActions.DiscreteActions;
                 if (!pastAction.IsEmpty())
             m_Allocator = allocator;
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
-            foreach (var infoSensorPair in infos)
+            for (var infoIndex = 0; infoIndex < infos.Count; infoIndex++)
+                var infoSensorPair = infos[infoIndex];
                 var agentInfo = infoSensorPair.agentInfo;
                 var maskList = agentInfo.discreteActionMasks;
                 for (var j = 0; j < maskSize; j++)
             m_Allocator = allocator;
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
         {
             TensorUtils.ResizeTensor(tensorProxy, batchSize, m_Allocator);
             TensorUtils.FillTensorWithRandomNormal(tensorProxy, m_RandomNormal);
             m_SensorIndices.Add(sensorIndex);
         }
 
-        public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos)
+        public void Generate(TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos)
-            foreach (var info in infos)
+            for (var infoIndex = 0; infoIndex < infos.Count; infoIndex++)
+                var info = infos[infoIndex];
                 if (info.agentInfo.done)
                 {
                     // If the agent is done, we might have a stale reference to the sensors
                 {
                     var tensorOffset = 0;
                     // Write each sensor consecutively to the tensor
-                    foreach (var sensorIndex in m_SensorIndices)
+                    for (var sensorIndexIndex = 0; sensorIndexIndex < m_SensorIndices.Count; sensorIndexIndex++)
+                        var sensorIndex = m_SensorIndices[sensorIndexIndex];
                         var sensor = info.sensors[sensorIndex];
                         m_ObservationWriter.SetTarget(tensorProxy, agentIndex, tensorOffset);
                         var numWritten = sensor.Write(m_ObservationWriter);

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

         TensorApplier m_TensorApplier;
 
         NNModel m_Model;
+        string m_ModelName;
-        IReadOnlyList<TensorProxy> m_InferenceOutputs;
+        List<TensorProxy> m_InferenceOutputs;
+        Dictionary<string, Tensor> m_InputsByName;
         Dictionary<int, List<float>> m_Memories = new Dictionary<int, List<float>>();
 
         SensorShapeValidator m_SensorShapeValidator = new SensorShapeValidator();
         {
             Model barracudaModel;
             m_Model = model;
+            m_ModelName = model.name;
             m_InferenceDevice = inferenceDevice;
             m_TensorAllocator = new TensorCachingAllocator();
             if (model != null)
                 seed, m_TensorAllocator, m_Memories, barracudaModel);
             m_TensorApplier = new TensorApplier(
                 actionSpec, seed, m_TensorAllocator, m_Memories, barracudaModel);
+            m_InputsByName = new Dictionary<string, Tensor>();
+            m_InferenceOutputs = new List<TensorProxy>();
         }
 
         public InferenceDevice InferenceDevice
             get { return m_Model; }
         }
 
-        static Dictionary<string, Tensor> PrepareBarracudaInputs(IEnumerable<TensorProxy> infInputs)
+        void PrepareBarracudaInputs(IReadOnlyList<TensorProxy> infInputs)
-            var inputs = new Dictionary<string, Tensor>();
-            foreach (var inp in infInputs)
+            m_InputsByName.Clear();
+            for (var i = 0; i < infInputs.Count; i++)
-                inputs[inp.name] = inp.data;
+                var inp = infInputs[i];
+                m_InputsByName[inp.name] = inp.data;
-
-            return inputs;
         }
 
         public void Dispose()
             m_TensorAllocator?.Reset(false);
         }
 
-        List<TensorProxy> FetchBarracudaOutputs(string[] names)
+        void FetchBarracudaOutputs(string[] names)
-            var outputs = new List<TensorProxy>();
+            m_InferenceOutputs.Clear();
-                outputs.Add(TensorUtils.TensorProxyFromBarracuda(output, n));
+                m_InferenceOutputs.Add(TensorUtils.TensorProxyFromBarracuda(output, n));
-
-            return outputs;
         }
 
         public void PutObservations(AgentInfo info, List<ISensor> sensors)
             }
 
             Profiler.BeginSample("ModelRunner.DecideAction");
+            Profiler.BeginSample(m_ModelName);
-            Profiler.BeginSample($"MLAgents.{m_Model.name}.GenerateTensors");
+            Profiler.BeginSample($"GenerateTensors");
-            Profiler.BeginSample($"MLAgents.{m_Model.name}.PrepareBarracudaInputs");
-            var inputs = PrepareBarracudaInputs(m_InferenceInputs);
+            Profiler.BeginSample($"PrepareBarracudaInputs");
+            PrepareBarracudaInputs(m_InferenceInputs);
-            Profiler.BeginSample($"MLAgents.{m_Model.name}.ExecuteGraph");
-            m_Engine.Execute(inputs);
+            Profiler.BeginSample($"ExecuteGraph");
+            m_Engine.Execute(m_InputsByName);
-            Profiler.BeginSample($"MLAgents.{m_Model.name}.FetchBarracudaOutputs");
-            m_InferenceOutputs = FetchBarracudaOutputs(m_OutputNames);
+            Profiler.BeginSample($"FetchBarracudaOutputs");
+            FetchBarracudaOutputs(m_OutputNames);
-            Profiler.BeginSample($"MLAgents.{m_Model.name}.ApplyTensors");
+            Profiler.BeginSample($"ApplyTensors");
-            Profiler.EndSample();
+            Profiler.EndSample(); // end name
+            Profiler.EndSample(); // end ModelRunner.DecideAction
 
             m_Infos.Clear();
 

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

             /// </param>
             /// <param name="actionIds"> List of Agents Ids that will be updated using the tensor's data</param>
             /// <param name="lastActions"> Dictionary of AgentId to Actions to be updated</param>
-            void Apply(TensorProxy tensorProxy, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions);
+            void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions);
         }
 
         readonly Dictionary<string, IApplier> m_Dict = new Dictionary<string, IApplier>();
         /// <exception cref="UnityAgentsException"> One of the tensor does not have an
         /// associated applier.</exception>
         public void ApplyTensors(
-            IEnumerable<TensorProxy> tensors, IEnumerable<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
+            IReadOnlyList<TensorProxy> tensors, IList<int> actionIds, Dictionary<int, ActionBuffers> lastActions)
-            foreach (var tensor in tensors)
+            for (var tensorIndex = 0; tensorIndex < tensors.Count; tensorIndex++)
+                var tensor = tensors[tensorIndex];
                 if (!m_Dict.ContainsKey(tensor.name))
                 {
                     throw new UnityAgentsException(

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

             /// the tensor's data.
             /// </param>
             void Generate(
-                TensorProxy tensorProxy, int batchSize, IEnumerable<AgentInfoSensorsPair> infos);
+                TensorProxy tensorProxy, int batchSize, IList<AgentInfoSensorsPair> infos);
         }
 
         readonly Dictionary<string, IGenerator> m_Dict = new Dictionary<string, IGenerator>();
         /// <exception cref="UnityAgentsException"> One of the tensor does not have an
         /// associated generator.</exception>
         public void GenerateTensors(
-            IEnumerable<TensorProxy> tensors, int currentBatchSize, IEnumerable<AgentInfoSensorsPair> infos)
+            IReadOnlyList<TensorProxy> tensors, int currentBatchSize, IList<AgentInfoSensorsPair> infos)
-            foreach (var tensor in tensors)
+            for (var tensorIndex = 0; tensorIndex < tensors.Count; tensorIndex++)
+                var tensor = tensors[tensorIndex];
                 if (!m_Dict.ContainsKey(tensor.name))
                 {
                     throw new UnityAgentsException(

com.unity.ml-agents/Runtime/Sensors/ObservationWriter.cs

         }
 
         /// <summary>
-        /// 1D write access at a specified index. Use AddRange if possible instead.
+        /// 1D write access at a specified index. Use AddList if possible instead.
         /// </summary>
         /// <param name="index">Index to write to.</param>
         public float this[int index]
         /// </summary>
         /// <param name="data"></param>
         /// <param name="writeOffset">Optional write offset.</param>
+        [Obsolete("Use AddList() for better performance")]
         public void AddRange(IEnumerable<float> data, int writeOffset = 0)
         {
             if (m_Data != null)
                 {
                     m_Proxy.data[m_Batch, index + m_Offset + writeOffset] = val;
                     index++;
+                }
+            }
+        }
+
+        public void AddList(IList<float> data, int writeOffset = 0)
+        {
+            if (m_Data != null)
+            {
+                for (var index = 0; index < data.Count; index++)
+                {
+                    var val = data[index];
+                    m_Data[index + m_Offset + writeOffset] = val;
+
+                }
+            }
+            else
+            {
+                for (var index = 0; index < data.Count; index++)
+                {
+                    var val = data[index];
+                    m_Proxy.data[m_Batch, index + m_Offset + writeOffset] = val;
                 }
             }
         }

com.unity.ml-agents/Runtime/Sensors/RayPerceptionSensor.cs

                     rayOutput.ToFloatArray(numDetectableTags, rayIndex, m_Observations);
                 }
                 // Finally, add the observations to the ObservationWriter
-                writer.AddRange(m_Observations);
+                writer.AddList(m_Observations);
             }
             return m_Observations.Length;
         }

com.unity.ml-agents/Runtime/Sensors/StackingSensor.cs

                 for (var i = 0; i < m_NumStackedObservations; i++)
                 {
                     var obsIndex = (m_CurrentIndex + 1 + i) % m_NumStackedObservations;
-                    writer.AddRange(m_StackedObservations[obsIndex], numWritten);
+                    writer.AddList(m_StackedObservations[obsIndex], numWritten);
                     numWritten += m_UnstackedObservationSize;
                 }
             }

com.unity.ml-agents/Runtime/Sensors/VectorSensor.cs

+using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using UnityEngine;
                     m_Observations.Add(0);
                 }
             }
-            writer.AddRange(m_Observations);
+            writer.AddList(m_Observations);
             return expectedObservations;
         }
 
         /// Adds a collection of float observations to the vector observations of the agent.
         /// </summary>
         /// <param name="observation">Observation.</param>
+        [Obsolete("Use AddObservation(IList<float>) for better performance.")]
+            }
+        }
+
+        /// <summary>
+        /// Adds a list or array of float observations to the vector observations of the agent.
+        /// </summary>
+        /// <param name="observation">Observation.</param>
+        public void AddObservation(IList<float> observation)
+        {
+            for (var i = 0; i < observation.Count; i++)
+            {
+                AddFloatObs(observation[i]);
             }
         }
 

com.unity.ml-agents/Tests/Editor/Sensor/ObservationWriterTests.cs

             writer[0] = 3f;
             Assert.AreEqual(new[] { 1f, 3f, 2f }, buffer);
 
-            // AddRange
+            // AddList
-            writer.AddRange(new[] { 4f, 5f });
+            writer.AddList(new[] { 4f, 5f });
-            // AddRange with offset
+            // AddList with offset
-            writer.AddRange(new[] { 6f, 7f });
+            writer.AddList(new[] { 6f, 7f });
             Assert.AreEqual(new[] { 4f, 6f, 7f }, buffer);
         }
 
             Assert.AreEqual(2f, t.data[1, 1]);
             Assert.AreEqual(3f, t.data[1, 2]);
 
-            // AddRange
+            // AddList
             t = new TensorProxy
             {
                 valueType = TensorProxy.TensorType.FloatingPoint,
             writer.SetTarget(t, 1, 1);
-            writer.AddRange(new[] { -1f, -2f });
+            writer.AddList(new[] { -1f, -2f });
             Assert.AreEqual(0f, t.data[0, 0]);
             Assert.AreEqual(0f, t.data[0, 1]);
             Assert.AreEqual(0f, t.data[0, 2]);

docs/Migrating.md

 ## Migrating to Release 13
 ### Implementing IHeuristic in your IActuator implementations
  - If you have any custom actuators, you can now implement the `IHeuristicProvider` interface to have your actuator
-handle the generation of actions when an Agent is running in heuristic mode.
+  handle the generation of actions when an Agent is running in heuristic mode.
+- `VectorSensor.AddObservation(IEnumerable<float>)` is deprecated. Use `VectorSensor.AddObservation(IList<float>)`
+  instead.
+- `ObservationWriter.AddRange()` is deprecated. Use `ObservationWriter.AddList()` instead.
 
 
 # Migrating