VectorSensor and StackedSensor (#2813)

* WIP VectorSensor and StackedSensor

* fix a few dumb mistakes

* more VectorSensor

* remove Update(), add util methods, hook into TensorGenerator

* WriteApdater to write to tensors and arrays

* write float observations

* used circular buffer for stacked obs

* cleanup

* fix unit tests

* docstrings

* undo accidental checkins

* rider suggestions, add range check

* bounds check before writing

* undo ProjectVersion.txt change

* fix unit tests

* unit test for VectorSensor

* StackingSensor tests

* missing meta file

* missing meta file

* WriteAdapter tests

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

                 done = true,
                 id = 5,
                 maxStepReached = true,
-                stackedVectorObservation = new List<float>() { 1f, 1f, 1f },
+                floatObservations = new List<float>() { 1f, 1f, 1f },
                 storedTextActions = "TestAction",
                 storedVectorActions = new[] { 0f, 1f },
                 textObservation = "TestAction",

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

 using System.Collections.Generic;
-using System.Linq;
+using System.Reflection;
+
-        static IEnumerable<Agent> GetFakeAgentInfos()
+        static IEnumerable<Agent> GetFakeAgents()
+            var acaGo = new GameObject("TestAcademy");
+            acaGo.AddComponent<TestAcademy>();
+            var aca = acaGo.GetComponent<TestAcademy>();
+            aca.resetParameters = new ResetParameters();
+
+            var bpA = goA.AddComponent<BehaviorParameters>();
+            bpA.brainParameters.vectorObservationSize = 3;
+            bpA.brainParameters.numStackedVectorObservations = 1;
+
+            var goB = new GameObject("goB");
+            var bpB = goB.AddComponent<BehaviorParameters>();
+            bpB.brainParameters.vectorObservationSize = 3;
+            bpB.brainParameters.numStackedVectorObservations = 1;
+            var agentB = goB.AddComponent<TestAgent>();
+
+            var agents = new List<Agent> { agentA, agentB };
+            foreach (var agent in agents)
+            {
+                var agentEnableMethod = typeof(Agent).GetMethod("OnEnableHelper",
+                    BindingFlags.Instance | BindingFlags.NonPublic);
+                agentEnableMethod?.Invoke(agent, new object[] { aca });
+            }
+            agentA.collectObservationsSensor.AddObservation(new Vector3(1, 2, 3));
+            agentB.collectObservationsSensor.AddObservation(new Vector3(4, 5, 6));
+
-                stackedVectorObservation = new[] { 1f, 2f, 3f }.ToList(),
-            var goB = new GameObject("goB");
-            var agentB = goB.AddComponent<TestAgent>();
+
-                stackedVectorObservation = new[] { 4f, 5f, 6f }.ToList(),
+
-            return new List<Agent> { agentA, agentB };
+            return agents;
         }
 
         [Test]
                 shape = new long[] { 2, 3 }
             };
             const int batchSize = 4;
-            var agentInfos = GetFakeAgentInfos();
+            var agentInfos = GetFakeAgents();
+            generator.AddSensorIndex(0);
+            generator.AddSensorIndex(1);
+            generator.AddSensorIndex(2);
             generator.Generate(inputTensor, batchSize, agentInfos);
             Assert.IsNotNull(inputTensor.data);
             Assert.AreEqual(inputTensor.data[0, 0], 1);
                 valueType = TensorProxy.TensorType.Integer
             };
             const int batchSize = 4;
-            var agentInfos = GetFakeAgentInfos();
+            var agentInfos = GetFakeAgents();
             var alloc = new TensorCachingAllocator();
             var generator = new PreviousActionInputGenerator(alloc);
 
                 valueType = TensorProxy.TensorType.FloatingPoint
             };
             const int batchSize = 4;
-            var agentInfos = GetFakeAgentInfos();
+            var agentInfos = GetFakeAgents();
             var alloc = new TensorCachingAllocator();
             var generator = new ActionMaskInputGenerator(alloc);
             generator.Generate(inputTensor, batchSize, agentInfos);

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

 using NUnit.Framework;
 using System.Reflection;
 using MLAgents.Sensor;
-using MLAgents.InferenceBrain;
 
 namespace MLAgents.Tests
 {
 
         public int[] GetFloatObservationShape()
         {
-            return new[] { 1 };
+            return new[] { 0 };
-        public void WriteToTensor(TensorProxy tensorProxy, int agentIndex) { }
+        public int Write(WriteAdapter adapter)
+        {
+            // No-op
+            return 0;
+        }
 
         public byte[] GetCompressedObservation()
         {

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

             var go = new GameObject("MyGameObject");
             var rayPer3D = go.AddComponent<RayPerception3D>();
             var result = rayPer3D.Perceive(1f, angles, tags);
-            Debug.Log(result.Count);
             Assert.IsTrue(result.Count == angles.Length * (tags.Length + 2));
         }
 

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

             string[] scenes = { "Assets/ML-Agents/Examples/3DBall/Scenes/3DBall.unity" };
             var buildResult = BuildPipeline.BuildPlayer(scenes, "testPlayer", BuildTarget.StandaloneOSX, BuildOptions.None);
 #if UNITY_2018_1_OR_NEWER
-            var isOK = buildResult.summary.result == BuildResult.Succeeded;
+            var isOk = buildResult.summary.result == BuildResult.Succeeded;
             var error = "";
             foreach (var stepInfo in buildResult.steps)
             {
             }
 #else
             var error = buildResult;
-            var isOK = string.IsNullOrEmpty(error);
+            var isOk = string.IsNullOrEmpty(error);
-            if (isOK)
+            if (isOk)
             {
                 EditorApplication.Exit(0);
             }

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

     public struct AgentInfo
     {
         /// <summary>
-        /// Most recent agent vector (i.e. numeric) observation.
-        /// </summary>
-        public List<float> vectorObservation;
-
-        /// <summary>
-        /// The previous agent vector observations, stacked. The length of the
-        /// history (i.e. number of vector observations to stack) is specified
-        /// in the Brain parameters.
-        /// </summary>
-        public List<float> stackedVectorObservation;
-
-        /// <summary>
+        // TODO struct?
+        public List<float> floatObservations;
+
         /// <summary>
         /// Most recent text observation.
         /// </summary>
         [FormerlySerializedAs("m_Sensors")]
         public List<ISensor> sensors;
 
+        public VectorSensor collectObservationsSensor;
+
+        WriteAdapter m_WriteAdapter = new WriteAdapter();
+
         /// MonoBehaviour function that is called when the attached GameObject
         /// becomes enabled or active.
         void OnEnable()
             if (m_Info.textObservation == null)
                 m_Info.textObservation = "";
             m_Action.textActions = "";
-            m_Info.vectorObservation =
-                new List<float>(param.vectorObservationSize);
-            m_Info.stackedVectorObservation =
-                new List<float>(param.vectorObservationSize
-                    * param.numStackedVectorObservations);
-            m_Info.stackedVectorObservation.AddRange(
-                new float[param.vectorObservationSize
-                          * param.numStackedVectorObservations]);
+            m_Info.floatObservations = new List<float>();
+            m_Info.floatObservations.AddRange(
+                new float[param.vectorObservationSize
+                    * param.numStackedVectorObservations]);
             m_Info.customObservation = null;
         }
 
         /// </summary>
         public void InitializeSensors()
         {
+            // Get all attached sensor components
             var attachedSensorComponents = GetComponents<SensorComponent>();
             sensors.Capacity += attachedSensorComponents.Length;
             foreach (var component in attachedSensorComponents)
 
+            // Support legacy CollectObservations
+            var param = m_PolicyFactory.brainParameters;
+            if (param.vectorObservationSize > 0)
+            {
+                collectObservationsSensor = new VectorSensor(param.vectorObservationSize);
+                if (param.numStackedVectorObservations > 1)
+                {
+                    var stackingSensor = new StackingSensor(collectObservationsSensor, param.numStackedVectorObservations);
+                    sensors.Add(stackingSensor);
+                }
+                else
+                {
+                    sensors.Add(collectObservationsSensor);
+                }
+            }
+
             // Sort the Sensors by name to ensure determinism
             sensors.Sort((x, y) => x.GetName().CompareTo(y.GetName()));
 
 
             m_Info.storedVectorActions = m_Action.vectorActions;
             m_Info.storedTextActions = m_Action.textActions;
-            m_Info.vectorObservation.Clear();
             m_Info.compressedObservations.Clear();
             m_ActionMasker.ResetMask();
             using (TimerStack.Instance.Scoped("CollectObservations"))
             m_Info.actionMasks = m_ActionMasker.GetMask();
 
-            var param = m_PolicyFactory.brainParameters;
-            if (m_Info.vectorObservation.Count != param.vectorObservationSize)
-            {
-                throw new UnityAgentsException(string.Format(
-                    "Vector Observation size mismatch in continuous " +
-                    "agent {0}. " +
-                    "Was Expecting {1} but received {2}. ",
-                    gameObject.name,
-                    param.vectorObservationSize,
-                    m_Info.vectorObservation.Count));
-            }
-
-            Utilities.ShiftLeft(m_Info.stackedVectorObservation, param.vectorObservationSize);
-            Utilities.ReplaceRange(m_Info.stackedVectorObservation, m_Info.vectorObservation,
-                m_Info.stackedVectorObservation.Count - m_Info.vectorObservation.Count);
+            // var param = m_PolicyFactory.brainParameters; // look, no brain params!
 
             m_Info.reward = m_Reward;
             m_Info.done = m_Done;
         /// </summary>
         public void GenerateSensorData()
         {
+
+            int floatsWritten = 0;
-            // TODO add bool argument indicating when to compress? For now, we always will compress.
-                var compressedObs = new CompressedObservation
+                if (sensor.GetCompressionType() == SensorCompressionType.None)
+                {
+                    m_WriteAdapter.SetTarget(m_Info.floatObservations, floatsWritten);
+                    floatsWritten += sensor.Write(m_WriteAdapter);
+                }
+                else
-                    Data = sensor.GetCompressedObservation(),
-                    Shape = sensor.GetFloatObservationShape(),
-                    CompressionType = sensor.GetCompressionType()
-                };
-                m_Info.compressedObservations.Add(compressedObs);
+                    var compressedObs = new CompressedObservation
+                    {
+                        Data = sensor.GetCompressedObservation(),
+                        Shape = sensor.GetFloatObservationShape(),
+                        CompressionType = sensor.GetCompressionType()
+                    };
+                    m_Info.compressedObservations.Add(compressedObs);
+                }
             }
         }
 
         /// <param name="observation">Observation.</param>
         protected void AddVectorObs(float observation)
         {
-            m_Info.vectorObservation.Add(observation);
+            collectObservationsSensor.AddObservation(observation);
         }
 
         /// <summary>
         /// <param name="observation">Observation.</param>
         protected void AddVectorObs(int observation)
         {
-            m_Info.vectorObservation.Add(observation);
+            collectObservationsSensor.AddObservation(observation);
         }
 
         /// <summary>
         /// <param name="observation">Observation.</param>
         protected void AddVectorObs(Vector3 observation)
         {
-            m_Info.vectorObservation.Add(observation.x);
-            m_Info.vectorObservation.Add(observation.y);
-            m_Info.vectorObservation.Add(observation.z);
+            collectObservationsSensor.AddObservation(observation);
         }
 
         /// <summary>
         /// <param name="observation">Observation.</param>
         protected void AddVectorObs(Vector2 observation)
         {
-            m_Info.vectorObservation.Add(observation.x);
-            m_Info.vectorObservation.Add(observation.y);
+            collectObservationsSensor.AddObservation(observation);
         }
 
         /// <summary>
         /// <param name="observation">Observation.</param>
         protected void AddVectorObs(IEnumerable<float> observation)
         {
-            m_Info.vectorObservation.AddRange(observation);
+            collectObservationsSensor.AddObservation(observation);
         }
 
         /// <summary>
         /// <param name="observation">Observation.</param>
         protected void AddVectorObs(Quaternion observation)
         {
-            m_Info.vectorObservation.Add(observation.x);
-            m_Info.vectorObservation.Add(observation.y);
-            m_Info.vectorObservation.Add(observation.z);
-            m_Info.vectorObservation.Add(observation.w);
+            collectObservationsSensor.AddObservation(observation);
         }
 
         /// <summary>
         /// <param name="observation"></param>
         protected void AddVectorObs(bool observation)
         {
-            m_Info.vectorObservation.Add(observation ? 1f : 0f);
+            collectObservationsSensor.AddObservation(observation);
-            var oneHotVector = new float[range];
-            oneHotVector[observation] = 1;
-            m_Info.vectorObservation.AddRange(oneHotVector);
+            collectObservationsSensor.AddOneHotObservation(observation, range);
         }
 
         /// <summary>

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

         {
             var agentInfoProto = new AgentInfoProto
             {
-                StackedVectorObservation = { ai.stackedVectorObservation },
+                StackedVectorObservation = { ai.floatObservations },
                 StoredVectorActions = { ai.storedVectorActions },
                 StoredTextActions = ai.storedTextActions,
                 TextObservation = ai.textObservation,

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/GeneratorImpl.cs

 using System;
 using Barracuda;
 using MLAgents.InferenceBrain.Utils;
+using MLAgents.Sensor;
+using UnityEngine;
 
 namespace MLAgents.InferenceBrain
 {
     public class VectorObservationGenerator : TensorGenerator.IGenerator
     {
         readonly ITensorAllocator m_Allocator;
+        List<int> m_SensorIndices = new List<int>();
+        WriteAdapter m_WriteAdapter = new WriteAdapter();
+
+        public void AddSensorIndex(int sensorIndex)
+        {
+            m_SensorIndices.Add(sensorIndex);
+        }
+
         public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<Agent> agents)
         {
             TensorUtils.ResizeTensor(tensorProxy, batchSize, m_Allocator);
             {
-                var info = agent.Info;
-                var vectorObs = info.stackedVectorObservation;
-                for (var j = 0; j < vecObsSizeT; j++)
+                var tensorOffset = 0;
+                // Write each sensor consecutively to the tensor
+                foreach (var sensorIndex in m_SensorIndices)
-                    tensorProxy.data[agentIndex, j] = vectorObs[j];
+                    m_WriteAdapter.SetTarget(tensorProxy, agentIndex, tensorOffset);
+                    var sensor = agent.sensors[sensorIndex];
+                    var numWritten = sensor.Write(m_WriteAdapter);
+                    tensorOffset += numWritten;
+                Debug.AssertFormat(
+                    tensorOffset == vecObsSizeT,
+                    "mismatch between vector observation size ({0}) and number of observations written ({1})",
+                    vecObsSizeT, tensorOffset
+                );
+
                 agentIndex++;
             }
         }
     /// </summary>
     public class VisualObservationInputGenerator : TensorGenerator.IGenerator
     {
-        readonly int m_Index;
-        readonly bool m_GrayScale;
+        readonly int m_SensorIndex;
+        WriteAdapter m_WriteAdapter = new WriteAdapter();
-            int index, ITensorAllocator allocator)
+            int sensorIndex, ITensorAllocator allocator)
-            m_Index = index;
+            m_SensorIndex = sensorIndex;
             m_Allocator = allocator;
         }
 
             var agentIndex = 0;
             foreach (var agent in agents)
             {
-                // TODO direct access to sensors list here - should we do it differently?
-                // TODO m_Index here is the visual observation index. Will work for now but not if we add more sensor types.
-                agent.sensors[m_Index].WriteToTensor(tensorProxy, agentIndex);
+                m_WriteAdapter.SetTarget(tensorProxy, agentIndex, 0);
+                agent.sensors[m_SensorIndex].Write(m_WriteAdapter);
                 agentIndex++;
             }
         }

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/ModelRunner.cs

                 // Just grab the first agent in the collection (any will suffice, really).
                 // We check for an empty Collection above, so this will always return successfully.
                 var firstAgent = m_Agents[0];
-                m_TensorGenerator.InitializeVisualObservations(firstAgent, m_TensorAllocator);
+                m_TensorGenerator.InitializeObservations(firstAgent, m_TensorAllocator);
                 m_VisualObservationsInitialized = true;
             }
 

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorGenerator.cs

                 new BatchSizeGenerator(allocator);
             m_Dict[TensorNames.SequenceLengthPlaceholder] =
                 new SequenceLengthGenerator(allocator);
-            m_Dict[TensorNames.VectorObservationPlacholder] =
-                new VectorObservationGenerator(allocator);
 
             if (barracudaModel != null)
             {
             m_Dict[TensorNames.ValueEstimateOutput] = new BiDimensionalOutputGenerator(allocator);
         }
 
-        public void InitializeVisualObservations(Agent agent, ITensorAllocator allocator)
+        public void InitializeObservations(Agent agent, ITensorAllocator allocator)
-            for (var visIndex = 0; visIndex < agent.sensors.Count; visIndex++)
+            // Loop through the sensors on a representative agent.
+            // For vector observations, add the index to the (single) VectorObservationGenerator
+            // For visual observations, make a VisualObservationInputGenerator
+            var visIndex = 0;
+            VectorObservationGenerator vecObsGen = null;
+            for (var sensorIndex = 0; sensorIndex < agent.sensors.Count; sensorIndex++)
-                // TODO handle non-visual Sensors too - need to index better
-                m_Dict[TensorNames.VisualObservationPlaceholderPrefix + visIndex] =
-                    new VisualObservationInputGenerator(visIndex, allocator);
+                var sensor = agent.sensors[sensorIndex];
+                var shape = sensor.GetFloatObservationShape();
+                // TODO generalize - we currently only have vector or visual, but can't handle "2D" observations
+                var isVectorSensor = (shape.Length == 1);
+                if (isVectorSensor)
+                {
+                    if (vecObsGen == null)
+                    {
+                        vecObsGen = new VectorObservationGenerator(allocator);
+                    }
+
+                    vecObsGen.AddSensorIndex(sensorIndex);
+                }
+                else
+                {
+                    m_Dict[TensorNames.VisualObservationPlaceholderPrefix + visIndex] =
+                        new VisualObservationInputGenerator(sensorIndex, allocator);
+                    visIndex++;
+                }
+            }
+
+            if (vecObsGen != null)
+            {
+                m_Dict[TensorNames.VectorObservationPlacholder] = vecObsGen;
             }
         }
 

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

 using System;
-using MLAgents.InferenceBrain;
 using UnityEngine;
 
 namespace MLAgents.Sensor
             }
         }
 
-        public void WriteToTensor(TensorProxy tensorProxy, int agentIndex)
+        public int Write(WriteAdapter adapter)
-                Utilities.TextureToTensorProxy(texture, tensorProxy, m_Grayscale, agentIndex);
+                var numWritten = Utilities.TextureToTensorProxy(texture, adapter, m_Grayscale);
+                return numWritten;
             }
         }
 

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

-using MLAgents.InferenceBrain;
-
 namespace MLAgents.Sensor
 {
     public enum SensorCompressionType
         int[] GetFloatObservationShape();
 
         /// <summary>
-        /// Write the observation data directly to the TensorProxy.
+        /// Write the observation data directly to the WriteAdapter.
-        /// <param name="tensorProxy"></param>
-        /// <param name="agentIndex"></param>
-        void WriteToTensor(TensorProxy tensorProxy, int agentIndex);
+        /// <param name="adapater"></param>
+        /// <returns>The number of elements written</returns>
+        int Write(WriteAdapter adapater);
 
         /// <summary>
         /// Return a compressed representation of the observation. For small observations, this should generally not be

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

 using System;
-using MLAgents.InferenceBrain;
-    class RenderTextureSensor : ISensor
+    public class RenderTextureSensor : ISensor
     {
         RenderTexture m_RenderTexture;
         int m_Width;
             }
         }
 
-        public void WriteToTensor(TensorProxy tensorProxy, int index)
+        public int Write(WriteAdapter adapter)
-                Utilities.TextureToTensorProxy(texture, tensorProxy, m_Grayscale, index);
+                var numWritten = Utilities.TextureToTensorProxy(texture, adapter, m_Grayscale);
+                return numWritten;
             }
         }
 

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

-using MLAgents.InferenceBrain;
 using UnityEngine;
 
 namespace MLAgents.Sensor
         public abstract string GetName();
 
         /// <summary>
-        /// Default implementation of WriteToTensor interface. This creates a temporary array, calls WriteObservation,
-        /// and then writes the results to the TensorProxy.
+        /// Default implementation of Write interface. This creates a temporary array, calls WriteObservation,
+        /// and then writes the results to the WriteAdapter.
-        /// <param name="tensorProxy"></param>
-        /// <param name="agentIndex"></param>
-        public virtual void WriteToTensor(TensorProxy tensorProxy, int agentIndex)
+        /// <param name="adapter"></param>
+        public virtual int Write(WriteAdapter adapter)
         {
             // TODO reuse buffer for similar agents, don't call GetFloatObservationShape()
             int[] shape = GetFloatObservationShape();
             float[] buffer = new float[numFloats];
             WriteObservation(buffer);
 
-            for (var i = 0; i < numFloats; i++)
-            {
-                tensorProxy.data[agentIndex, i] = buffer[i];
-            }
+            adapter.AddRange(buffer);
+
+            return numFloats;
         }
 
         public virtual byte[] GetCompressedObservation()

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

 using UnityEngine;
 using System.Collections.Generic;
-using MLAgents.InferenceBrain;
+using MLAgents.Sensor;
-        /// <summary>
-        /// Converts a list of Texture2D into a TensorProxy.
-        /// </summary>
-        /// <param name="textures">
-        /// The list of textures to be put into the tensor.
-        /// Note that the textures must have same width and height.
-        /// </param>
-        /// <param name="tensorProxy">
-        /// TensorProxy to fill with Texture data.
-        /// </param>
-        /// <param name="grayScale">
-        /// If set to <c>true</c> the textures will be converted to grayscale before
-        /// being stored in the tensor.
-        /// </param>
-        public static void TextureToTensorProxy(
-            List<Texture2D> textures,
-            TensorProxy tensorProxy,
-            bool grayScale)
-        {
-            var numTextures = textures.Count;
-            var width = textures[0].width;
-            var height = textures[0].height;
-
-            for (var t = 0; t < numTextures; t++)
-            {
-                var texture = textures[t];
-                Debug.Assert(width == texture.width, "All Textures must have the same dimension");
-                Debug.Assert(height == texture.height, "All Textures must have the same dimension");
-                TextureToTensorProxy(texture, tensorProxy, grayScale, t);
-            }
-        }
-        /// Puts a Texture2D into a TensorProxy.
+        /// Puts a Texture2D into a WriteAdapter.
-        /// <param name="tensorProxy">
-        /// TensorProxy to fill with Texture data.
+        /// <param name="adapter">
+        /// Adapter to fill with Texture data.
-        /// <param name="textureOffset">
-        /// Index of the texture being written.
-        /// </param>
-        public static void TextureToTensorProxy(
+        /// <returns>The number of floats written</returns>
+        public static int TextureToTensorProxy(
-            TensorProxy tensorProxy,
-            bool grayScale,
-            int textureOffset = 0)
+            WriteAdapter adapter,
+            bool grayScale)
-            var data = tensorProxy.data;
-            var t = textureOffset;
             var texturePixels = texture.GetPixels32();
             // During training, we convert from Texture to PNG before sending to the trainer, which has the
             // effect of flipping the image. We need another flip here at inference time to match this.
                     var currentPixel = texturePixels[(height - h - 1) * width + w];
                     if (grayScale)
                     {
-                        data[t, h, w, 0] =
+                        adapter[h, w, 0] =
-                        data[t, h, w, 0] = currentPixel.r / 255.0f;
-                        data[t, h, w, 1] = currentPixel.g / 255.0f;
-                        data[t, h, w, 2] = currentPixel.b / 255.0f;
+                        adapter[h, w, 0] = currentPixel.r / 255.0f;
+                        adapter[h, w, 1] = currentPixel.g / 255.0f;
+                        adapter[h, w, 2] = currentPixel.b / 255.0f;
+            return height * width * (grayScale ? 1 : 3);
         }
 
         /// <summary>

UnitySDK/Assets/ML-Agents/Editor/Tests/Sensor.meta

+fileFormatVersion: 2
+guid: 1b196836e6e3a4361bc62265ec88ebed
+folderAsset: yes
+DefaultImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

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

+namespace MLAgents.Sensor
+{
+    /// <summary>
+    /// Sensor that wraps around another Sensor to provide temporal stacking.
+    /// Conceptually, consecutive observations are stored left-to-right, which is how they're output
+    /// For example, 4 stacked sets of observations would be output like
+    ///   |  t = now - 3  |  t = now -3  |  t = now - 2  |  t = now  |
+    /// Internally, a circular buffer of arrays is used. The m_CurrentIndex represents the most recent observation.
+    /// </summary>
+    public class StackingSensor : ISensor
+    {
+        /// <summary>
+        /// The wrapped sensor.
+        /// </summary>
+        ISensor m_WrappedSensor;
+
+        /// <summary>
+        /// Number of stacks to save
+        /// </summary>
+        int m_NumStackedObservations;
+        int m_UnstackedObservationSize;
+
+        string m_Name;
+        int[] m_Shape;
+
+        /// <summary>
+        /// Buffer of previous observations
+        /// </summary>
+        float[][] m_StackedObservations;
+
+        int m_CurrentIndex;
+        WriteAdapter m_LocalAdapter = new WriteAdapter();
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="wrapped">The wrapped sensor</param>
+        /// <param name="numStackedObservations">Number of stacked observations to keep</param>
+        public StackingSensor(ISensor wrapped, int numStackedObservations)
+        {
+            // TODO ensure numStackedObservations > 1
+            m_WrappedSensor = wrapped;
+            m_NumStackedObservations = numStackedObservations;
+
+            m_Name = $"StackingSensor_size{numStackedObservations}_{wrapped.GetName()}";
+
+            var shape = wrapped.GetFloatObservationShape();
+            m_Shape = new int[shape.Length];
+
+            m_UnstackedObservationSize = 1;
+            for (int d = 0; d < shape.Length; d++)
+            {
+                m_Shape[d] = shape[d];
+                m_UnstackedObservationSize *= shape[d];
+            }
+
+            // TODO support arbitrary stacking dimension
+            m_Shape[0] *= numStackedObservations;
+            m_StackedObservations = new float[numStackedObservations][];
+            for (var i = 0; i < numStackedObservations; i++)
+            {
+                m_StackedObservations[i] = new float[m_UnstackedObservationSize];
+            }
+        }
+
+        public int Write(WriteAdapter adapter)
+        {
+            // First, call the wrapped sensor's write method. Make sure to use our own adapater, not the passed one.
+            m_LocalAdapter.SetTarget(m_StackedObservations[m_CurrentIndex], 0);
+            m_WrappedSensor.Write(m_LocalAdapter);
+
+            // Now write the saved observations (oldest first)
+            var numWritten = 0;
+            for (var i = 0; i < m_NumStackedObservations; i++)
+            {
+                var obsIndex = (m_CurrentIndex + 1 + i) % m_NumStackedObservations;
+                adapter.AddRange(m_StackedObservations[obsIndex], numWritten);
+                numWritten += m_UnstackedObservationSize;
+            }
+
+            // Finally update the index of the "current" buffer.
+            m_CurrentIndex = (m_CurrentIndex + 1) % m_NumStackedObservations;
+            return numWritten;
+        }
+
+        public int[] GetFloatObservationShape()
+        {
+            return m_Shape;
+        }
+
+        public string GetName()
+        {
+            return m_Name;
+        }
+
+        public virtual byte[] GetCompressedObservation()
+        {
+            return null;
+        }
+
+        public virtual SensorCompressionType GetCompressionType()
+        {
+            return SensorCompressionType.None;
+        }
+
+        // TODO support stacked compressed observations (byte stream)
+
+    }
+}

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

+fileFormatVersion: 2
+guid: 8b7a6e88d47d4438ad67e1862566462c
+timeCreated: 1572299581

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

+using System.Collections.Generic;
+using UnityEngine;
+
+namespace MLAgents.Sensor
+{
+    public class VectorSensor : ISensor
+    {
+        // TODO use float[] instead
+        // TOOD allow setting float[]
+        List<float> m_Observations;
+        int[] m_Shape;
+        string m_Name;
+
+        public VectorSensor(int observationSize, string name = null)
+        {
+            if (name == null)
+            {
+                name = $"VectorSensor_size{observationSize}";
+            }
+
+            m_Observations = new List<float>(observationSize);
+            m_Name = name;
+            m_Shape = new[] { observationSize };
+        }
+
+        public int Write(WriteAdapter adapter)
+        {
+            var expectedObservations = m_Shape[0];
+            if (m_Observations.Count > expectedObservations)
+            {
+                // Too many observations, truncate
+                Debug.LogWarningFormat(
+                    "More observations ({0}) made than vector observation size ({1}). The observations will be truncated.",
+                    m_Observations.Count, expectedObservations
+                );
+                m_Observations.RemoveRange(expectedObservations, m_Observations.Count - expectedObservations);
+            }
+            else if (m_Observations.Count < expectedObservations)
+            {
+                // Not enough observations; pad with zeros.
+                Debug.LogWarningFormat(
+                    "Fewer observations ({0}) made than vector observation size ({1}). The observations will be padded.",
+                    m_Observations.Count, expectedObservations
+                );
+                for (int i = m_Observations.Count; i < expectedObservations; i++)
+                {
+                    m_Observations.Add(0);
+                }
+            }
+            adapter.AddRange(m_Observations);
+            Clear();
+            return expectedObservations;
+        }
+
+        public int[] GetFloatObservationShape()
+        {
+            return m_Shape;
+        }
+
+        public string GetName()
+        {
+            return m_Name;
+        }
+
+        public virtual byte[] GetCompressedObservation()
+        {
+            return null;
+        }
+
+        public virtual SensorCompressionType GetCompressionType()
+        {
+            return SensorCompressionType.None;
+        }
+
+        void Clear()
+        {
+            m_Observations.Clear();
+        }
+
+        void AddFloatObs(float obs)
+        {
+            m_Observations.Add(obs);
+        }
+
+        // Compatibility methods with Agent observation. These should be removed eventually.
+
+        /// <summary>
+        /// Adds a float observation to the vector observations of the agent.
+        /// </summary>
+        /// <param name="observation">Observation.</param>
+        public void AddObservation(float observation)
+        {
+            AddFloatObs(observation);
+        }
+
+        /// <summary>
+        /// Adds an integer observation to the vector observations of the agent.
+        /// </summary>
+        /// <param name="observation">Observation.</param>
+        public void AddObservation(int observation)
+        {
+            AddFloatObs(observation);
+        }
+
+        /// <summary>
+        /// Adds an Vector3 observation to the vector observations of the agent.
+        /// </summary>
+        /// <param name="observation">Observation.</param>
+        public void AddObservation(Vector3 observation)
+        {
+            AddFloatObs(observation.x);
+            AddFloatObs(observation.y);
+            AddFloatObs(observation.z);
+        }
+
+        /// <summary>
+        /// Adds an Vector2 observation to the vector observations of the agent.
+        /// </summary>
+        /// <param name="observation">Observation.</param>
+        public void AddObservation(Vector2 observation)
+        {
+            AddFloatObs(observation.x);
+            AddFloatObs(observation.y);
+        }
+
+        /// <summary>
+        /// Adds a collection of float observations to the vector observations of the agent.
+        /// </summary>
+        /// <param name="observation">Observation.</param>
+        public void AddObservation(IEnumerable<float> observation)
+        {
+            foreach (var f in observation)
+            {
+                AddFloatObs(f);
+            }
+        }
+
+        /// <summary>
+        /// Adds a quaternion observation to the vector observations of the agent.
+        /// </summary>
+        /// <param name="observation">Observation.</param>
+        public void AddObservation(Quaternion observation)
+        {
+            AddFloatObs(observation.x);
+            AddFloatObs(observation.y);
+            AddFloatObs(observation.z);
+            AddFloatObs(observation.w);
+        }
+
+        /// <summary>
+        /// Adds a boolean observation to the vector observation of the agent.
+        /// </summary>
+        /// <param name="observation"></param>
+        public void AddObservation(bool observation)
+        {
+            AddFloatObs(observation ? 1f : 0f);
+        }
+
+
+        public void AddOneHotObservation(int observation, int range)
+        {
+            for (var i = 0; i < range; i++)
+            {
+                AddFloatObs(i == observation ? 1.0f : 0.0f);
+            }
+        }
+    }
+}

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

+fileFormatVersion: 2
+guid: e3966c9961b343108808d91a4d140a68
+timeCreated: 1572300800

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

+using System.Collections.Generic;
+using MLAgents.InferenceBrain;
+
+namespace MLAgents.Sensor
+{
+    /// <summary>
+    /// Allows sensors to write to both TensorProxy and float arrays/lists.
+    /// </summary>
+    public class WriteAdapter
+    {
+        IList<float> m_Data;
+        int m_Offset;
+
+        TensorProxy m_Proxy;
+        int m_Batch;
+
+        /// <summary>
+        /// Set the adapter to write to an IList at the given channelOffset.
+        /// </summary>
+        /// <param name="data"></param>
+        /// <param name="offset"></param>
+        public void SetTarget(IList<float> data, int offset)
+        {
+            m_Data = data;
+            m_Offset = offset;
+            m_Proxy = null;
+            m_Batch = -1;
+        }
+
+        /// <summary>
+        /// Set the adapter to write to a TensorProxy at the given batch and channel offset.
+        /// </summary>
+        /// <param name="tensorProxy"></param>
+        /// <param name="batchIndex"></param>
+        /// <param name="channelOffset"></param>
+        public void SetTarget(TensorProxy tensorProxy, int batchIndex, int channelOffset)
+        {
+            m_Proxy = tensorProxy;
+            m_Batch = batchIndex;
+            m_Offset = channelOffset;
+            m_Data = null;
+        }
+
+        /// <summary>
+        /// 1D write access at a specified index. Use AddRange if possible instead.
+        /// </summary>
+        /// <param name="index">Index to write to</param>
+        public float this[int index]
+        {
+            set
+            {
+                if (m_Data != null)
+                {
+                    m_Data[index + m_Offset] = value;
+                }
+                else
+                {
+                    m_Proxy.data[m_Batch, index + m_Offset] = value;
+                }
+            }
+        }
+
+        /// <summary>
+        /// 3D write access at the specified height, width, and channel. Only usable with a TensorProxy target.
+        /// </summary>
+        /// <param name="h"></param>
+        /// <param name="w"></param>
+        /// <param name="ch"></param>
+        public float this[int h, int w, int ch]
+        {
+            set
+            {
+                // Only TensorProxy supports 3D access
+                m_Proxy.data[m_Batch, h, w, ch + m_Offset] = value;
+            }
+        }
+
+        /// <summary>
+        /// Write the range of floats
+        /// </summary>
+        /// <param name="data"></param>
+        /// <param name="writeOffset">Optional write offset</param>
+        public void AddRange(IEnumerable<float> data, int writeOffset = 0)
+        {
+            if (m_Data != null)
+            {
+                int index = 0;
+                foreach (var val in data)
+                {
+                    m_Data[index + m_Offset + writeOffset] = val;
+                    index++;
+                }
+            }
+            else
+            {
+                int index = 0;
+                foreach (var val in data)
+                {
+                    m_Proxy.data[m_Batch, index + m_Offset + writeOffset] = val;
+                    index++;
+                }
+            }
+        }
+    }
+}

UnitySDK/Assets/ML-Agents/Scripts/Sensor/WriteAdapter.cs.meta

+fileFormatVersion: 2
+guid: 86bad2e6dded4a62853752a1713981f2
+timeCreated: 1572540197

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

+using NUnit.Framework;
+using UnityEngine;
+using MLAgents.Sensor;
+
+namespace MLAgents.Tests
+{
+    public class StackingSensorTests
+    {
+        [Test]
+        public void TestCtor()
+        {
+            ISensor wrapped = new VectorSensor(4);
+            ISensor sensor = new StackingSensor(wrapped, 4);
+            Assert.AreEqual("StackingSensor_size4_VectorSensor_size4", sensor.GetName());
+            Assert.AreEqual(sensor.GetFloatObservationShape(), new [] {16});
+        }
+
+        [Test]
+        public void TestStacking()
+        {
+            VectorSensor wrapped = new VectorSensor(2);
+            ISensor sensor = new StackingSensor(wrapped, 3);
+
+            wrapped.AddObservation(new [] {1f, 2f});
+            SensorTestHelper.CompareObservation(sensor, new [] {0f, 0f, 0f, 0f, 1f, 2f});
+
+            wrapped.AddObservation(new [] {3f, 4f});
+            SensorTestHelper.CompareObservation(sensor, new [] {0f, 0f, 1f, 2f, 3f, 4f});
+
+            wrapped.AddObservation(new [] {5f, 6f});
+            SensorTestHelper.CompareObservation(sensor, new [] {1f, 2f, 3f, 4f, 5f, 6f});
+
+            wrapped.AddObservation(new [] {7f, 8f});
+            SensorTestHelper.CompareObservation(sensor, new [] {3f, 4f, 5f, 6f, 7f, 8f});
+
+            wrapped.AddObservation(new [] {9f, 10f});
+            SensorTestHelper.CompareObservation(sensor, new [] {5f, 6f, 7f, 8f, 9f, 10f});
+        }
+
+
+    }
+}

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

+fileFormatVersion: 2
+guid: 7b071fdf91474d18a05ea20175c6b3bd
+timeCreated: 1572564843

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

+using NUnit.Framework;
+using UnityEngine;
+using MLAgents.Sensor;
+
+namespace MLAgents.Tests
+{
+    public class SensorTestHelper
+    {
+        public static void CompareObservation(ISensor sensor, float[] expected)
+        {
+            var numExpected = expected.Length;
+            const float fill = -1337f;
+            var output = new float[numExpected];
+            for (var i = 0; i < numExpected; i++)
+            {
+                output[i] = fill;
+            }
+            Assert.AreEqual(fill, output[0]);
+
+            WriteAdapter writer = new WriteAdapter();
+            writer.SetTarget(output, 0);
+
+            // Make sure WriteAdapter didn't touch anything
+            Assert.AreEqual(fill, output[0]);
+
+            sensor.Write(writer);
+            for (var i = 0; i < numExpected; i++)
+            {
+                Assert.AreEqual(expected[i], output[i]);
+            }
+        }
+    }
+
+    public class VectorSensorTests
+    {
+        [Test]
+        public void TestCtor()
+        {
+            ISensor sensor = new VectorSensor(4);
+            Assert.AreEqual("VectorSensor_size4", sensor.GetName());
+
+            sensor = new VectorSensor(3, "test_sensor");
+            Assert.AreEqual("test_sensor", sensor.GetName());
+        }
+
+        [Test]
+        public void TestWrite()
+        {
+            var sensor = new VectorSensor(4);
+            sensor.AddObservation(1f);
+            sensor.AddObservation(2f);
+            sensor.AddObservation(3f);
+            sensor.AddObservation(4f);
+
+            SensorTestHelper.CompareObservation(sensor, new[] { 1f, 2f, 3f, 4f });
+        }
+
+        [Test]
+        public void TestAddObservationFloat()
+        {
+            var sensor = new VectorSensor(1);
+            sensor.AddObservation(1.2f);
+            SensorTestHelper.CompareObservation(sensor, new []{1.2f});
+        }
+
+        [Test]
+        public void TestAddObservationInt()
+        {
+            var sensor = new VectorSensor(1);
+            sensor.AddObservation(42);
+            SensorTestHelper.CompareObservation(sensor, new []{42f});
+        }
+
+        [Test]
+        public void TestAddObservationVec()
+        {
+            var sensor = new VectorSensor(3);
+            sensor.AddObservation(new Vector3(1,2,3));
+            SensorTestHelper.CompareObservation(sensor, new []{1f, 2f, 3f});
+
+            sensor = new VectorSensor(2);
+            sensor.AddObservation(new Vector2(4,5));
+            SensorTestHelper.CompareObservation(sensor, new[] { 4f, 5f });
+        }
+
+        [Test]
+        public void TestAddObservationQuaternion()
+        {
+            var sensor = new VectorSensor(4);
+            sensor.AddObservation(Quaternion.identity);
+            SensorTestHelper.CompareObservation(sensor, new []{0f, 0f, 0f, 1f});
+        }
+
+        [Test]
+        public void TestWriteEnumerable()
+        {
+            var sensor = new VectorSensor(4);
+            sensor.AddObservation(new [] {1f, 2f, 3f, 4f});
+
+            SensorTestHelper.CompareObservation(sensor, new[] { 1f, 2f, 3f, 4f });
+        }
+
+        [Test]
+        public void TestAddObservationBool()
+        {
+            var sensor = new VectorSensor(1);
+            sensor.AddObservation(true);
+            SensorTestHelper.CompareObservation(sensor, new []{1f});
+        }
+
+        [Test]
+        public void TestAddObservationOneHot()
+        {
+            var sensor = new VectorSensor(4);
+            sensor.AddOneHotObservation(2, 4);
+            SensorTestHelper.CompareObservation(sensor, new []{0f, 0f, 1f, 0f});
+        }
+
+        [Test]
+        public void TestWriteTooMany()
+        {
+            var sensor = new VectorSensor(2);
+            sensor.AddObservation(new [] {1f, 2f, 3f, 4f});
+
+            SensorTestHelper.CompareObservation(sensor, new[] { 1f, 2f});
+        }
+
+        [Test]
+        public void TestWriteNotEnough()
+        {
+            var sensor = new VectorSensor(4);
+            sensor.AddObservation(new [] {1f, 2f});
+
+            // Make sure extra zeros are added
+            SensorTestHelper.CompareObservation(sensor, new[] { 1f, 2f, 0f, 0f});
+        }
+    }
+}

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

+fileFormatVersion: 2
+guid: 18c0d390ce4c5464ab48b96db0392eb0
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

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

+using NUnit.Framework;
+using UnityEngine;
+using MLAgents.Sensor;
+
+using Barracuda;
+using MLAgents.InferenceBrain;
+using MLAgents.InferenceBrain.Utils;
+
+
+namespace MLAgents.Tests
+{
+    public class WriteAdapterTests
+    {
+        [Test]
+        public void TestWritesToIList()
+        {
+            WriteAdapter writer = new WriteAdapter();
+            var buffer = new[] { 0f, 0f, 0f };
+
+            writer.SetTarget(buffer, 0);
+            // Elementwise writes
+            writer[0] = 1f;
+            writer[2] = 2f;
+            Assert.AreEqual(new[] { 1f, 0f, 2f }, buffer);
+
+            // Elementwise writes with offset
+            writer.SetTarget(buffer, 1);
+            writer[0] = 3f;
+            Assert.AreEqual(new[] { 1f, 3f, 2f }, buffer);
+
+            // AddRange
+            writer.SetTarget(buffer, 0);
+            writer.AddRange(new [] {4f, 5f});
+            Assert.AreEqual(new[] { 4f, 5f, 2f }, buffer);
+
+            // AddRange with offset
+            writer.SetTarget(buffer, 1);
+            writer.AddRange(new [] {6f, 7f});
+            Assert.AreEqual(new[] { 4f, 6f, 7f }, buffer);
+        }
+
+        [Test]
+        public void TestWritesToTensor()
+        {
+            WriteAdapter writer = new WriteAdapter();
+            var t = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                data = new Tensor(2, 3)
+            };
+            writer.SetTarget(t, 0, 0);
+            Assert.AreEqual(0f, t.data[0, 0]);
+            writer[0] = 1f;
+            Assert.AreEqual(1f, t.data[0, 0]);
+
+            writer.SetTarget(t, 1, 1);
+            writer[0] = 2f;
+            writer[1] = 3f;
+            // [0, 0] shouldn't change
+            Assert.AreEqual(1f, t.data[0, 0]);
+            Assert.AreEqual(2f, t.data[1, 1]);
+            Assert.AreEqual(3f, t.data[1, 2]);
+
+            // AddRange
+            t = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                data = new Tensor(2, 3)
+            };
+
+            writer.SetTarget(t, 1, 1);
+            writer.AddRange(new [] {-1f, -2f});
+            Assert.AreEqual(0f, t.data[0, 0]);
+            Assert.AreEqual(0f, t.data[0, 1]);
+            Assert.AreEqual(0f, t.data[0, 2]);
+            Assert.AreEqual(0f, t.data[1, 0]);
+            Assert.AreEqual(-1f, t.data[1, 1]);
+            Assert.AreEqual(-2f, t.data[1, 2]);
+        }
+
+        [Test]
+        public void TestWritesToTensor3D()
+        {
+            WriteAdapter writer = new WriteAdapter();
+            var t = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                data = new Tensor(2, 2, 2, 3)
+            };
+
+            writer.SetTarget(t, 0, 0);
+            writer[1, 0, 1] = 1f;
+            Assert.AreEqual(1f, t.data[0, 1, 0, 1]);
+
+            writer.SetTarget(t, 0, 1);
+            writer[1, 0, 0] = 2f;
+            Assert.AreEqual(2f, t.data[0, 1, 0, 1]);
+        }
+    }
+}

UnitySDK/Assets/ML-Agents/Editor/Tests/Sensor/WriterAdapterTests.cs.meta

+fileFormatVersion: 2
+guid: 3de9cbda816e4d7b907e765577dd54f7
+timeCreated: 1572568337