Made ModelParamLoader static and made a modelRunner for batching at i…

…nference (#2718)

* Created the model runner and uses a shared interface with the communicator.

Fixing bugs with dealocation

Removing unnecessary code

Added code comments

Renaming

* Addressing comments

* Modified the constructor of ModelRunner

* Addressing comments

* remaming the _verbose variable

* Addressing comments : Removed the Verbose check in the LearningBrainEditor

UnitySDK/Assets/ML-Agents/Editor/LearningBrainEditor.cs

@@ -1,5 +1,7 @@
 using UnityEngine;
 using UnityEditor;
+using Barracuda;
+
 
 namespace MLAgents
 {
@@ -57,12 +59,18 @@ public override void OnInspectorGUI()
             }
             if (m_RequireReload && m_TimeSinceModelReload > k_TimeBetweenModelReloads)
             {
-                brain.ReloadModel();
                 m_RequireReload = false;
                 m_TimeSinceModelReload = 0;
             }
             // Display all failed checks
-            var failedChecks = brain.GetModelFailedChecks();
+            D.logEnabled = false;
+            Barracuda.Model barracudaModel = null;
+            if (brain.model != null)
+            {
+                barracudaModel = ModelLoader.Load(brain.model.Value);
+            }
+            var failedChecks = InferenceBrain.BarracudaModelParamLoader.CheckModel(
+                barracudaModel, brain.brainParameters);
             foreach (var check in failedChecks)
             {
                 if (check != null)

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

@@ -1,10 +1,11 @@
 using UnityEngine;
-using System.Linq;
+using System.Collections.Generic;
 using UnityEngine.Serialization;
 #if UNITY_EDITOR
 using UnityEditor;
-
 #endif
+using MLAgents.InferenceBrain;
+using Barracuda;
 
 /**
  * Welcome to Unity Machine Learning Agents (ML-Agents).
@@ -179,6 +180,8 @@ bool IsCommunicatorOn
         /// Pointer to the communicator currently in use by the Academy.
         public ICommunicator Communicator;
 
+        private List<ModelRunner> m_ModelRunners = new List<ModelRunner>();
+
         // Flag used to keep track of the first time the Academy is reset.
         bool m_FirstAcademyReset;
 
@@ -300,7 +303,8 @@ private void InitializeEnvironment()
                     Communicator = null;
                 }
 
-                if (Communicator != null){
+                if (Communicator != null)
+                {
                     Communicator.QuitCommandReceived += OnQuitCommandReceived;
                     Communicator.ResetCommandReceived += OnResetCommand;
                     Communicator.RLInputReceived += OnRLInputReceived;
@@ -313,13 +317,13 @@ private void InitializeEnvironment()
 
             SetIsInference(!IsCommunicatorOn);
 
-            BrainDecideAction += () => {};
-            DestroyAction += () => {};
-            AgentSetStatus += i => {};
-            AgentResetIfDone += () => {};
-            AgentSendState += () => {};
-            AgentAct += () => {};
-            AgentForceReset += () => {};
+            BrainDecideAction += () => { };
+            DestroyAction += () => { };
+            AgentSetStatus += i => { };
+            AgentResetIfDone += () => { };
+            AgentSendState += () => { };
+            AgentAct += () => { };
+            AgentForceReset += () => { };
 
             ConfigureEnvironment();
         }
@@ -555,6 +559,26 @@ void FixedUpdate()
             EnvironmentStep();
         }
 
+        /// <summary>
+        /// Creates or retrieves an existing ModelRunner that uses the same NNModel and the InferenceDevice as
+        /// provided.
+        /// </summary>
+        /// <param name="model"> The NNModel the ModelRunner must use </param>
+        /// <param name="brainParameters"> The brainParameters used to create the ModelRunner </param>
+        /// <param name="inferenceDevice"> The inference device (CPU or GPU) the ModelRunner will use </param>
+        /// <returns> The ModelRunner compatible with the input settings</returns>
+        public ModelRunner GetOrCreateModelRunner(NNModel model, BrainParameters brainParameters, InferenceDevice inferenceDevice)
+        {
+            var modelRunner = m_ModelRunners.Find(x => x.HasModel(model, inferenceDevice));
+            if (modelRunner == null)
+            {
+                modelRunner = new ModelRunner(
+                    model, brainParameters, inferenceDevice);
+                m_ModelRunners.Add(modelRunner);
+            }
+            return modelRunner;
+        }
+
         /// <summary>
         /// Cleanup function
         /// </summary>

UnitySDK/Assets/ML-Agents/Scripts/Grpc/RpcCommunicator.cs

@@ -130,7 +130,8 @@ public void SubscribeBrain(string brainKey, BrainParameters brainParameters)
             m_CurrentUnityRlOutput.AgentInfos.Add(
                 brainKey,
                 new CommunicatorObjects.UnityRLOutputProto.Types.ListAgentInfoProto());
-            if (m_CurrentUnityRlInitializationOutput == null){
+            if (m_CurrentUnityRlInitializationOutput == null)
+            {
                 m_CurrentUnityRlInitializationOutput = new CommunicatorObjects.UnityRLInitializationOutputProto();
             }
             m_CurrentUnityRlInitializationOutput.BrainParameters.Add(brainParameters.ToProto(brainKey, true));
@@ -175,7 +176,7 @@ private UnityInputProto Initialize(UnityOutputProto unityOutput,
         /// <summary>
         /// Close the communicator gracefully on both sides of the communication.
         /// </summary>
-        public void Close()
+        public void Dispose()
         {
 # if UNITY_EDITOR || UNITY_STANDALONE_WIN || UNITY_STANDALONE_OSX || UNITY_STANDALONE_LINUX
             if (!m_IsOpen)
@@ -206,19 +207,19 @@ private void SendCommandEvent(CommandProto command, EnvironmentParametersProto e
             switch (command)
             {
                 case CommandProto.Quit:
-                {
-                    QuitCommandReceived?.Invoke();
-                    return;
-                }
+                    {
+                        QuitCommandReceived?.Invoke();
+                        return;
+                    }
                 case CommandProto.Reset:
-                {
-                    ResetCommandReceived?.Invoke(environmentParametersProto.ToEnvironmentResetParameters());
-                    return;
-                }
+                    {
+                        ResetCommandReceived?.Invoke(environmentParametersProto.ToEnvironmentResetParameters());
+                        return;
+                    }
                 default:
-                {
-                    return;
-                }
+                    {
+                        return;
+                    }
             }
         }
 
@@ -231,8 +232,16 @@ private void SendRLInputReceivedEvent(bool isTraining)
 
         #region Sending and retreiving data
 
-        public void PutObservations(
-            string brainKey, IEnumerable<Agent> agents)
+        /// <summary>
+        /// Sends the observations. If at least one brain has an agent in need of
+        /// a decision or if the academy is done, the data is sent via
+        /// Communicator. Else, a new step is realized. The data can only be
+        /// sent once all the brains that were part of initialization have tried
+        /// to send information.
+        /// </summary>
+        /// <param name="key">Batch Key.</param>
+        /// <param name="agents">Agent info.</param>
+        public void PutObservations(string brainKey, ICollection<Agent> agents)
         {
             // The brain tried called GiveBrainInfo, update m_hasQueried
             m_HasQueried[brainKey] = true;
@@ -409,7 +418,7 @@ private void HandleOnPlayModeChanged(PlayModeStateChange state)
             // This method is run whenever the playmode state is changed.
             if (state == PlayModeStateChange.ExitingPlayMode)
             {
-                Close();
+                Dispose();
             }
         }
 

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

@@ -108,7 +108,7 @@ Since the messages are sent back and forth with exchange and simultaneously when
     UnityOutput and UnityInput can be extended to provide functionalities beyond RL
     UnityRLOutput and UnityRLInput can be extended to provide new RL functionalities
      */
-    public interface ICommunicator
+    public interface ICommunicator : IBatchedDecisionMaker
     {
         /// <summary>
         /// Quit was received by the communicator.
@@ -141,27 +141,17 @@ public interface ICommunicator
         /// <param name="brainParameters">The Parameters for the Brain being registered</param>
         void SubscribeBrain(string name, BrainParameters brainParameters);
 
-        /// <summary>
-        /// Sends the observations. If at least one brain has an agent in need of
-        /// a decision or if the academy is done, the data is sent via
-        /// Communicator. Else, a new step is realized. The data can only be
-        /// sent once all the brains that were part of initialization have tried
-        /// to send information.
-        /// </summary>
-        /// <param name="key">Batch Key.</param>
-        /// <param name="agents">Agent info.</param>
-        void PutObservations(string key, IEnumerable<Agent> agents);
-
         /// <summary>
         /// Gets the AgentActions based on the batching key.
         /// </summary>
         /// <param name="key">A key to identify which actions to get</param>
         /// <returns></returns>
         Dictionary<Agent, AgentAction> GetActions(string key);
 
-        /// <summary>
-        /// Close the communicator gracefully on both sides of the communication.
-        /// </summary>
-        void Close();
+    }
+
+    public interface IBatchedDecisionMaker : IDisposable
+    {
+        void PutObservations(string key, ICollection<Agent> agents);
     }
 }