using System.IO.Abstractions; using System.Text.RegularExpressions; using UnityEngine; using System.IO; using Unity.MLAgents.Policies; using UnityEngine.Serialization; namespace Unity.MLAgents.Demonstrations { /// /// The Demonstration Recorder component facilitates the recording of demonstrations /// used for imitation learning. /// /// Add this component to the [GameObject] containing an /// to enable recording the agent for imitation learning. You must implement the /// function of the agent to provide manual control /// in order to record demonstrations. /// /// See [Imitation Learning - Recording Demonstrations] for more information. /// /// [GameObject]: https://docs.unity3d.com/Manual/GameObjects.html /// [Imitation Learning - Recording Demonstrations]: https://github.com/Unity-Technologies/ml-agents/blob/release_17_docs/docs//Learning-Environment-Design-Agents.md#recording-demonstrations /// [RequireComponent(typeof(Agent))] [AddComponentMenu("ML Agents/Demonstration Recorder", (int)MenuGroup.Default)] public class DemonstrationRecorder : MonoBehaviour { /// /// Whether or not to record demonstrations. /// [FormerlySerializedAs("record")] [Tooltip("Whether or not to record demonstrations.")] public bool Record; /// /// Number of steps to record. The editor will stop playing when it reaches this threshold. /// Set to zero to record indefinitely. /// [Tooltip("Number of steps to record. The editor will stop playing when it reaches this threshold. " + "Set to zero to record indefinitely.")] public int NumStepsToRecord; /// /// Base demonstration file name. If multiple files are saved, the additional filenames /// will have a sequence of unique numbers appended. /// [FormerlySerializedAs("demonstrationName")] [Tooltip("Base demonstration file name. If multiple files are saved, the additional " + "filenames will have a unique number appended.")] public string DemonstrationName; /// /// Directory to save the demo files. Will default to a "Demonstrations/" folder in the /// Application data path if not specified. /// [FormerlySerializedAs("demonstrationDirectory")] [Tooltip("Directory to save the demo files. Will default to " + "{Application.dataPath}/Demonstrations if not specified.")] public string DemonstrationDirectory; DemonstrationWriter m_DemoWriter; internal const int MaxNameLength = 16; const string k_ExtensionType = ".demo"; const string k_DefaultDirectoryName = "Demonstrations"; IFileSystem m_FileSystem; Agent m_Agent; void OnEnable() { m_Agent = GetComponent(); } void Update() { if (!Record) { return; } LazyInitialize(); // Quit when num steps to record is reached if (NumStepsToRecord > 0 && m_DemoWriter.NumSteps >= NumStepsToRecord) { Application.Quit(0); #if UNITY_EDITOR UnityEditor.EditorApplication.isPlaying = false; #endif } } /// /// Creates demonstration store for use in recording. /// Has no effect if the demonstration store was already created. /// internal DemonstrationWriter LazyInitialize(IFileSystem fileSystem = null) { if (m_DemoWriter != null) { return m_DemoWriter; } if (m_Agent == null) { m_Agent = GetComponent(); } m_FileSystem = fileSystem ?? new FileSystem(); var behaviorParams = GetComponent(); if (string.IsNullOrEmpty(DemonstrationName)) { DemonstrationName = behaviorParams.BehaviorName; } if (string.IsNullOrEmpty(DemonstrationDirectory)) { DemonstrationDirectory = Path.Combine(Application.dataPath, k_DefaultDirectoryName); } DemonstrationName = SanitizeName(DemonstrationName, MaxNameLength); var filePath = MakeDemonstrationFilePath(m_FileSystem, DemonstrationDirectory, DemonstrationName); var stream = m_FileSystem.File.Create(filePath); m_DemoWriter = new DemonstrationWriter(stream); AddDemonstrationWriterToAgent(m_DemoWriter); return m_DemoWriter; } /// /// Removes all characters except alphanumerics from demonstration name. /// Shorten name if it is longer than the maxNameLength. /// internal static string SanitizeName(string demoName, int maxNameLength) { var rgx = new Regex("[^a-zA-Z0-9 -]"); demoName = rgx.Replace(demoName, ""); // If the string is too long, it will overflow the metadata. if (demoName.Length > maxNameLength) { demoName = demoName.Substring(0, maxNameLength); } return demoName; } /// /// Gets a unique path for the DemonstrationName in the DemonstrationDirectory. /// /// /// /// /// internal static string MakeDemonstrationFilePath( IFileSystem fileSystem, string demonstrationDirectory, string demonstrationName ) { // 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; } /// /// Close the DemonstrationWriter and remove it from the Agent. /// Has no effect if the DemonstrationWriter is already closed (or wasn't opened) /// public void Close() { if (m_DemoWriter != null) { RemoveDemonstrationWriterFromAgent(m_DemoWriter); m_DemoWriter.Close(); m_DemoWriter = null; } } /// /// Clean up the DemonstrationWriter when shutting down or destroying the Agent. /// void OnDestroy() { Close(); } /// /// Add additional DemonstrationWriter to the Agent. It is still up to the user to Close this /// DemonstrationWriters when recording is done. /// /// public void AddDemonstrationWriterToAgent(DemonstrationWriter demoWriter) { var behaviorParams = GetComponent(); demoWriter.Initialize( DemonstrationName, behaviorParams.BrainParameters, behaviorParams.FullyQualifiedBehaviorName ); m_Agent.DemonstrationWriters.Add(demoWriter); } /// /// Remove additional DemonstrationWriter to the Agent. It is still up to the user to Close this /// DemonstrationWriters when recording is done. /// /// public void RemoveDemonstrationWriterFromAgent(DemonstrationWriter demoWriter) { m_Agent.DemonstrationWriters.Remove(demoWriter); } } }