write observations directly to protobuf (#3229)

* write observations directly to protobuf

* docstring and comment about Capacity

Author: Chris Elion

UnitySDK/Assets/ML-Agents/Editor/Tests/DemonstrationTests.cs

@@ -3,6 +3,7 @@
 using System.IO.Abstractions.TestingHelpers;
 using System.Reflection;
 using MLAgents.CommunicatorObjects;
+using MLAgents.Sensor;
 
 namespace MLAgents.Tests
 {
@@ -64,7 +65,7 @@ public void TestStoreInitalize()
                 storedVectorActions = new[] { 0f, 1f },
             };
 
-            demoStore.Record(agentInfo, new System.Collections.Generic.List<Sensor.Observation>());
+            demoStore.Record(agentInfo, new System.Collections.Generic.List<ISensor>());
             demoStore.Close();
         }
 

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

@@ -235,13 +235,6 @@ public AgentInfo Info
         /// </summary>
         public VectorSensor collectObservationsSensor;
 
-        /// <summary>
-        /// Internal buffer used for generating float observations.
-        /// </summary>
-        float[] m_VectorSensorBuffer;
-
-        WriteAdapter m_WriteAdapter = new WriteAdapter();
-
         /// MonoBehaviour function that is called when the attached GameObject
         /// becomes enabled or active.
         void OnEnable()
@@ -558,8 +551,6 @@ void SendInfoToBrain()
             }
             m_Info.actionMasks = m_ActionMasker.GetMask();
 
-            // var param = m_PolicyFactory.brainParameters; // look, no brain params!
-
             m_Info.reward = m_Reward;
             m_Info.done = m_Done;
             m_Info.maxStepReached = m_MaxStepReached;
@@ -569,19 +560,7 @@ void SendInfoToBrain()
 
             if (m_Recorder != null && m_Recorder.record && Application.isEditor)
             {
-
-                if (m_VectorSensorBuffer == null)
-                {
-                    // Create a buffer for writing uncompressed (i.e. float) sensor data to
-                    m_VectorSensorBuffer = new float[sensors.GetSensorFloatObservationSize()];
-                }
-
-                // This is a bit of a hack - if we're in inference mode, observations won't be generated
-                // But we need these to be generated for the recorder. So generate them here.
-                var observations = new List<Observation>();
-                GenerateSensorData(sensors, m_VectorSensorBuffer, m_WriteAdapter, observations);
-
-                m_Recorder.WriteExperience(m_Info, observations);
+                m_Recorder.WriteExperience(m_Info, sensors);
             }
 
         }

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

@@ -70,9 +70,9 @@ public static string SanitizeName(string demoName, int maxNameLength)
         /// <summary>
         /// Forwards AgentInfo to Demonstration Store.
         /// </summary>
-        public void WriteExperience(AgentInfo info, List<Observation> observations)
+        public void WriteExperience(AgentInfo info, List<ISensor> sensors)
         {
-            m_DemoStore.Record(info, observations);
+            m_DemoStore.Record(info, sensors);
         }
 
         public void Close()

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

@@ -20,6 +20,7 @@ public class DemonstrationStore
         DemonstrationMetaData m_MetaData;
         Stream m_Writer;
         float m_CumulativeReward;
+        WriteAdapter m_WriteAdapter = new WriteAdapter();
 
         public DemonstrationStore(IFileSystem fileSystem)
         {
@@ -92,7 +93,7 @@ void WriteBrainParameters(string brainName, BrainParameters brainParameters)
         /// <summary>
         /// Write AgentInfo experience to file.
         /// </summary>
-        public void Record(AgentInfo info, List<Observation> observations)
+        public void Record(AgentInfo info, List<ISensor> sensors)
         {
             // Increment meta-data counters.
             m_MetaData.numberExperiences++;
@@ -102,8 +103,13 @@ public void Record(AgentInfo info, List<Observation> observations)
                 EndEpisode();
             }
 
-            // Write AgentInfo to file.
-            var agentProto = info.ToInfoActionPairProto(observations);
+            // Generate observations and add AgentInfo to file.
+            var agentProto = info.ToInfoActionPairProto();
+            foreach (var sensor in sensors)
+            {
+                agentProto.AgentInfo.Observations.Add(sensor.GetObservationProto(m_WriteAdapter));
+            }
+
             agentProto.WriteDelimitedTo(m_Writer);
         }
 

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

@@ -16,9 +16,9 @@ public static class GrpcExtensions
         /// Converts a AgentInfo to a protobuf generated AgentInfoActionPairProto
         /// </summary>
         /// <returns>The protobuf version of the AgentInfoActionPairProto.</returns>
-        public static AgentInfoActionPairProto ToInfoActionPairProto(this AgentInfo ai, List<Observation> observations)
+        public static AgentInfoActionPairProto ToInfoActionPairProto(this AgentInfo ai)
         {
-            var agentInfoProto = ai.ToAgentInfoProto(observations);
+            var agentInfoProto = ai.ToAgentInfoProto();
 
             var agentActionProto = new AgentActionProto
             {
@@ -36,7 +36,7 @@ public static AgentInfoActionPairProto ToInfoActionPairProto(this AgentInfo ai,
         /// Converts a AgentInfo to a protobuf generated AgentInfoProto
         /// </summary>
         /// <returns>The protobuf version of the AgentInfo.</returns>
-        public static AgentInfoProto ToAgentInfoProto(this AgentInfo ai, List<Observation> observations)
+        public static AgentInfoProto ToAgentInfoProto(this AgentInfo ai)
         {
             var agentInfoProto = new AgentInfoProto
             {
@@ -51,14 +51,6 @@ public static AgentInfoProto ToAgentInfoProto(this AgentInfo ai, List<Observatio
                 agentInfoProto.ActionMask.AddRange(ai.actionMasks);
             }
 
-            if (observations != null)
-            {
-                foreach (var obs in observations)
-                {
-                    agentInfoProto.Observations.Add(obs.ToProto());
-                }
-            }
-
             return agentInfoProto;
         }
 
@@ -197,5 +189,49 @@ public static ObservationProto ToProto(this Observation obs)
             obsProto.Shape.AddRange(obs.Shape);
             return obsProto;
         }
+
+        /// <summary>
+        /// Generate an ObservationProto for the sensor using the provided WriteAdapter.
+        /// This is equivalent to producing an Observation and calling Observation.ToProto(),
+        /// but avoid some intermediate memory allocations.
+        /// </summary>
+        /// <param name="sensor"></param>
+        /// <param name="writeAdapter"></param>
+        /// <returns></returns>
+        public static ObservationProto GetObservationProto(this ISensor sensor, WriteAdapter writeAdapter)
+        {
+            var shape = sensor.GetObservationShape();
+            ObservationProto observationProto = null;
+            if (sensor.GetCompressionType() == SensorCompressionType.None)
+            {
+                var numFloats = sensor.ObservationSize();
+                var floatDataProto = new ObservationProto.Types.FloatData();
+                // Resize the float array
+                // TODO upgrade protobuf versions so that we can set the Capacity directly - see https://github.com/protocolbuffers/protobuf/pull/6530
+                for (var i = 0; i < numFloats; i++)
+                {
+                    floatDataProto.Data.Add(0.0f);
+                }
+
+                writeAdapter.SetTarget(floatDataProto.Data, sensor.GetObservationShape(), 0);
+                sensor.Write(writeAdapter);
+
+                observationProto = new ObservationProto
+                {
+                    FloatData = floatDataProto,
+                    CompressionType = (CompressionTypeProto)SensorCompressionType.None,
+                };
+            }
+            else
+            {
+                observationProto = new ObservationProto
+                {
+                    CompressedData = ByteString.CopyFrom(sensor.GetCompressedObservation()),
+                    CompressionType = (CompressionTypeProto)sensor.GetCompressionType(),
+                };
+            }
+            observationProto.Shape.AddRange(shape);
+            return observationProto;
+        }
     }
 }

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

@@ -35,8 +35,6 @@ public struct IdCallbackPair
 
         List<string> m_BehaviorNames = new List<string>();
         bool m_NeedCommunicateThisStep;
-        float[] m_VectorObservationBuffer = new float[0];
-        List<Observation> m_ObservationBuffer = new List<Observation>();
         WriteAdapter m_WriteAdapter = new WriteAdapter();
         Dictionary<string, SensorShapeValidator> m_SensorShapeValidators = new Dictionary<string, SensorShapeValidator>();
         Dictionary<string, List<IdCallbackPair>> m_ActionCallbacks = new Dictionary<string, List<IdCallbackPair>>();
@@ -239,18 +237,12 @@ public void DecideBatch()
         }
 
         /// <summary>
-        /// Sends the observations of one Agent. 
+        /// Sends the observations of one Agent.
         /// </summary>
         /// <param name="brainKey">Batch Key.</param>
         /// <param name="agent">Agent info.</param>
         public void PutObservations(string brainKey, AgentInfo info, List<ISensor> sensors, Action<AgentAction> action)
         {
-            int numFloatObservations = sensors.GetSensorFloatObservationSize();
-            if (m_VectorObservationBuffer.Length < numFloatObservations)
-            {
-                m_VectorObservationBuffer = new float[numFloatObservations];
-            }
-
 # if DEBUG
             if (!m_SensorShapeValidators.ContainsKey(brainKey))
             {
@@ -259,16 +251,21 @@ public void PutObservations(string brainKey, AgentInfo info, List<ISensor> senso
             m_SensorShapeValidators[brainKey].ValidateSensors(sensors);
 #endif
 
-            using (TimerStack.Instance.Scoped("GenerateSensorData"))
-            {
-                Agent.GenerateSensorData(sensors, m_VectorObservationBuffer, m_WriteAdapter, m_ObservationBuffer);
-            }
             using (TimerStack.Instance.Scoped("AgentInfo.ToProto"))
             {
-                var agentInfoProto = info.ToAgentInfoProto(m_ObservationBuffer);
+                var agentInfoProto = info.ToAgentInfoProto();
+
+                using (TimerStack.Instance.Scoped("GenerateSensorData"))
+                {
+                    foreach (var sensor in sensors)
+                    {
+                        var obsProto = sensor.GetObservationProto(m_WriteAdapter);
+                        agentInfoProto.Observations.Add(obsProto);
+                    }
+                }
                 m_CurrentUnityRlOutput.AgentInfos[brainKey].Value.Add(agentInfoProto);
             }
-            m_ObservationBuffer.Clear();
+
             m_NeedCommunicateThisStep = true;
             if (!m_ActionCallbacks.ContainsKey(brainKey))
             {
@@ -451,7 +448,7 @@ void UpdateSentBrainParameters(UnityRLInitializationOutputProto output)
         #region Handling side channels
 
         /// <summary>
-        /// Registers a side channel to the communicator. The side channel will exchange 
+        /// Registers a side channel to the communicator. The side channel will exchange
         /// messages with its Python equivalent.
         /// </summary>
         /// <param name="sideChannel"> The side channel to be registered.</param>