Refactor DemonstrationStore/Recorder

com.unity.ml-agents/Runtime/Agent.cs

@@ -10,7 +10,7 @@ namespace MLAgents
     /// Struct that contains all the information for an Agent, including its
     /// observations, actions and current status, that is sent to the Brain.
     /// </summary>
-    public struct AgentInfo
+    internal struct AgentInfo
     {
         /// <summary>
         /// Keeps track of the last vector action taken by the Brain.
@@ -159,7 +159,6 @@ internal struct AgentParameters
         /// Whether or not the agent requests a decision.
         bool m_RequestDecision;
 
-
         /// Keeps track of the number of steps taken by the agent in this episode.
         /// Note that this value is different for each agent, and may not overlap
         /// with the step counter in the Academy, since agents reset based on
@@ -178,9 +177,11 @@ internal struct AgentParameters
         ActionMasker m_ActionMasker;
 
         /// <summary>
-        /// Demonstration recorder.
+        /// Set of DemonstrationStores that the Agent will write its step information to.
+        /// If you use a DemonstrationRecorder component, this will automatically register its DemonstrationStore.
+        /// You can also add your own DemonstrationStore by calling DemonstrationRecorder.AddDemonstrationStoreToAgent()
         /// </summary>
-        DemonstrationRecorder m_Recorder;
+        internal ISet<DemonstrationStore> DemonstrationStores = new HashSet<DemonstrationStore>();
 
         /// <summary>
         /// List of sensors used to generate observations.
@@ -231,8 +232,6 @@ public void LazyInitialize()
             // Grab the "static" properties for the Agent.
             m_EpisodeId = EpisodeIdCounter.GetEpisodeId();
             m_PolicyFactory = GetComponent<BehaviorParameters>();
-            m_Recorder = GetComponent<DemonstrationRecorder>();
-
 
             m_Info = new AgentInfo();
             m_Action = new AgentAction();
@@ -253,6 +252,8 @@ public void LazyInitialize()
         /// becomes disabled or inactive.
         void OnDisable()
         {
+            DemonstrationStores.Clear();
+
             // If Academy.Dispose has already been called, we don't need to unregister with it.
             // We don't want to even try, because this will lazily create a new Academy!
             if (Academy.IsInitialized)
@@ -277,9 +278,10 @@ void NotifyAgentDone(bool maxStepReached = false)
             // We request a decision so Python knows the Agent is done immediately
             m_Brain?.RequestDecision(m_Info, sensors);
 
-            if (m_Recorder != null && m_Recorder.record && Application.isEditor)
+            // We also have to write any to any DemonstationStores so that they get the "done" flag.
+            foreach(var demoWriter in DemonstrationStores)
             {
-                m_Recorder.WriteExperience(m_Info, sensors);
+                demoWriter.Record(m_Info, sensors);
             }
 
             UpdateRewardStats();
@@ -522,9 +524,10 @@ void SendInfoToBrain()
 
             m_Brain.RequestDecision(m_Info, sensors);
 
-            if (m_Recorder != null && m_Recorder.record && Application.isEditor)
+            // If we have any DemonstrationStores, write the AgentInfo and sensors to them.
+            foreach(var demoWriter in DemonstrationStores)
             {
-                m_Recorder.WriteExperience(m_Info, sensors);
+                demoWriter.Record(m_Info, sensors);
             }
         }
 
@@ -536,6 +539,7 @@ void UpdateSensors()
             }
         }
 
+
         /// <summary>
         /// Collects the vector observations of the agent.
         /// The agent observation describes the current environment from the

com.unity.ml-agents/Runtime/Demonstration.cs

@@ -8,7 +8,7 @@ namespace MLAgents
     /// Used for imitation learning, or other forms of learning from data.
     /// </summary>
     [Serializable]
-    public class Demonstration : ScriptableObject
+    internal class Demonstration : ScriptableObject
     {
         public DemonstrationMetaData metaData;
         public BrainParameters brainParameters;
@@ -26,7 +26,7 @@ public void Initialize(BrainParameters brainParams,
     /// Kept in a struct for easy serialization and deserialization.
     /// </summary>
     [Serializable]
-    public class DemonstrationMetaData
+    internal class DemonstrationMetaData
     {
         public int numberExperiences;
         public int numberEpisodes;

com.unity.ml-agents/Runtime/DemonstrationRecorder.cs

@@ -1,7 +1,7 @@
 using System.IO.Abstractions;
 using System.Text.RegularExpressions;
 using UnityEngine;
-using System.Collections.Generic;
+using System.IO;
 
 namespace MLAgents
 {
@@ -12,47 +12,84 @@ namespace MLAgents
     [AddComponentMenu("ML Agents/Demonstration Recorder", (int)MenuGroup.Default)]
     public class DemonstrationRecorder : MonoBehaviour
     {
+        [Tooltip("Whether or not to record demonstrations.")]
         public bool record;
+
+        [Tooltip("Base demonstration file name. Will have numbers appended to make unique.")]
         public string demonstrationName;
-        string m_FilePath;
+
+        [Tooltip("Base directory to write the demo files. If null, will use {Application.dataPath}/Demonstrations.")]
+        public string demonstrationDirectory;
+
         DemonstrationStore m_DemoStore;
-        public const int MaxNameLength = 16;
+        internal const int MaxNameLength = 16;
 
-        void Start()
+        const string k_ExtensionType = ".demo";
+        IFileSystem m_FileSystem;
+
+        Agent m_Agent;
+
+        void OnEnable()
         {
-            if (Application.isEditor && record)
-            {
-                InitializeDemoStore();
-            }
+            m_Agent = GetComponent<Agent>();
         }
 
         void Update()
         {
-            if (Application.isEditor && record && m_DemoStore == null)
+            if (record)
             {
-                InitializeDemoStore();
+                LazyInitialize();
             }
         }
 
         /// <summary>
         /// Creates demonstration store for use in recording.
+        /// Has no effect if the demonstration store was already created.
         /// </summary>
-        public void InitializeDemoStore(IFileSystem fileSystem = null)
+        internal DemonstrationStore LazyInitialize(IFileSystem fileSystem = null)
         {
-            m_DemoStore = new DemonstrationStore(fileSystem);
+            if (m_DemoStore != null)
+            {
+                return m_DemoStore;
+            }
+
+            if (m_Agent == null)
+            {
+                m_Agent = GetComponent<Agent>();
+            }
+
+            m_FileSystem = fileSystem ?? new FileSystem();
             var behaviorParams = GetComponent<BehaviorParameters>();
+            if (string.IsNullOrEmpty(demonstrationName))
+            {
+                demonstrationName = behaviorParams.behaviorName;
+            }
+            if (string.IsNullOrEmpty(demonstrationDirectory))
+            {
+                demonstrationDirectory = Path.Combine(Application.dataPath, "Demonstrations");
+            }
+
             demonstrationName = SanitizeName(demonstrationName, MaxNameLength);
+            var filePath = MakeDemonstrationFilePath(m_FileSystem, demonstrationDirectory, demonstrationName);
+            var stream = m_FileSystem.File.Create(filePath);
+            m_DemoStore = new DemonstrationStore(stream);
+
             m_DemoStore.Initialize(
                 demonstrationName,
                 behaviorParams.brainParameters,
-                behaviorParams.fullyQualifiedBehaviorName);
+                behaviorParams.fullyQualifiedBehaviorName
+            );
+
+            AddDemonstrationStoreToAgent(m_DemoStore);
+
+            return m_DemoStore;
         }
 
         /// <summary>
         /// Removes all characters except alphanumerics from demonstration name.
         /// Shorten name if it is longer than the maxNameLength.
         /// </summary>
-        public static string SanitizeName(string demoName, int maxNameLength)
+        internal static string SanitizeName(string demoName, int maxNameLength)
         {
             var rgx = new Regex("[^a-zA-Z0-9 -]");
             demoName = rgx.Replace(demoName, "");
@@ -65,31 +102,78 @@ public static string SanitizeName(string demoName, int maxNameLength)
         }
 
         /// <summary>
-        /// Forwards AgentInfo to Demonstration Store.
+        /// Gets a unique path for the demonstrationName in the demonstrationDirectory.
         /// </summary>
-        public void WriteExperience(AgentInfo info, List<ISensor> sensors)
+        /// <param name="fileSystem"></param>
+        /// <param name="demonstrationDirectory"></param>
+        /// <param name="demonstrationName"></param>
+        /// <returns></returns>
+        internal static string MakeDemonstrationFilePath(
+            IFileSystem fileSystem, string demonstrationDirectory, string demonstrationName
+        )
         {
-            m_DemoStore?.Record(info, sensors);
+            // Create the directory if it doesn't already exist
+            if (!fileSystem.Directory.Exists(demonstrationDirectory))
+            {
+                fileSystem.Directory.CreateDirectory(demonstrationDirectory);
+            }
+
+            var literalName = demonstrationName;
+            var filePath = Path.Combine(demonstrationDirectory, literalName + k_ExtensionType);
+            var uniqueNameCounter = 0;
+            while (fileSystem.File.Exists(filePath))
+            {
+                // TODO should we use a timestamp instead of a counter here? This loops an increasing number of times
+                // as the number of demos increases.
+                literalName = demonstrationName + "_" + uniqueNameCounter;
+                filePath = Path.Combine(demonstrationDirectory, literalName + k_ExtensionType);
+                uniqueNameCounter++;
+            }
+
+            return filePath;
         }
 
+        /// <summary>
+        /// Close the DemonstrationStore and remove it from the Agent.
+        /// Has no effect if the DemonstrationStore is already closed (or wasn't opened)
+        /// </summary>
         public void Close()
         {
             if (m_DemoStore != null)
             {
+                RemoveDemonstrationStoreFromAgent(m_DemoStore);
+
                 m_DemoStore.Close();
                 m_DemoStore = null;
             }
         }
 
         /// <summary>
-        /// Closes Demonstration store.
+        /// Clean up the DemonstrationStore when shutting down or destroying the Agent.
         /// </summary>
-        void OnApplicationQuit()
+        void OnDestroy()
         {
-            if (Application.isEditor && record)
-            {
-                Close();
-            }
+            Close();
+        }
+
+        /// <summary>
+        /// Add additional DemonstrationStore to the Agent. It is still up to the user to Close this
+        /// DemonstrationStores when recording is done.
+        /// </summary>
+        /// <param name="demoStore"></param>
+        public void AddDemonstrationStoreToAgent(DemonstrationStore demoStore)
+        {
+            m_Agent.DemonstrationStores.Add(demoStore);
+        }
+
+        /// <summary>
+        /// Remove additional DemonstrationStore to the Agent. It is still up to the user to Close this
+        /// DemonstrationStores when recording is done.
+        /// </summary>
+        /// <param name="demoStore"></param>
+        public void RemoveDemonstrationStoreFromAgent(DemonstrationStore demoStore)
+        {
+            m_Agent.DemonstrationStores.Remove(demoStore);
         }
     }
 }