[MLA-488] Fix observations in demonstration drawer

com.unity.ml-agents/Editor/DemonstrationDrawer.cs

@@ -1,3 +1,4 @@
+using System.Collections.Generic;
 using System.Text;
 using UnityEditor;
 using MLAgents.Demonstrations;
@@ -7,19 +8,21 @@
 namespace MLAgents.Editor
 {
     /// <summary>
-    /// Renders a custom UI for Demonstration Scriptable Object.
+    /// Renders a custom UI for DemonstrationSummary ScriptableObject.
     /// </summary>
-    [CustomEditor(typeof(Demonstration))]
+    [CustomEditor(typeof(DemonstrationSummary))]
     [CanEditMultipleObjects]
     internal class DemonstrationEditor : UnityEditor.Editor
     {
         SerializedProperty m_BrainParameters;
         SerializedProperty m_DemoMetaData;
+        SerializedProperty m_ObservationShapes;
 
         void OnEnable()
         {
             m_BrainParameters = serializedObject.FindProperty("brainParameters");
             m_DemoMetaData = serializedObject.FindProperty("metaData");
+            m_ObservationShapes = serializedObject.FindProperty("observationSummaries");
         }
 
         /// <summary>
@@ -28,25 +31,25 @@ void OnEnable()
         void MakeMetaDataProperty(SerializedProperty property)
         {
             var nameProp = property.FindPropertyRelative("demonstrationName");
-            var expProp = property.FindPropertyRelative("numberExperiences");
-            var epiProp = property.FindPropertyRelative("numberEpisodes");
-            var rewProp = property.FindPropertyRelative("meanReward");
+            var experiencesProp = property.FindPropertyRelative("numberExperiences");
+            var episodesProp = property.FindPropertyRelative("numberEpisodes");
+            var rewardsProp = property.FindPropertyRelative("meanReward");
 
             var nameLabel = nameProp.displayName + ": " + nameProp.stringValue;
-            var expLabel = expProp.displayName + ": " + expProp.intValue;
-            var epiLabel = epiProp.displayName + ": " + epiProp.intValue;
-            var rewLabel = rewProp.displayName + ": " + rewProp.floatValue;
+            var experiencesLabel = experiencesProp.displayName + ": " + experiencesProp.intValue;
+            var episodesLabel = episodesProp.displayName + ": " + episodesProp.intValue;
+            var rewardsLabel = rewardsProp.displayName + ": " + rewardsProp.floatValue;
 
             EditorGUILayout.LabelField(nameLabel);
-            EditorGUILayout.LabelField(expLabel);
-            EditorGUILayout.LabelField(epiLabel);
-            EditorGUILayout.LabelField(rewLabel);
+            EditorGUILayout.LabelField(experiencesLabel);
+            EditorGUILayout.LabelField(episodesLabel);
+            EditorGUILayout.LabelField(rewardsLabel);
         }
 
         /// <summary>
-        /// Constructs label for action size array.
+        /// Constructs label for a serialized integer array.
         /// </summary>
-        static string BuildActionArrayLabel(SerializedProperty actionSizeProperty)
+        static string BuildIntArrayLabel(SerializedProperty actionSizeProperty)
         {
             var actionSize = actionSizeProperty.arraySize;
             var actionLabel = new StringBuilder("[ ");
@@ -64,35 +67,62 @@ static string BuildActionArrayLabel(SerializedProperty actionSizeProperty)
         }
 
         /// <summary>
-        /// Renders Inspector UI for Brain Parameters of Demonstration.
+        /// Renders Inspector UI for BrainParameters of a DemonstrationSummary.
+        /// Only the Action size and type are used from the BrainParameters.
         /// </summary>
-        void MakeBrainParametersProperty(SerializedProperty property)
+        void MakeActionsProperty(SerializedProperty property)
         {
-            var vecObsSizeProp = property.FindPropertyRelative("vectorObservationSize");
-            var numStackedProp = property.FindPropertyRelative("numStackedVectorObservations");
             var actSizeProperty = property.FindPropertyRelative("vectorActionSize");
             var actSpaceTypeProp = property.FindPropertyRelative("vectorActionSpaceType");
 
-            var vecObsSizeLabel = vecObsSizeProp.displayName + ": " + vecObsSizeProp.intValue;
-            var numStackedLabel = numStackedProp.displayName + ": " + numStackedProp.intValue;
             var vecActSizeLabel =
-                actSizeProperty.displayName + ": " + BuildActionArrayLabel(actSizeProperty);
+                actSizeProperty.displayName + ": " + BuildIntArrayLabel(actSizeProperty);
             var actSpaceTypeLabel = actSpaceTypeProp.displayName + ": " +
                 (SpaceType)actSpaceTypeProp.enumValueIndex;
 
-            EditorGUILayout.LabelField(vecObsSizeLabel);
-            EditorGUILayout.LabelField(numStackedLabel);
             EditorGUILayout.LabelField(vecActSizeLabel);
             EditorGUILayout.LabelField(actSpaceTypeLabel);
         }
 
+        /// <summary>
+        /// Render the observation shapes of a DemonstrationSummary.
+        /// </summary>
+        /// <param name="obsSummariesProperty"></param>
+        void MakeObservationsProperty(SerializedProperty obsSummariesProperty)
+        {
+            var shapesLabels = new List<string>();
+            var numObservations = obsSummariesProperty.arraySize;
+            for (var i = 0; i < numObservations; i++)
+            {
+                var summary = obsSummariesProperty.GetArrayElementAtIndex(i);
+                var shapeProperty = summary.FindPropertyRelative("shape");
+                shapesLabels.Add(BuildIntArrayLabel(shapeProperty));
+            }
+
+            var shapeLabel = $"Shapes: {string.Join(",  ", shapesLabels)}";
+            EditorGUILayout.LabelField(shapeLabel);
+
+        }
+
         public override void OnInspectorGUI()
         {
             serializedObject.Update();
+
             EditorGUILayout.LabelField("Meta Data", EditorStyles.boldLabel);
+            EditorGUI.indentLevel++;
             MakeMetaDataProperty(m_DemoMetaData);
-            EditorGUILayout.LabelField("Brain Parameters", EditorStyles.boldLabel);
-            MakeBrainParametersProperty(m_BrainParameters);
+            EditorGUI.indentLevel--;
+
+            EditorGUILayout.LabelField("Observations", EditorStyles.boldLabel);
+            EditorGUI.indentLevel++;
+            MakeObservationsProperty(m_ObservationShapes);
+            EditorGUI.indentLevel--;
+
+            EditorGUILayout.LabelField("Actions", EditorStyles.boldLabel);
+            EditorGUI.indentLevel++;
+            MakeActionsProperty(m_BrainParameters);
+            EditorGUI.indentLevel--;
+
             serializedObject.ApplyModifiedProperties();
         }
     }

com.unity.ml-agents/Editor/DemonstrationImporter.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.IO;
 using MLAgents.CommunicatorObjects;
 using UnityEditor;
@@ -26,7 +27,7 @@ public override void OnImportAsset(AssetImportContext ctx)
 
             try
             {
-                // Read first two proto objects containing metadata and brain parameters.
+                // Read first three proto objects containing metadata, brain parameters, and observations.
                 Stream reader = File.OpenRead(ctx.assetPath);
 
                 var metaDataProto = DemonstrationMetaProto.Parser.ParseDelimitedFrom(reader);
@@ -36,17 +37,30 @@ public override void OnImportAsset(AssetImportContext ctx)
                 var brainParamsProto = BrainParametersProto.Parser.ParseDelimitedFrom(reader);
                 var brainParameters = brainParamsProto.ToBrainParameters();
 
+                // Read the first AgentInfoActionPair so that we can get the observation sizes.
+                List<ObservationSummary> observationSummaries;
+                try
+                {
+                    var agentInfoActionPairProto = AgentInfoActionPairProto.Parser.ParseDelimitedFrom(reader);
+                    observationSummaries = agentInfoActionPairProto.GetObservationSummaries();
+                }
+                catch
+                {
+                    // Just in case there weren't any AgentInfoActionPair or they couldn't be read.
+                    observationSummaries = new List<ObservationSummary>();
+                }
+
                 reader.Close();
 
-                var demonstration = ScriptableObject.CreateInstance<Demonstration>();
-                demonstration.Initialize(brainParameters, metaData);
-                userData = demonstration.ToString();
+                var demonstrationSummary = ScriptableObject.CreateInstance<DemonstrationSummary>();
+                demonstrationSummary.Initialize(brainParameters, metaData, observationSummaries);
+                userData = demonstrationSummary.ToString();
 
                 var texture = (Texture2D)
                     AssetDatabase.LoadAssetAtPath(k_IconPath, typeof(Texture2D));
 
-                ctx.AddObjectToAsset(ctx.assetPath, demonstration, texture);
-                ctx.SetMainObject(demonstration);
+                ctx.AddObjectToAsset(ctx.assetPath, demonstrationSummary, texture);
+                ctx.SetMainObject(demonstrationSummary);
             }
             catch
             {

com.unity.ml-agents/Runtime/Communicator/GrpcExtensions.cs

@@ -17,6 +17,7 @@ namespace MLAgents
 {
     internal static class GrpcExtensions
     {
+        #region AgentInfo
         /// <summary>
         /// Converts a AgentInfo to a protobuf generated AgentInfoActionPairProto
         /// </summary>
@@ -59,6 +60,24 @@ public static AgentInfoProto ToAgentInfoProto(this AgentInfo ai)
             return agentInfoProto;
         }
 
+        public static List<ObservationSummary> GetObservationSummaries(this AgentInfoActionPairProto infoActionPair)
+        {
+            List<ObservationSummary> summariesOut = new List<ObservationSummary>();
+            var agentInfo = infoActionPair.AgentInfo;
+            foreach (var obs in agentInfo.Observations)
+            {
+                var summary = new ObservationSummary();
+                summary.shape = obs.Shape.ToArray();
+                summariesOut.Add(summary);
+            }
+
+            return summariesOut;
+        }
+
+
+        #endregion
+
+        #region BrainParameters
         /// <summary>
         /// Converts a Brain into to a Protobuf BrainInfoProto so it can be sent
         /// </summary>
@@ -80,6 +99,25 @@ public static BrainParametersProto ToProto(this BrainParameters bp, string name,
             return brainParametersProto;
         }
 
+        /// <summary>
+        /// Convert a BrainParametersProto to a BrainParameters struct.
+        /// </summary>
+        /// <param name="bpp">An instance of a brain parameters protobuf object.</param>
+        /// <returns>A BrainParameters struct.</returns>
+        public static BrainParameters ToBrainParameters(this BrainParametersProto bpp)
+        {
+            var bp = new BrainParameters
+            {
+                vectorActionSize = bpp.VectorActionSize.ToArray(),
+                vectorActionDescriptions = bpp.VectorActionDescriptions.ToArray(),
+                vectorActionSpaceType = (SpaceType)bpp.VectorActionSpaceType
+            };
+            return bp;
+        }
+
+        #endregion
+
+        #region DemonstrationMetaData
         /// <summary>
         /// Convert metadata object to proto object.
         /// </summary>
@@ -114,22 +152,7 @@ public static DemonstrationMetaData ToDemonstrationMetaData(this DemonstrationMe
             }
             return dm;
         }
-
-        /// <summary>
-        /// Convert a BrainParametersProto to a BrainParameters struct.
-        /// </summary>
-        /// <param name="bpp">An instance of a brain parameters protobuf object.</param>
-        /// <returns>A BrainParameters struct.</returns>
-        public static BrainParameters ToBrainParameters(this BrainParametersProto bpp)
-        {
-            var bp = new BrainParameters
-            {
-                vectorActionSize = bpp.VectorActionSize.ToArray(),
-                vectorActionDescriptions = bpp.VectorActionDescriptions.ToArray(),
-                vectorActionSpaceType = (SpaceType)bpp.VectorActionSpaceType
-            };
-            return bp;
-        }
+        #endregion
 
         public static UnityRLInitParameters ToUnityRLInitParameters(this UnityRLInitializationInputProto inputProto)
         {
@@ -141,6 +164,7 @@ public static UnityRLInitParameters ToUnityRLInitParameters(this UnityRLInitiali
             };
         }
 
+        #region AgentAction
         public static AgentAction ToAgentAction(this AgentActionProto aap)
         {
             return new AgentAction
@@ -158,7 +182,9 @@ public static List<AgentAction> ToAgentActionList(this UnityRLInputProto.Types.L
             }
             return agentActions;
         }
+        #endregion
 
+        #region Observations
         public static ObservationProto ToProto(this Observation obs)
         {
             ObservationProto obsProto = null;
@@ -248,5 +274,6 @@ public static ObservationProto GetObservationProto(this ISensor sensor, WriteAda
             observationProto.Shape.AddRange(shape);
             return observationProto;
         }
+        #endregion
     }
 }

com.unity.ml-agents/Runtime/Demonstrations/DemonstrationMetaData.cs

@@ -0,0 +1,20 @@
+using System;
+using UnityEngine;
+using MLAgents.Policies;
+
+namespace MLAgents.Demonstrations
+{
+    /// <summary>
+    /// Demonstration meta-data.
+    /// Kept in a struct for easy serialization and deserialization.
+    /// </summary>
+    [Serializable]
+    internal class DemonstrationMetaData
+    {
+        public int numberExperiences;
+        public int numberEpisodes;
+        public float meanReward;
+        public string demonstrationName;
+        public const int ApiVersion = 1;
+    }
+}