From 103b3255961e7a14111962bacf31718c7cdb0a91 Mon Sep 17 00:00:00 2001
From: Chris Elion <chris.elion@unity3d.com>
Date: Wed, 17 Feb 2021 10:53:04 -0800
Subject: [PATCH] pass sensor dimension flags to analytics

---
 com.unity.ml-agents/Runtime/Analytics/Events.cs      |  3 ++-
 .../Editor/Analytics/InferenceAnalyticsTests.cs      |  2 ++
 .../Tests/Editor/ParameterLoaderTest.cs              | 12 +++++++++++-
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Analytics/Events.cs b/com.unity.ml-agents/Runtime/Analytics/Events.cs
index ed22d3c67d..78cbe6c627 100644
--- a/com.unity.ml-agents/Runtime/Analytics/Events.cs
+++ b/com.unity.ml-agents/Runtime/Analytics/Events.cs
@@ -72,11 +72,12 @@ internal struct EventObservationSpec
         public static EventObservationSpec FromSensor(ISensor sensor)
         {
             var shape = sensor.GetObservationShape();
+            var dimProps = (sensor as IDimensionPropertiesSensor)?.GetDimensionProperties();
             var dimInfos = new EventObservationDimensionInfo[shape.Length];
             for (var i = 0; i < shape.Length; i++)
             {
                 dimInfos[i].Size = shape[i];
-                // TODO copy flags when we have them
+                dimInfos[i].Flags = dimProps != null ? (int)dimProps[i] : 0;
             }
 
             var builtInSensorType =
diff --git a/com.unity.ml-agents/Tests/Editor/Analytics/InferenceAnalyticsTests.cs b/com.unity.ml-agents/Tests/Editor/Analytics/InferenceAnalyticsTests.cs
index 489c64f2f0..8d353f0dc6 100644
--- a/com.unity.ml-agents/Tests/Editor/Analytics/InferenceAnalyticsTests.cs
+++ b/com.unity.ml-agents/Tests/Editor/Analytics/InferenceAnalyticsTests.cs
@@ -58,6 +58,8 @@ public void TestModelEvent()
             Assert.AreEqual(2, continuousEvent.ObservationSpecs.Count);
             Assert.AreEqual(3, continuousEvent.ObservationSpecs[0].DimensionInfos.Length);
             Assert.AreEqual(20, continuousEvent.ObservationSpecs[0].DimensionInfos[0].Size);
+            Assert.AreEqual((int)DimensionProperty.TranslationalEquivariance, continuousEvent.ObservationSpecs[0].DimensionInfos[0].Flags);
+            Assert.AreEqual((int)DimensionProperty.None, continuousEvent.ObservationSpecs[0].DimensionInfos[2].Flags);
             Assert.AreEqual("None", continuousEvent.ObservationSpecs[0].CompressionType);
             Assert.AreEqual(Test3DSensor.k_BuiltInSensorType, continuousEvent.ObservationSpecs[0].BuiltInSensorType);
             Assert.AreNotEqual(null, continuousEvent.ModelHash);
diff --git a/com.unity.ml-agents/Tests/Editor/ParameterLoaderTest.cs b/com.unity.ml-agents/Tests/Editor/ParameterLoaderTest.cs
index 7c1a3cca20..dc96d28a79 100644
--- a/com.unity.ml-agents/Tests/Editor/ParameterLoaderTest.cs
+++ b/com.unity.ml-agents/Tests/Editor/ParameterLoaderTest.cs
@@ -24,7 +24,7 @@ public override int[] GetObservationShape()
             return Sensor.GetObservationShape();
         }
     }
-    public class Test3DSensor : ISensor, IBuiltInSensor
+    public class Test3DSensor : ISensor, IBuiltInSensor, IDimensionPropertiesSensor
     {
         int m_Width;
         int m_Height;
@@ -77,6 +77,16 @@ public BuiltInSensorType GetBuiltInSensorType()
         {
             return (BuiltInSensorType)k_BuiltInSensorType;
         }
+
+        public DimensionProperty[] GetDimensionProperties()
+        {
+            return new[]
+            {
+                DimensionProperty.TranslationalEquivariance,
+                DimensionProperty.TranslationalEquivariance,
+                DimensionProperty.None
+            };
+        }
     }
 
     [TestFixture]