Multiple LSTM cell handling added to Barracuda code path

UnitySDK/Assets/ML-Agents/Scripts/Agent.cs

             action.memories.AddRange(memories);
         }
 
+        public List<float> GetMemoriesAction()
+        {
+            return action.memories;
+        }
+
         /// <summary>
         /// Updates the text action.
         /// </summary>

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/ApplierImpl.cs

 using System.Collections.Generic;
+using System.Linq;
 using MLAgents.InferenceBrain.Utils;
 using UnityEngine;
 
     
     public class BarracudaMemoryOutputApplier : TensorApplier.Applier
     {
-        private bool firstHalf = true;
+        private int memoriesCount;
+        private int memoryIndex;
-        public BarracudaMemoryOutputApplier(bool firstHalf)
+        public BarracudaMemoryOutputApplier(int memoriesCount, int memoryIndex)
-            this.firstHalf = firstHalf;
+            this.memoriesCount = memoriesCount;
+            this.memoryIndex = memoryIndex;
         }
         
         public void Apply(Tensor tensor, Dictionary<Agent, AgentInfo> agentInfo)
-            var memorySize = tensor.Shape[tensor.Shape.Length - 1];
+            var memorySize = (int)tensor.Shape[tensor.Shape.Length - 1];
+            
-                var memory = new List<float>();
-                for (var j = 0; j < memorySize; j++)
-                {
-                    memory.Add(tensorDataMemory[agentIndex, j]);
-                }
+                var memory = agent.GetMemoriesAction();
-                if (firstHalf)
+                if (memory == null || memory.Count < memorySize * memoriesCount)
-                    agent.UpdateMemoriesAction(memory);
+                    memory = new List<float>();
+                    memory.AddRange(Enumerable.Repeat(0f, memorySize * memoriesCount));
-                else
+
+                for (var j = 0; j < memorySize; j++)
-                    agent.AppendMemoriesAction(memory);
+                    memory[memorySize * memoryIndex + j] = tensorDataMemory[agentIndex, j];
+                
+                agent.UpdateMemoriesAction(memory);
+
                 agentIndex++;
             }
         }

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/BarracudaModelParamLoader.cs

             var memory = GetIntScalar(TensorNames.MemorySize);
             if (memory > 0)
             {
-                names.Add(TensorNames.RecurrentOutput_C);
-                names.Add(TensorNames.RecurrentOutput_H);
+                foreach (var mem in _model.memories)
+                {
+                    names.Add(mem.output);
+                }
             }
 
             names.Sort();
             // If the model has a non-negative memory size but requires a recurrent input
             if (memory > 0)
             {
-                if (!tensorsNames.Contains(TensorNames.RecurrentInPlaceholder_H) ||
-                    !tensorsNames.Contains(TensorNames.RecurrentInPlaceholder_C))
+                if (!tensorsNames.Any(x => x.EndsWith("_h")) ||
+                    !tensorsNames.Any(x => x.EndsWith("_c")))
                 {
                     _failedModelChecks.Add(
                         "The model does not contain a Recurrent Input Node but has memory_size.");
             {
                 var memOutputs = _model.memories.Select(x => x.output).ToList();
                 
-                if (!memOutputs.Contains(TensorNames.RecurrentOutput_H) || 
-                    !memOutputs.Contains(TensorNames.RecurrentOutput_C))
+                if (!memOutputs.Any(x => x.EndsWith("_h")) || 
+                    !memOutputs.Any(x => x.EndsWith("_c")))
                 {
                     _failedModelChecks.Add(
                         "The model does not contain a Recurrent Output Node but has memory_size.");
                     {TensorNames.RandomNormalEpsilonPlaceholder, ((tensor) => null)},
                     {TensorNames.ActionMaskPlaceholder, ((tensor) => null)},
                     {TensorNames.SequenceLengthPlaceholder, ((tensor) => null)},
-                    {TensorNames.RecurrentInPlaceholder_H, ((tensor) => null)},
-                    {TensorNames.RecurrentInPlaceholder_C, ((tensor) => null)},
+                    {TensorNames.RecurrentInPlaceholder, ((tensor) => null)},
+
+            foreach (var mem in _model.memories)
+                tensorTester[mem.input] = ((tensor) => null);
+            
             for (var obsIndex = 0; obsIndex < _brainParameters.cameraResolutions.Length; obsIndex++)
             {
                 var index = obsIndex;

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/GeneratorImpl.cs

 using System.Collections.Generic;
 using System;
 using System.Linq;
+using Barracuda;
 using MLAgents.InferenceBrain.Utils;
 
 namespace MLAgents.InferenceBrain
     
     public class BarracudaRecurrentInputGenerator : TensorGenerator.Generator
     {
-        private bool firstHalf = true;
+        private int memoriesCount;
+        private int memoryIndex;
-        public BarracudaRecurrentInputGenerator(bool firstHalf)
+        public BarracudaRecurrentInputGenerator(int memoriesCount, int memoryIndex)
-            this.firstHalf = firstHalf;
+            this.memoriesCount = memoriesCount;
+            this.memoryIndex = memoryIndex;
-            var memorySize = tensor.Shape[tensor.Shape.Length - 1];
+            
+            var memorySize = (int)tensor.Shape[tensor.Shape.Length - 1];
+            
-                var memory = agentInfo[agent].memories;
+                var memory = agentInfo[agent].memories;         
-                int offset = 0;
-                if (!firstHalf)
-                {
-                    offset = memory.Count - (int)memorySize;
-                }
+                int offset = memorySize * memoryIndex;
                 
                 if (memory == null)
                 {

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorApplier.cs

-using System.Collections.Generic;
+#define ENABLE_BARRACUDA
+using System.Collections.Generic;
 
 namespace MLAgents.InferenceBrain
 {    
         /// <param name="bp"> The BrainParameters used to determine what Appliers will be
         /// used</param>
         /// <param name="seed"> The seed the Appliers will be initialized with.</param>
-        public TensorApplier(BrainParameters bp, int seed)
+        public TensorApplier(BrainParameters bp, int seed, object barracudaModel = null)
         {
             _dict[TensorNames.ValueEstimateOutput] = new ValueEstimateApplier();
             if (bp.vectorActionSpaceType == SpaceType.continuous)
             }
             _dict[TensorNames.RecurrentOutput] = new MemoryOutputApplier();
             
-            _dict[TensorNames.RecurrentOutput_C] = new BarracudaMemoryOutputApplier(true);
-            _dict[TensorNames.RecurrentOutput_H] = new BarracudaMemoryOutputApplier(false);
+            #if ENABLE_BARRACUDA
+            Barracuda.Model model = (Barracuda.Model) barracudaModel;
+
+            for (var i = 0; i < model?.memories.Length; i++)
+            {
+                _dict[model.memories[i].output] = new BarracudaMemoryOutputApplier(model.memories.Length, i);
+            }
+            #endif
         }
 
         /// <summary>

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorGenerator.cs

-using System.Collections.Generic;
+#define ENABLE_BARRACUDA
+using System.Collections.Generic;
+using Barracuda;
 
 namespace MLAgents.InferenceBrain
 {
         /// <param name="bp"> The BrainParameters used to determine what Generators will be
         /// used</param>
         /// <param name="seed"> The seed the Generators will be initialized with.</param>
-        public TensorGenerator(BrainParameters bp, int seed)
+        public TensorGenerator(BrainParameters bp, int seed, object barracudaModel = null)
         {
             // Generator for Inputs
             _dict[TensorNames.BatchSizePlaceholder] = new BatchSizeGenerator();
             
-            _dict[TensorNames.RecurrentInPlaceholder_C] = new BarracudaRecurrentInputGenerator(true);
-            _dict[TensorNames.RecurrentInPlaceholder_H] = new BarracudaRecurrentInputGenerator(false);
+            #if ENABLE_BARRACUDA
+            Barracuda.Model model = (Barracuda.Model) barracudaModel;
+            for (var i = 0; i < model?.memories.Length; i++)
+            {
+                _dict[model.memories[i].input] = new BarracudaRecurrentInputGenerator(model.memories.Length, i);
+            }
+            #endif
             
             _dict[TensorNames.PreviousActionPlaceholder] = new PreviousActionInputGenerator();
             _dict[TensorNames.ActionMaskPlaceholder] = new ActionMaskInputGenerator();

UnitySDK/Assets/ML-Agents/Scripts/LearningBrain.cs

             _modelParamLoader = BarracudaModelParamLoader.GetLoaderAndCheck(_engine, _barracudaModel, brainParameters);
             _inferenceInputs = _modelParamLoader.GetInputTensors();
             _outputNames = _modelParamLoader.GetOutputNames();
-            _tensorGenerator = new TensorGenerator(brainParameters, seed);
-            _tensorApplier = new TensorApplier(brainParameters, seed);
+            _tensorGenerator = new TensorGenerator(brainParameters, seed, _barracudaModel);
+            _tensorApplier = new TensorApplier(brainParameters, seed, _barracudaModel);
 #endif
         }