Multiple LSTM cell handling added to Barracuda code path

mantasp · mantasp · commit 3e6ba7882649 · 2019-06-06T17:01:48.000+03:00
diff --git a/UnitySDK/Assets/ML-Agents/Scripts/Agent.cs b/UnitySDK/Assets/ML-Agents/Scripts/Agent.cs
@@ -918,6 +918,11 @@ public void AppendMemoriesAction(List<float> memories)
             action.memories.AddRange(memories);
         }
 
+        public List<float> GetMemoriesAction()
+        {
+            return action.memories;
+        }
+
         /// <summary>
         /// Updates the text action.
         /// </summary>
diff --git a/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/ApplierImpl.cs b/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/ApplierImpl.cs
@@ -1,4 +1,5 @@
 ﻿using System.Collections.Generic;
+using System.Linq;
 using MLAgents.InferenceBrain.Utils;
 using UnityEngine;
 
@@ -99,34 +100,38 @@ public void Apply(Tensor tensor, Dictionary<Agent, AgentInfo> agentInfo)
     
     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 tensorDataMemory = tensor.Data as float[,];
             var agentIndex = 0;
-            var memorySize = tensor.Shape[tensor.Shape.Length - 1];
+            var memorySize = (int)tensor.Shape[tensor.Shape.Length - 1];
+            
             foreach (var agent in agentInfo.Keys)
             {
-                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++;
             }
         }
diff --git a/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/BarracudaModelParamLoader.cs b/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/BarracudaModelParamLoader.cs
@@ -109,8 +109,10 @@ public string[] GetOutputNames()
             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();
@@ -264,8 +266,8 @@ private void CheckInputTensorPresence(int memory, ModelActionType isContinuous)
             // 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.");
@@ -302,8 +304,8 @@ private void CheckOutputTensorPresence(int memory)
             {
                 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.");
@@ -325,9 +327,12 @@ private void CheckInputTensorShape()
                     {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;
diff --git a/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/GeneratorImpl.cs b/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/GeneratorImpl.cs
@@ -1,6 +1,7 @@
 ﻿using System.Collections.Generic;
 using System;
 using System.Linq;
+using Barracuda;
 using MLAgents.InferenceBrain.Utils;
 
 namespace MLAgents.InferenceBrain
@@ -119,28 +120,28 @@ public void Generate(Tensor tensor, int batchSize, Dictionary<Agent, AgentInfo>
     
     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;
         }
         
         public void Generate(Tensor tensor, int batchSize, Dictionary<Agent, AgentInfo> agentInfo)
         {
             tensor.Shape[0] = batchSize;
-            var memorySize = tensor.Shape[tensor.Shape.Length - 1];
+            
+            var memorySize = (int)tensor.Shape[tensor.Shape.Length - 1];
+            
             tensor.Data = new float[batchSize, memorySize];
             var agentIndex = 0;
             foreach (var agent in agentInfo.Keys)
             {
-                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)
                 {
diff --git a/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorApplier.cs b/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorApplier.cs
@@ -1,4 +1,5 @@
-﻿using System.Collections.Generic;
+﻿#define ENABLE_BARRACUDA
+using System.Collections.Generic;
 
 namespace MLAgents.InferenceBrain
 {    
@@ -37,7 +38,7 @@ public interface Applier
         /// <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)
@@ -51,8 +52,14 @@ public TensorApplier(BrainParameters bp, int seed)
             }
             _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>
diff --git a/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorGenerator.cs b/UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorGenerator.cs
@@ -1,4 +1,6 @@
-﻿using System.Collections.Generic;
+﻿#define ENABLE_BARRACUDA
+using System.Collections.Generic;
+using Barracuda;
 
 namespace MLAgents.InferenceBrain
 {
@@ -37,16 +39,21 @@ public interface Generator
         /// <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.SequenceLengthPlaceholder] = new SequenceLengthGenerator();
             _dict[TensorNames.VectorObservationPlacholder] = new VectorObservationGenerator();
             _dict[TensorNames.RecurrentInPlaceholder] = new RecurrentInputGenerator();
             
-            _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();
diff --git a/UnitySDK/Assets/ML-Agents/Scripts/LearningBrain.cs b/UnitySDK/Assets/ML-Agents/Scripts/LearningBrain.cs
@@ -126,8 +126,8 @@ public void ReloadModel(int seed = 0)
             _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
         }
         

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`using System.Collections.Generic;`
	`2`	`+using System.Linq;`
`2`	`3`	`using MLAgents.InferenceBrain.Utils;`
`3`	`4`	`using UnityEngine;`
`4`	`5`
`@@ -99,34 +100,38 @@ public void Apply(Tensor tensor, Dictionary<Agent, AgentInfo> agentInfo)`
`99`	`100`
`100`	`101`	`public class BarracudaMemoryOutputApplier : TensorApplier.Applier`
`101`	`102`	`{`
`102`		`- private bool firstHalf = true;`
	`103`	`+ private int memoriesCount;`
	`104`	`+ private int memoryIndex;`
`103`	`105`
`104`		`- public BarracudaMemoryOutputApplier(bool firstHalf)`
	`106`	`+ public BarracudaMemoryOutputApplier(int memoriesCount, int memoryIndex)`
`105`	`107`	`{`
`106`		`- this.firstHalf = firstHalf;`
	`108`	`+ this.memoriesCount = memoriesCount;`
	`109`	`+ this.memoryIndex = memoryIndex;`
`107`	`110`	`}`
`108`	`111`
`109`	`112`	`public void Apply(Tensor tensor, Dictionary<Agent, AgentInfo> agentInfo)`
`110`	`113`	`{`
`111`	`114`	`var tensorDataMemory = tensor.Data as float[,];`
`112`	`115`	`var agentIndex = 0;`
`113`		`- var memorySize = tensor.Shape[tensor.Shape.Length - 1];`
	`116`	`+ var memorySize = (int)tensor.Shape[tensor.Shape.Length - 1];`
	`117`	`+`
`114`	`118`	`foreach (var agent in agentInfo.Keys)`
`115`	`119`	`{`
`116`		`- var memory = new List<float>();`
`117`		`- for (var j = 0; j < memorySize; j++)`
`118`		`- {`
`119`		`- memory.Add(tensorDataMemory[agentIndex, j]);`
`120`		`- }`
	`120`	`+ var memory = agent.GetMemoriesAction();`
`121`	`121`
`122`		`- if (firstHalf)`
	`122`	`+ if (memory == null \|\| memory.Count < memorySize * memoriesCount)`
`123`	`123`	`{`
`124`		`- agent.UpdateMemoriesAction(memory);`
	`124`	`+ memory = new List<float>();`
	`125`	`+ memory.AddRange(Enumerable.Repeat(0f, memorySize * memoriesCount));`
`125`	`126`	`}`
`126`		`- else`
	`127`	`+`
	`128`	`+ for (var j = 0; j < memorySize; j++)`
`127`	`129`	`{`
`128`		`- agent.AppendMemoriesAction(memory);`
	`130`	`+ memory[memorySize * memoryIndex + j] = tensorDataMemory[agentIndex, j];`
`129`	`131`	`}`
	`132`	`+`
	`133`	`+ agent.UpdateMemoriesAction(memory);`
	`134`	`+`
`130`	`135`	`agentIndex++;`
`131`	`136`	`}`
`132`	`137`	`}`