Add ISensor.Update() (#2852)

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

 using NUnit.Framework;
 using UnityEngine;
 using System.IO.Abstractions.TestingHelpers;
+using System.Reflection;
+using MLAgents.CommunicatorObjects;
+using Google.Protobuf;
 
 namespace MLAgents.Tests
 {
 
             demoStore.Record(agentInfo);
             demoStore.Close();
+        }
+
+        public class ObservationAgent : TestAgent
+        {
+            public override void CollectObservations()
+            {
+                collectObservationsCalls += 1;
+                AddVectorObs(1f);
+                AddVectorObs(2f);
+                AddVectorObs(3f);
+            }
+        }
+
+        [Test]
+        public void TestAgentWrite()
+        {
+            var agentGo1 = new GameObject("TestAgent");
+            var bpA = agentGo1.AddComponent<BehaviorParameters>();
+            bpA.brainParameters.vectorObservationSize = 3;
+            bpA.brainParameters.numStackedVectorObservations = 1;
+            bpA.brainParameters.vectorActionDescriptions = new[] { "TestActionA", "TestActionB" };
+            bpA.brainParameters.vectorActionSize = new[] { 2, 2 };
+            bpA.brainParameters.vectorActionSpaceType = SpaceType.Discrete;
+
+            agentGo1.AddComponent<ObservationAgent>();
+            var agent1 = agentGo1.GetComponent<ObservationAgent>();
+
+            agentGo1.AddComponent<DemonstrationRecorder>();
+            var demoRecorder = agentGo1.GetComponent<DemonstrationRecorder>();
+            var fileSystem = new MockFileSystem();
+            demoRecorder.demonstrationName = "TestBrain";
+            demoRecorder.record = true;
+            demoRecorder.InitializeDemoStore(fileSystem);
+
+            var acaGo = new GameObject("TestAcademy");
+            acaGo.AddComponent<TestAcademy>();
+            var aca = acaGo.GetComponent<TestAcademy>();
+            aca.resetParameters = new ResetParameters();
+
+            var academyInitializeMethod = typeof(Academy).GetMethod("InitializeEnvironment",
+                BindingFlags.Instance | BindingFlags.NonPublic);
+            var agentEnableMethod = typeof(Agent).GetMethod("OnEnable",
+                BindingFlags.Instance | BindingFlags.NonPublic);
+            var agentSendInfo = typeof(Agent).GetMethod("SendInfo",
+                BindingFlags.Instance | BindingFlags.NonPublic);
+
+            agentEnableMethod?.Invoke(agent1, new object[] { });
+            academyInitializeMethod?.Invoke(aca, new object[] { });
+
+            // Step the agent
+            agent1.RequestDecision();
+            agentSendInfo?.Invoke(agent1, new object[] { });
+
+            demoRecorder.Close();
+
+            // Read back the demo file and make sure observations were written
+            var reader = fileSystem.File.OpenRead("Assets/Demonstrations/TestBrain.demo");
+            reader.Seek(DemonstrationStore.MetaDataBytes + 1, 0);
+            BrainParametersProto.Parser.ParseDelimitedFrom(reader);
+
+            var agentInfoProto = AgentInfoProto.Parser.ParseDelimitedFrom(reader);
+            var obs = agentInfoProto.StackedVectorObservation;
+            Assert.AreEqual(obs.Count, bpA.brainParameters.vectorObservationSize);
+            for (var i = 0; i < obs.Count; i++)
+            {
+                Assert.AreEqual((float) i+1, obs[i]);
+            }
+
         }
     }
 }

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

         {
             return sensorName;
         }
+
+        public void Update() { }
     }
 
     public class EditModeTestGeneration

UnitySDK/Assets/ML-Agents/Editor/Tests/Sensor/StackingSensorTests.cs

             wrapped.AddObservation(new [] {1f, 2f});
             SensorTestHelper.CompareObservation(sensor, new [] {0f, 0f, 0f, 0f, 1f, 2f});
 
+            sensor.Update();
+            sensor.Update();
+            sensor.Update();
+            sensor.Update();
+            SensorTestHelper.CompareObservation(sensor, new [] {5f, 6f, 7f, 8f, 9f, 10f});
+
+            // Check that if we don't call Update(), the same observations are produced
             SensorTestHelper.CompareObservation(sensor, new [] {5f, 6f, 7f, 8f, 9f, 10f});
         }
 

UnitySDK/Assets/ML-Agents/Editor/Tests/Sensor/VectorSensorTests.cs

             sensor.AddObservation(4f);
 
             SensorTestHelper.CompareObservation(sensor, new[] { 1f, 2f, 3f, 4f });
+            // Check that if we don't call Update(), the same observations are produced
+            SensorTestHelper.CompareObservation(sensor, new[] { 1f, 2f, 3f, 4f });
+
+            // Check that Update() clears the data
+            sensor.Update();
+            SensorTestHelper.CompareObservation(sensor, new[] { 0f, 0f, 0f, 0f });
+
         }
 
         [Test]

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

             m_Info.storedVectorActions = m_Action.vectorActions;
             m_Info.compressedObservations.Clear();
             m_ActionMasker.ResetMask();
+            UpdateSensors();
             using (TimerStack.Instance.Scoped("CollectObservations"))
             {
                 CollectObservations();
                 m_Recorder.WriteExperience(m_Info);
             }
 
+        }
+
+        void UpdateSensors()
+        {
+            for (var i = 0; i < sensors.Count; i++)
+            {
+                sensors[i].Update();
+            }
         }
 
         /// <summary>

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

+using System.IO.Abstractions;
+using System.Text.RegularExpressions;
-using System.Text.RegularExpressions;
 
 namespace MLAgents
 {
         /// <summary>
         /// Creates demonstration store for use in recording.
         /// </summary>
-        void InitializeDemoStore()
+        public void InitializeDemoStore(IFileSystem fileSystem = null)
-            m_DemoStore = new DemonstrationStore();
+            m_DemoStore = new DemonstrationStore(fileSystem);
+            var behaviorParams = GetComponent<BehaviorParameters>();
-                GetComponent<BehaviorParameters>().brainParameters,
-                GetComponent<BehaviorParameters>().behaviorName);
+                behaviorParams.brainParameters,
+                behaviorParams.behaviorName);
             Monitor.Log("Recording Demonstration of Agent: ", m_RecordingAgent.name);
         }
 
             m_DemoStore.Record(info);
         }
 
+        public void Close()
+        {
+            if (m_DemoStore != null)
+            {
+                m_DemoStore.Close();
+                m_DemoStore = null;
+            }
+        }
+
-            if (Application.isEditor && record && m_DemoStore != null)
+            if (Application.isEditor && record)
-                m_DemoStore.Close();
+                Close();
             }
         }
     }

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

 using System.IO;
 using System.IO.Abstractions;
 using Google.Protobuf;
+using UnityEngine;
 
 namespace MLAgents
 {
 
         public DemonstrationStore(IFileSystem fileSystem)
         {
-            m_FileSystem = fileSystem;
-        }
-
-        public DemonstrationStore()
-        {
-            m_FileSystem = new FileSystem();
+            if (fileSystem != null)
+            {
+                m_FileSystem = fileSystem;
+            }
+            else
+            {
+                m_FileSystem = new FileSystem();
+            }
         }
 
         /// <summary>

UnitySDK/Assets/ML-Agents/Scripts/Sensor/CameraSensor.cs

             }
         }
 
+        public void Update() { }
+
         public SensorCompressionType GetCompressionType()
         {
             return SensorCompressionType.PNG;

UnitySDK/Assets/ML-Agents/Scripts/Sensor/ISensor.cs

         /// <summary>
         /// Write the observation data directly to the WriteAdapter.
         /// This is considered an advanced interface; for a simpler approach, use SensorBase and override WriteFloats instead.
+        /// Note that this (and GetCompressedObservation) may be called multiple times per agent step, so should not
+        /// mutate any internal state.
         /// </summary>
         /// <param name="adapater"></param>
         /// <returns>The number of elements written</returns>
         /// </summary>
         /// <returns></returns>
         byte[] GetCompressedObservation();
+
+        /// <summary>
+        /// Update any internal state of the sensor. This is called once per each agent step.
+        /// </summary>
+        void Update();
 
         /// <summary>
         /// Return the compression type being used. If no compression is used, return SensorCompressionType.None

UnitySDK/Assets/ML-Agents/Scripts/Sensor/RenderTextureSensor.cs

             }
         }
 
+        public void Update() { }
+
         public SensorCompressionType GetCompressionType()
         {
             return SensorCompressionType.PNG;

UnitySDK/Assets/ML-Agents/Scripts/Sensor/SensorBase.cs

             return numFloats;
         }
 
+        public void Update() { }
+
         public virtual byte[] GetCompressedObservation()
         {
             return null;

UnitySDK/Assets/ML-Agents/Scripts/Sensor/StackingSensor.cs

                 numWritten += m_UnstackedObservationSize;
             }
 
-            // Finally update the index of the "current" buffer.
+            return numWritten;
+        }
+
+        /// <summary>
+        /// Updates the index of the "current" buffer.
+        /// </summary>
+        public void Update()
+        {
+            m_WrappedSensor.Update();
-            return numWritten;
         }
 
         public int[] GetFloatObservationShape()

UnitySDK/Assets/ML-Agents/Scripts/Sensor/VectorSensor.cs

                 }
             }
             adapter.AddRange(m_Observations);
-            Clear();
+        }
+
+        public void Update()
+        {
+            Clear();
         }
 
         public int[] GetFloatObservationShape()